This is an automated email from the ASF dual-hosted git repository. beto pushed a commit to branch client-info-db-extra in repository https://gitbox.apache.org/repos/asf/superset.git
commit efd725025d1f300f48692b1aeded4ef657684a0a Author: Beto Dealmeida <[email protected]> AuthorDate: Thu Jun 6 08:44:08 2024 -0400 feat: OAuth2 client initial work --- superset/models/core.py | 62 ++++++++++++++++++++++++++++++------------------ superset/utils/oauth2.py | 14 ++++++++++- 2 files changed, 52 insertions(+), 24 deletions(-) diff --git a/superset/models/core.py b/superset/models/core.py index e6d97a197b..719328afd6 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -29,7 +29,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 @@ -78,7 +78,7 @@ from superset.superset_typing import OAuth2ClientConfig, ResultSetColumnType from superset.utils import cache as cache_util, core as utils, json from superset.utils.backports import StrEnum from superset.utils.core import DatasourceName, 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"] @@ -116,7 +116,9 @@ class ConfigurationMethod(StrEnum): DYNAMIC_FORM = "dynamic_form" -class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable=too-many-public-methods +class Database( + Model, AuditMixinNullable, ImportExportMixin +): # pylint: disable=too-many-public-methods """An ORM object that stores Database related information""" __tablename__ = "dbs" @@ -390,9 +392,7 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable return ( username if (username := get_username()) - else object_url.username - if self.impersonate_user - else None + else object_url.username if self.impersonate_user else None ) @contextmanager @@ -554,17 +554,23 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable nullpool=nullpool, source=source, ) as engine: - with closing(engine.raw_connection()) as conn: - # pre-session queries are used to set the selected schema and, in the - # future, the selected catalog - for prequery in self.db_engine_spec.get_prequeries( - catalog=catalog, - schema=schema, - ): - cursor = conn.cursor() - cursor.execute(prequery) + try: + with closing(engine.raw_connection()) as conn: + # pre-session queries are used to set the selected schema and, in the + # future, the selected catalog + for prequery in self.db_engine_spec.get_prequeries( + catalog=catalog, + schema=schema, + ): + cursor = conn.cursor() + cursor.execute(prequery) - yield conn + yield conn + + except Exception as ex: + if self.is_oauth2_enabled(): + self.db_engine_spec.start_oauth2_dance(self) + raise ex def get_default_catalog(self) -> str | None: """ @@ -1063,20 +1069,30 @@ 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() + config = json.loads(self.encrypted_extra or "{}") + oauth2_client_info = config.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. """ + config = json.loads(self.encrypted_extra or "{}") + if oauth2_client_info := config.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..bc4805fd81 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=lambda: url_for("DatabaseRestApi.oauth2", _external=True), + ) + authorization_request_uri = fields.String(required=True) + token_request_uri = fields.String(required=True)
