This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch db-oauth2-client-info
in repository https://gitbox.apache.org/repos/asf/superset.git

commit d64c49c0855a43a2cb5aa773ade711efdcdb1698
Author: Beto Dealmeida <[email protected]>
AuthorDate: Sun May 5 17:15:58 2024 -0400

    feat: read OAuth2 info from database extra
---
 superset/models/core.py  | 26 +++++++++++++++++---------
 superset/utils/oauth2.py | 14 +++++++++++++-
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/superset/models/core.py b/superset/models/core.py
index 9a4a1de403..79309cdb3d 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -30,7 +30,7 @@ from contextlib import closing, contextmanager, nullcontext, 
suppress
 from copy import deepcopy
 from datetime import datetime
 from functools import lru_cache
-from typing import Any, Callable, TYPE_CHECKING
+from typing import Any, Callable, cast, TYPE_CHECKING
 
 import numpy
 import pandas as pd
@@ -79,7 +79,7 @@ from superset.superset_typing import OAuth2ClientConfig, 
ResultSetColumnType
 from superset.utils import cache as cache_util, core as utils
 from superset.utils.backports import StrEnum
 from superset.utils.core import get_username
-from superset.utils.oauth2 import get_oauth2_access_token
+from superset.utils.oauth2 import get_oauth2_access_token, 
OAuth2ClientConfigSchema
 
 config = app.config
 custom_password_store = config["SQLALCHEMY_CUSTOM_PASSWORD_STORE"]
@@ -1022,20 +1022,28 @@ class Database(Model, AuditMixinNullable, 
ImportExportMixin):  # pylint: disable
         """
         Is OAuth2 enabled in the database for authentication?
 
-        Currently this looks for a global config at the DB engine spec level, 
but in the
-        future we want to be allow admins to create custom OAuth2 clients from 
the
-        Superset UI, and assign them to specific databases.
+        Currently this checks for configuration stored in the database 
`extra`, and then
+        for a global config at the DB engine spec level. In the future we want 
to allow
+        admins to create custom OAuth2 clients from the Superset UI, and 
assign them to
+        specific databases.
         """
-        return self.db_engine_spec.is_oauth2_enabled()
+        oauth2_client_info = self.get_extra().get("oauth2_client_info", {})
+        return bool(oauth2_client_info) or 
self.db_engine_spec.is_oauth2_enabled()
 
     def get_oauth2_config(self) -> OAuth2ClientConfig | None:
         """
         Return OAuth2 client configuration.
 
-        This includes client ID, client secret, scope, redirect URI, 
endpointsm etc.
-        Currently this reads the global DB engine spec config, but in the 
future it
-        should first check if there's a custom client assigned to the database.
+        Currently this checks for configuration stored in the database 
`extra`, and then
+        for a global config at the DB engine spec level. In the future we want 
to allow
+        admins to create custom OAuth2 clients from the Superset UI, and 
assign them to
+        specific databases.
         """
+        if oauth2_client_info := self.get_extra().get("oauth2_client_info"):
+            schema = OAuth2ClientConfigSchema()
+            client_config = schema.load(oauth2_client_info)
+            return cast(OAuth2ClientConfig, client_config)
+
         return self.db_engine_spec.get_oauth2_config()
 
 
diff --git a/superset/utils/oauth2.py b/superset/utils/oauth2.py
index 9cc58a0b7f..7b7440b059 100644
--- a/superset/utils/oauth2.py
+++ b/superset/utils/oauth2.py
@@ -22,7 +22,7 @@ from typing import Any, TYPE_CHECKING
 
 import backoff
 import jwt
-from flask import current_app
+from flask import current_app, url_for
 from marshmallow import EXCLUDE, fields, post_load, Schema
 
 from superset import db
@@ -180,3 +180,15 @@ def decode_oauth2_state(encoded_state: str) -> OAuth2State:
     state = oauth2_state_schema.load(payload)
 
     return state
+
+
+class OAuth2ClientConfigSchema(Schema):
+    id = fields.String(required=True)
+    secret = fields.String(required=True)
+    scope = fields.String(required=True)
+    redirect_uri = fields.String(
+        required=False,
+        load_default=url_for("DatabaseRestApi.oauth2", _external=True),
+    )
+    authorization_request_uri = fields.String(required=True)
+    token_request_uri = fields.String(required=True)

Reply via email to