You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by be...@apache.org on 2024/05/10 15:25:03 UTC
(superset) branch db-oauth2-client-info created (now f31084d002)
This is an automated email from the ASF dual-hosted git repository.
beto pushed a change to branch db-oauth2-client-info
in repository https://gitbox.apache.org/repos/asf/superset.git
at f31084d002 WIP
This branch includes the following new commits:
new d64c49c085 feat: read OAuth2 info from database extra
new f31084d002 WIP
The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(superset) 02/02: WIP
Posted by be...@apache.org.
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 f31084d002d42d2a9c8765a602ec1f90de5a3a7f
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Wed May 8 10:18:54 2024 -0400
WIP
---
superset/db_engine_specs/snowflake.py | 84 +++++++++++++++++++++++++++++++++++
1 file changed, 84 insertions(+)
diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py
index 83d382cda1..a17297217e 100644
--- a/superset/db_engine_specs/snowflake.py
+++ b/superset/db_engine_specs/snowflake.py
@@ -14,6 +14,9 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+
+from __future__ import annotations
+
import json
import logging
import re
@@ -22,6 +25,7 @@ from re import Pattern
from typing import Any, Optional, TYPE_CHECKING, TypedDict
from urllib import parse
+import requests
from apispec import APISpec
from apispec.ext.marshmallow import MarshmallowPlugin
from cryptography.hazmat.backends import default_backend
@@ -29,9 +33,11 @@ from cryptography.hazmat.primitives import serialization
from flask import current_app
from flask_babel import gettext as __
from marshmallow import fields, Schema
+from requests.auth import HTTPBasicAuth
from sqlalchemy import types
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.engine.url import URL
+from sqlalchemy.exc import ProgrammingError
from superset.constants import TimeGrain, USER_AGENT
from superset.databases.utils import make_url_safe
@@ -39,6 +45,7 @@ from superset.db_engine_specs.base import BaseEngineSpec, BasicPropertiesType
from superset.db_engine_specs.postgres import PostgresBaseEngineSpec
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.models.sql_lab import Query
+from superset.superset_typing import OAuth2ClientConfig, OAuth2TokenResponse
if TYPE_CHECKING:
from superset.models.core import Database
@@ -87,6 +94,11 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
supports_dynamic_schema = True
supports_catalog = True
+ supports_oauth2 = True
+ oauth2_scope = "refresh_token session:role:SYSADMIN"
+ oauth2_authorization_request_uri = None
+ oauth2_token_request_uri = None
+
_time_grain_expressions = {
None: "{col}",
TimeGrain.SECOND: "DATE_TRUNC('SECOND', {col})",
@@ -391,3 +403,75 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
f"must be listed in 'ALLOWED_EXTRA_AUTHENTICATIONS' config"
)
connect_args["auth"] = snowflake_auth(**auth_params)
+
+ @classmethod
+ def update_impersonation_config(
+ cls,
+ connect_args: dict[str, Any],
+ uri: str,
+ username: str | None,
+ access_token: str | None,
+ ) -> None:
+ if access_token:
+ connect_args.update(
+ {
+ "authenticator": "oauth",
+ "token": access_token,
+ },
+ )
+
+ @classmethod
+ def execute(
+ cls,
+ cursor: Any,
+ query: str,
+ database: Database,
+ **kwargs: Any,
+ ) -> None:
+ try:
+ cursor.execute(query)
+ except ProgrammingError as ex:
+ if database.is_oauth2_enabled() and "User is empty" in str(ex):
+ cls.start_oauth2_dance(database)
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+ except Exception as ex:
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+
+ @classmethod
+ def get_oauth2_token(
+ cls,
+ config: OAuth2ClientConfig,
+ code: str,
+ ) -> OAuth2TokenResponse:
+ timeout = current_app.config["DATABASE_OAUTH2_TIMEOUT"].total_seconds()
+ uri = config["token_request_uri"]
+ response = requests.post(
+ uri,
+ data={
+ "code": code,
+ "redirect_uri": config["redirect_uri"],
+ "grant_type": "authorization_code",
+ },
+ auth=HTTPBasicAuth(config["id"], config["secret"]),
+ timeout=timeout,
+ )
+ return response.json()
+
+ @classmethod
+ def get_oauth2_fresh_token(
+ cls,
+ config: OAuth2ClientConfig,
+ refresh_token: str,
+ ) -> OAuth2TokenResponse:
+ timeout = current_app.config["DATABASE_OAUTH2_TIMEOUT"].total_seconds()
+ uri = config["token_request_uri"]
+ response = requests.post(
+ uri,
+ data={
+ "refresh_token": refresh_token,
+ "grant_type": "authorization_code",
+ },
+ auth=HTTPBasicAuth(config["id"], config["secret"]),
+ timeout=timeout,
+ )
+ return response.json()
(superset) 01/02: feat: read OAuth2 info from database extra
Posted by be...@apache.org.
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 <ro...@dealmeida.net>
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)