From bedb8dce78dbada84aa5595a09e3a8f9324f99f0 Mon Sep 17 00:00:00 2001 From: Alexander Kharkevich Date: Wed, 16 Oct 2024 22:35:58 -0400 Subject: [PATCH] feat: add redis cache support --- README.md | 16 ++++++++++++++++ mlflow_oidc_auth/config.py | 21 ++++++++++++++++++++- pyproject.toml | 3 ++- scripts/docker-compose.yaml | 14 ++++++++++++++ scripts/run-dev-server.sh | 2 +- 5 files changed, 53 insertions(+), 3 deletions(-) create mode 100644 scripts/docker-compose.yaml diff --git a/README.md b/README.md index 455d522..b97b603 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ python3 -m pip install mlflow-oidc-auth # Configuration The plugin required the following environment variables but also supported `.env` file +## Application configuration | Parameter | Description| |---|---| | OIDC_REDIRECT_URI | Application redirect/callback url (https://example.com/callback) | @@ -36,6 +37,21 @@ The plugin required the following environment variables but also supported `.env | LOG_LEVEL | Application log level | | OIDC_USERS_DB_URI | Database connection string | +## Application session storage configuration +| Parameter | Description | Default | +|---|---|---| +| SESSION_TYPE | Flask session type (filesystem or redis supported) | filesystem | +| SESSION_FILE_DIR | The directory where session files are stored | flask_session | +| SESSION_PERMANENT | Whether use permanent session or not | False | +| PERMANENT_SESSION_LIFETIME | Server-side session expiration time (in seconds) | 86400 | +| SESSION_KEY_PREFIX | A prefix that is added before all session keys | mlflow_oidc: | +| REDIS_HOST | Redis hostname | localhost | +| REDIS_PORT | Redis port | 6379 | +| REDIS_DB | Redis DB number | 0 | +| REDIS_USERNAME | Redis username | None | +| REDIS_PASSWORD | Redis password | None | +| REDIS_SSL | Use SSL | false | + # Configuration examples ## Okta diff --git a/mlflow_oidc_auth/config.py b/mlflow_oidc_auth/config.py index 95506bf..1121d78 100644 --- a/mlflow_oidc_auth/config.py +++ b/mlflow_oidc_auth/config.py @@ -13,7 +13,6 @@ class AppConfig: DEFAULT_MLFLOW_PERMISSION = os.environ.get("DEFAULT_MLFLOW_PERMISSION", "MANAGE") SECRET_KEY = os.environ.get("SECRET_KEY", secrets.token_hex(16)) - SESSION_TYPE = "cachelib" OIDC_USERS_DB_URI = os.environ.get("OIDC_USERS_DB_URI", "sqlite:///auth.db") OIDC_GROUP_NAME = os.environ.get("OIDC_GROUP_NAME", "mlflow") OIDC_ADMIN_GROUP_NAME = os.environ.get("OIDC_ADMIN_GROUP_NAME", "mlflow-admin") @@ -36,6 +35,26 @@ class AppConfig: OIDC_CLIENT_ID = os.environ.get("OIDC_CLIENT_ID", None) OIDC_CLIENT_SECRET = os.environ.get("OIDC_CLIENT_SECRET", None) + # https://flask-session.readthedocs.io/en/latest/config.html + SESSION_TYPE = os.environ.get("SESSION_TYPE", "filesystem") + SESSION_PERMANENT = os.environ.get("SESSION_PERMANENT", str(False)).lower() in ("true", "1", "t") + SESSION_KEY_PREFIX = os.environ.get("SESSION_KEY_PREFIX", "mlflow_oidc:") + PERMANENT_SESSION_LIFETIME = os.environ.get("PERMANENT_SESSION_LIFETIME", 86400) + if SESSION_TYPE == "filesystem": + SESSION_FILE_DIR = os.environ.get("SESSION_FILE_DIR", "./flask_session/") + elif SESSION_TYPE == "redis": + import redis + SESSION_REDIS = redis.Redis( + host=os.environ.get("REDIS_HOST", "localhost"), + port=os.environ.get("REDIS_PORT", 6379), + db=os.environ.get("REDIS_DB", 0), + password=os.environ.get("REDIS_PASSWORD", None), + ssl=os.environ.get("REDIS_SSL", str(False)).lower() in ("true", "1", "t"), + username=os.environ.get("REDIS_USERNAME", None), + ) + else: + raise ValueError(f"Invalid session type: {SESSION_TYPE}") + @staticmethod def get_property(property_name): app.logger.debug(f"Getting property {property_name}") diff --git a/pyproject.toml b/pyproject.toml index 5ef1434..c905976 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,8 @@ dependencies = [ [project.optional-dependencies] full = ["mlflow<3,>=2.11.1"] -dev = ["black==24.8.0", "pre-commit==3.5.0"] +caching-redis = ["redis[hiredis]<6"] +dev = ["black<25", "pre-commit<5"] [[project.maintainers]] name = "Data Platform folks" diff --git a/scripts/docker-compose.yaml b/scripts/docker-compose.yaml new file mode 100644 index 0000000..d208a7a --- /dev/null +++ b/scripts/docker-compose.yaml @@ -0,0 +1,14 @@ +# docker compose to test redis cache integration +services: + redis: + image: redis:6.0.9 + container_name: redis + ports: + - "6379:6379" + redisinsight: + image: redislabs/redisinsight:latest + container_name: redisinsight + ports: + - "5540:5540" + environment: + - REDIS_URI=redis://redis:6379 diff --git a/scripts/run-dev-server.sh b/scripts/run-dev-server.sh index ab5bc30..143028b 100755 --- a/scripts/run-dev-server.sh +++ b/scripts/run-dev-server.sh @@ -13,7 +13,7 @@ python_preconfigure() { source venv/bin/activate python3 -m pip install --upgrade pip python3 -m pip install build setuptools - python3 -m pip install --editable=".[full]" + python3 -m pip install --editable=".[full, caching-redis]" fi }