feat: support default group permission

.gitignore

@@ -189,3 +189,6 @@ integration-tests/
 
 # until it ready
 web-react/
+
+# VS Code
+.vscode

README.md

@@ -18,39 +18,38 @@ python3 -m pip install mlflow-oidc-auth
 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) |
-| OIDC_DISCOVERY_URL     | OIDC Discovery URL |
-| OIDC_CLIENT_SECRET     | OIDC Client Secret |
-| OIDC_CLIENT_ID         |  OIDC Client ID |
-| OIDC_GROUP_DETECTION_PLUGIN | OIDC plugin to detect groups |
-| OIDC_PROVIDER_DISPLAY_NAME | any text to display |
-| OIDC_SCOPE | OIDC scope |
-| OIDC_GROUP_NAME | User group name to be allowed login to MLFlow, currently supported groups in OIDC claims and Microsoft Entra ID groups |
-| OIDC_ADMIN_GROUP_NAME | User group name to be allowed login to MLFlow manage and define permissions, currently supported groups in OIDC claims and Microsoft Entra ID groups |
-| OIDC_AUTHORIZATION_URL | OIDC Auth URL (if discovery URL is not defined) |
-| OIDC_TOKEN_URL         | OIDC Token URL (if discovery URL is not defined) |
-| OIDC_USER_URL          | OIDC User info URL (if discovery URL is not defined) |
-| SECRET_KEY             | Key to perform cookie encryption |
-| LOG_LEVEL                   | Application log level |
-| OIDC_USERS_DB_URI | Database connection string |
-| OIDC_ALEMBIC_VERSION_TABLE  | Name of the table to use for alembic versions (defaults to alembic_version if not provided)                                                          |
+| Parameter | Description| Default | Mandatory |
+|---|---|---|---|
+| OIDC_REDIRECT_URI      |  Application redirect/callback url (https://example.com/callback) | None | Yes |
+| OIDC_DISCOVERY_URL     | OIDC Discovery URL | None | Yes |
+| OIDC_CLIENT_SECRET     | OIDC Client Secret | None | Yes |
+| OIDC_CLIENT_ID         |  OIDC Client ID | None | Yes |
+| OIDC_GROUP_DETECTION_PLUGIN | OIDC plugin to detect groups | None | No |
+| OIDC_PROVIDER_DISPLAY_NAME | any text to display | "Login with OIDC" | No |
+| OIDC_SCOPE | OIDC scope | "openid,email,profile" | No |
+| OIDC_GROUP_NAME | User group name to be allowed login to MLFlow, currently supported groups in OIDC claims and Microsoft Entra ID groups | "mlflow" | No |
+| OIDC_ADMIN_GROUP_NAME | User group name to be allowed login to MLFlow manage and define permissions, currently supported groups in OIDC claims and Microsoft Entra ID groups | "mlflow-admin" | No |
+| SECRET_KEY             | Key to perform cookie encryption | A secret key will be generated | No |
+| LOG_LEVEL                   | Application log level | "INFO" | No |
+| OIDC_USERS_DB_URI | Database connection string | "sqlite:///auth.db" | No |
+| OIDC_ALEMBIC_VERSION_TABLE  | Name of the table to use for alembic versions | "alembic_version" | No |
+| DEFAULT_MLFLOW_PERMISSION         | Default fallback permission on all resources  | "MANAGE" | No |
+| DEFAULT_MLFLOW_GROUP_PERMISSION   | Default group permission assigned on resource creation, no permission will be assigned if unspecified | None | No |
 
 ## 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 |
+| Parameter | Description | Default | Mandatory |
+|---|---|---|---|
+| SESSION_TYPE | Flask session type (filesystem or redis supported) | filesystem | No |
+| SESSION_FILE_DIR | The directory where session files are stored | flask_session | No |
+| SESSION_PERMANENT | Whether use permanent session or not | False | No |
+| PERMANENT_SESSION_LIFETIME | Server-side session expiration time (in seconds) | 86400 | No |
+| SESSION_KEY_PREFIX | A prefix that is added before all session keys | mlflow_oidc: | No |
+| REDIS_HOST | Redis hostname | localhost | No |
+| REDIS_PORT | Redis port | 6379 | No |
+| REDIS_DB | Redis DB number | 0 | No |
+| REDIS_USERNAME | Redis username | None | No |
+| REDIS_PASSWORD | Redis password | None | No |
+| REDIS_SSL | Use SSL | false | No |
 
 # Configuration examples
 

docs/configuration/index.md

@@ -2,35 +2,35 @@
 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) |
-| OIDC_DISCOVERY_URL     | OIDC Discovery URL |
-| OIDC_CLIENT_SECRET     | OIDC Client Secret |
-| OIDC_CLIENT_ID         |  OIDC Client ID |
-| OIDC_GROUP_DETECTION_PLUGIN | OIDC plugin to detect groups |
-| OIDC_PROVIDER_DISPLAY_NAME | any text to display |
-| OIDC_SCOPE | OIDC scope |
-| OIDC_GROUP_NAME | User group name to be allowed login to MLFlow, currently supported groups in OIDC claims and Microsoft Entra ID groups |
-| OIDC_ADMIN_GROUP_NAME | User group name to be allowed login to MLFlow manage and define permissions, currently supported groups in OIDC claims and Microsoft Entra ID groups |
-| OIDC_AUTHORIZATION_URL | OIDC Auth URL (if discovery URL is not defined) |
-| OIDC_TOKEN_URL         | OIDC Token URL (if discovery URL is not defined) |
-| OIDC_USER_URL          | OIDC User info URL (if discovery URL is not defined) |
-| SECRET_KEY             | Key to perform cookie encryption |
-| LOG_LEVEL                   | Application log level |
-| OIDC_USERS_DB_URI | Database connection string |
+| Parameter | Description| Default | Mandatory |
+|---|---|---|---|
+| OIDC_REDIRECT_URI      |  Application redirect/callback url (https://example.com/callback) | None | Yes |
+| OIDC_DISCOVERY_URL     | OIDC Discovery URL | None | Yes |
+| OIDC_CLIENT_SECRET     | OIDC Client Secret | None | Yes |
+| OIDC_CLIENT_ID         |  OIDC Client ID | None | Yes |
+| OIDC_GROUP_DETECTION_PLUGIN | OIDC plugin to detect groups | None | No |
+| OIDC_PROVIDER_DISPLAY_NAME | any text to display | "Login with OIDC" | No |
+| OIDC_SCOPE | OIDC scope | "openid,email,profile" | No |
+| OIDC_GROUP_NAME | User group name to be allowed login to MLFlow, currently supported groups in OIDC claims and Microsoft Entra ID groups | "mlflow" | No |
+| OIDC_ADMIN_GROUP_NAME | User group name to be allowed login to MLFlow manage and define permissions, currently supported groups in OIDC claims and Microsoft Entra ID groups | "mlflow-admin" | No |
+| SECRET_KEY             | Key to perform cookie encryption | A secret key will be generated | No |
+| LOG_LEVEL                   | Application log level | "INFO" | No |
+| OIDC_USERS_DB_URI | Database connection string | "sqlite:///auth.db" | No |
+| OIDC_ALEMBIC_VERSION_TABLE  | Name of the table to use for alembic versions | "alembic_version" | No |
+| DEFAULT_MLFLOW_PERMISSION         | Default fallback permission on all resources  | "MANAGE" | No |
+| DEFAULT_MLFLOW_GROUP_PERMISSION   | Default group permission assigned on resource creation, no permission will be assigned if unspecified | None | No |
 
 ## 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 |
+| Parameter | Description | Default | Mandatory |
+|---|---|---|---|
+| SESSION_TYPE | Flask session type (filesystem or redis supported) | filesystem | No |
+| SESSION_FILE_DIR | The directory where session files are stored | flask_session | No |
+| SESSION_PERMANENT | Whether use permanent session or not | False | No |
+| PERMANENT_SESSION_LIFETIME | Server-side session expiration time (in seconds) | 86400 | No |
+| SESSION_KEY_PREFIX | A prefix that is added before all session keys | mlflow_oidc: | No |
+| REDIS_HOST | Redis hostname | localhost | No |
+| REDIS_PORT | Redis port | 6379 | No |
+| REDIS_DB | Redis DB number | 0 | No |
+| REDIS_USERNAME | Redis username | None | No |
+| REDIS_PASSWORD | Redis password | None | No |
+| REDIS_SSL | Use SSL | false | No |

mlflow_oidc_auth/config.py

@@ -17,6 +17,7 @@ def get_bool_env_variable(variable, default_value):
 class AppConfig:
     def __init__(self):
         self.DEFAULT_MLFLOW_PERMISSION = os.environ.get("DEFAULT_MLFLOW_PERMISSION", "MANAGE")
+        self.DEFAULT_MLFLOW_GROUP_PERMISSION = os.environ.get("DEFAULT_MLFLOW_GROUP_PERMISSION", None)
         self.SECRET_KEY = os.environ.get("SECRET_KEY", secrets.token_hex(16))
         self.OIDC_USERS_DB_URI = os.environ.get("OIDC_USERS_DB_URI", "sqlite:///auth.db")
         self.OIDC_GROUP_NAME = [group.strip() for group in os.environ.get("OIDC_GROUP_NAME", "mlflow").split(",")]

mlflow_oidc_auth/hooks/after_request.py

@@ -14,6 +14,7 @@
 from mlflow.utils.proto_json_utils import message_to_json, parse_dict
 from mlflow.utils.search_utils import SearchUtils
 
+from mlflow_oidc_auth.config import config
 from mlflow_oidc_auth.permissions import MANAGE
 from mlflow_oidc_auth.store import store
 from mlflow_oidc_auth.utils import (
@@ -25,26 +26,38 @@
     fetch_registered_models_paginated,
     fetch_readable_registered_models,
     fetch_readable_experiments,
+    get_user_groups,
 )
 
 
-def _set_can_manage_experiment_permission(resp: Response):
-    response_message = CreateExperiment.Response()  # type: ignore
+def _set_initial_experiment_permission(resp: Response):
+    response_message = CreateExperiment.Response()
     parse_dict(resp.json, response_message)
     experiment_id = response_message.experiment_id
     username = get_username()
     store.create_experiment_permission(experiment_id, username, MANAGE.name)
+    user_groups = get_user_groups(username)
+    if permission := config.DEFAULT_MLFLOW_GROUP_PERMISSION:
+        for group_name in user_groups:
+            store.create_group_experiment_permission(group_name, experiment_id, permission)
 
 
-def _set_can_manage_registered_model_permission(resp: Response):
+def _set_initial_registered_model_permission(resp: Response):
     response_message = CreateRegisteredModel.Response()  # type: ignore
     parse_dict(resp.json, response_message)
-    name = response_message.registered_model.name
+    model_name = response_message.registered_model.name
     username = get_username()
-    store.create_registered_model_permission(name, username, MANAGE.name)
+    store.create_registered_model_permission(model_name, username, MANAGE.name)
+    user_groups = get_user_groups(username)
+    if permission := config.DEFAULT_MLFLOW_GROUP_PERMISSION:
+        for group_name in user_groups:
+            store.create_group_model_permission(group_name, model_name, permission)
 
 
-def _delete_can_manage_registered_model_permission(resp: Response):
+# TODO: Should a _delete_experiment_permission be added?
+
+
+def _delete_registered_model_permission(resp: Response):
     """
     Delete registered model permission when the model is deleted.
 
@@ -53,7 +66,6 @@ def _delete_can_manage_registered_model_permission(resp: Response):
     we have to delete the permission record when the model is deleted otherwise it
     conflicts with the new model registered with the same name.
     """
-    # Get model name from request context because it's not available in the response
     model_name = get_model_name()
     store.wipe_group_model_permissions(model_name)
     store.wipe_registered_model_permissions(model_name)
@@ -130,9 +142,9 @@ def _filter_search_registered_models(resp: Response):
 
 
 AFTER_REQUEST_PATH_HANDLERS = {
-    CreateExperiment: _set_can_manage_experiment_permission,
-    CreateRegisteredModel: _set_can_manage_registered_model_permission,
-    DeleteRegisteredModel: _delete_can_manage_registered_model_permission,
+    CreateExperiment: _set_initial_experiment_permission,
+    CreateRegisteredModel: _set_initial_registered_model_permission,
+    DeleteRegisteredModel: _delete_registered_model_permission,
     SearchExperiments: _filter_search_experiments,
     SearchRegisteredModels: _filter_search_registered_models,
 }

mlflow_oidc_auth/tests/test_utils.py

@@ -1,12 +1,15 @@
 import unittest
 from unittest.mock import MagicMock, patch
+from flask import Flask, session, request
 
-from flask import Flask
 from mlflow.exceptions import MlflowException
 from mlflow.protos.databricks_pb2 import BAD_REQUEST, INVALID_PARAMETER_VALUE, RESOURCE_DOES_NOT_EXIST
 
 from mlflow_oidc_auth.permissions import Permission
 from mlflow_oidc_auth.utils import (
+    get_is_admin,
+    get_user_groups,
+    get_permission_from_store_or_default,
     PermissionResult,
     can_manage_experiment,
     can_manage_registered_model,
@@ -49,6 +52,7 @@
 class TestUtils(unittest.TestCase):
     def setUp(self):
         self.app = Flask(__name__)
+        self.app.secret_key = 'test_secret_key'
         self.app.config["TESTING"] = True
         self.app_context = self.app.app_context()
         self.app_context.push()
@@ -199,6 +203,72 @@ def mock_func():
             mock_get_is_admin.return_value = True
             self.assertEqual(mock_func(), "success")
 
+    @patch("importlib.import_module")
+    @patch("mlflow_oidc_auth.utils.app")
+    @patch('mlflow_oidc_auth.utils.config')
+    @patch('mlflow_oidc_auth.utils.validate_token')
+    def test_get_user_groups_from_plugin(self, mock_validate_token,
+                                               mock_config, mock_app,
+                                               mock_import_module):
+        mock_validate_token.return_value = {"test_oidc_groups":
+                                            ['group1', 'group2']}
+        mock_config.OIDC_GROUPS_ATTRIBUTE = "test_oidc_groups"
+        mock_config.OIDC_GROUP_NAME = ["group1", "group2"]
+        mock_config.OIDC_GROUP_DETECTION_PLUGIN = "groups_plugin"
+        mock_plugin = MagicMock()
+        mock_plugin.get_user_groups.return_value = ['group1', 'group2']
+        mock_import_module.return_value = mock_plugin
+        mock_app.logger.debug = MagicMock()
+
+        with self.app.test_request_context(headers={'Authorization':
+                                                    'Bearer test_token'}):
+            groups = get_user_groups()
+            assert groups == ['group1', 'group2']
+            mock_app.logger.debug.assert_called_once_with(
+                f"Groups from plugin: {groups}"
+                )
+
+    @patch("mlflow_oidc_auth.utils.app")
+    @patch('mlflow_oidc_auth.utils.config')
+    @patch('mlflow_oidc_auth.utils.validate_token')
+    def test_get_user_groups_from_bearer_token(self, mock_validate_token,
+                                               mock_config, mock_app):
+        mock_validate_token.return_value = {"test_oidc_groups":
+                                            ['group1', 'group2']}
+        mock_config.OIDC_GROUPS_ATTRIBUTE = "test_oidc_groups"
+        mock_config.OIDC_GROUP_NAME = ["group1", "group2"]
+        mock_config.OIDC_GROUP_DETECTION_PLUGIN = None
+        mock_app.logger.debug = MagicMock()
+
+        with self.app.test_request_context(headers={'Authorization':
+                                                    'Bearer test_token'}):
+            groups = get_user_groups()
+            assert groups == ['group1', 'group2']
+            mock_app.logger.debug.assert_called_once_with(
+                f"Groups from bearer token: {groups}"
+                )
+
+    @patch("mlflow_oidc_auth.utils.app")
+    @patch('mlflow_oidc_auth.utils.store.get_groups_for_user')
+    def test_get_user_groups_from_store(self, mock_get_groups_for_user,
+                                        mock_app):
+        mock_get_groups_for_user.return_value = ['group1', 'group2']
+        mock_app.logger.debug = MagicMock()
+
+        with self.app.test_request_context():
+            groups = get_user_groups(username='test_user')
+            assert groups == ['group1', 'group2']
+            mock_app.logger.debug.assert_called_once_with(
+                f"Groups from store: {groups}"
+                )
+
+    @patch("mlflow_oidc_auth.utils.app")
+    def test_get_user_groups_no_groups_found(self, mock_app):
+        with self.app.test_request_context():
+            groups = get_user_groups()
+            assert groups == []
+            mock_app.logger.debug.assert_not_called()
+
     @patch("mlflow_oidc_auth.utils.store")
     @patch("mlflow_oidc_auth.utils.get_is_admin")
     @patch("mlflow_oidc_auth.utils.get_username")