feat: add redis cache support

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

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}")

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"

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

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
 }