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/06/05 19:04:03 UTC
(superset) branch db-oauth2-client-info updated (bbc0956585 -> e4998d823a)
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
discard bbc0956585 WIP
new e4998d823a WIP
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (bbc0956585)
\
N -- N -- N refs/heads/db-oauth2-client-info (e4998d823a)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 1 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.
Summary of changes:
superset/db_engine_specs/snowflake.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
(superset) 01/01: 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 e4998d823a1213c556096ac0f6597fa33be70b6d
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Wed May 8 10:18:54 2024 -0400
WIP
---
superset/commands/database/update.py | 4 ++
superset/databases/api.py | 1 -
superset/db_engine_specs/base.py | 18 +++--
superset/db_engine_specs/snowflake.py | 128 +++++++++++++++++++++++++++++++---
superset/models/core.py | 14 ++--
superset/sql_lab.py | 9 +++
superset/utils/oauth2.py | 2 +-
7 files changed, 152 insertions(+), 24 deletions(-)
diff --git a/superset/commands/database/update.py b/superset/commands/database/update.py
index 5e0968954c..31c1a1e7cf 100644
--- a/superset/commands/database/update.py
+++ b/superset/commands/database/update.py
@@ -133,6 +133,10 @@ class UpdateDatabaseCommand(BaseCommand):
try:
schemas = database.get_all_schema_names(ssh_tunnel=ssh_tunnel)
except Exception as ex:
+ # XXX conditional
+ if 1:
+ db.session.commit()
+ database.db_engine_spec.start_oauth2_dance(database)
db.session.rollback()
raise DatabaseConnectionFailedError() from ex
diff --git a/superset/databases/api.py b/superset/databases/api.py
index a77019123b..2fca84a357 100644
--- a/superset/databases/api.py
+++ b/superset/databases/api.py
@@ -453,7 +453,6 @@ class DatabaseRestApi(BaseSupersetModelRestApi):
@expose("/<int:pk>", methods=("PUT",))
@protect()
- @safe
@statsd_metrics
@event_logger.log_this_with_context(
action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.put",
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 3cc1315129..65a4cfbd64 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -131,7 +131,9 @@ builtin_time_grains: dict[str | None, str] = {
}
-class TimestampExpression(ColumnClause): # pylint: disable=abstract-method, too-many-ancestors
+class TimestampExpression(
+ ColumnClause
+): # pylint: disable=abstract-method, too-many-ancestors
def __init__(self, expr: str, col: ColumnClause, **kwargs: Any) -> None:
"""Sqlalchemy class that can be used to render native column elements respecting
engine-specific quoting rules as part of a string-based expression.
@@ -389,9 +391,9 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
max_column_name_length: int | None = None
try_remove_schema_from_table_name = True # pylint: disable=invalid-name
run_multiple_statements_as_one = False
- custom_errors: dict[
- Pattern[str], tuple[str, SupersetErrorType, dict[str, Any]]
- ] = {}
+ custom_errors: dict[Pattern[str], tuple[str, SupersetErrorType, dict[str, Any]]] = (
+ {}
+ )
# Whether the engine supports file uploads
# if True, database will be listed as option in the upload file form
@@ -1597,9 +1599,11 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
@classmethod
def _get_fields(cls, cols: list[ResultSetColumnType]) -> list[Any]:
return [
- literal_column(query_as)
- if (query_as := c.get("query_as"))
- else column(c["column_name"])
+ (
+ literal_column(query_as)
+ if (query_as := c.get("query_as"))
+ else column(c["column_name"])
+ )
for c in cols
]
diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py
index 83d382cda1..baeb4f4762 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,16 +33,20 @@ 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 import security_manager
from superset.constants import TimeGrain, USER_AGENT
from superset.databases.utils import make_url_safe
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
@@ -57,12 +65,12 @@ logger = logging.getLogger(__name__)
class SnowflakeParametersSchema(Schema):
- username = fields.Str(required=True)
- password = fields.Str(required=True)
+ username = fields.Str(required=False)
+ password = fields.Str(required=False)
account = fields.Str(required=True)
database = fields.Str(required=True)
- role = fields.Str(required=True)
- warehouse = fields.Str(required=True)
+ role = fields.Str(required=False)
+ warehouse = fields.Str(required=False)
class SnowflakeParametersType(TypedDict):
@@ -87,6 +95,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})",
@@ -123,17 +136,29 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
),
}
- @staticmethod
- def get_extra_params(database: "Database") -> dict[str, Any]:
+ @classmethod
+ def get_extra_params(cls, database: "Database") -> dict[str, Any]:
"""
Add a user agent to be used in the requests.
"""
extra: dict[str, Any] = BaseEngineSpec.get_extra_params(database)
engine_params: dict[str, Any] = extra.setdefault("engine_params", {})
connect_args: dict[str, Any] = engine_params.setdefault("connect_args", {})
-
connect_args.setdefault("application", USER_AGENT)
+ # populate OAuth2 URLs if not set, since they can be inferred from the account
+ if oauth2_client_info := extra.get("oauth2_client_info"):
+ account = database.url_object.host
+ oauth2_client_info.setdefault(
+ "authorization_request_uri",
+ f"https://{account}.snowflakecomputing.com/oauth/authorize",
+ )
+ oauth2_client_info.setdefault(
+ "token_request_uri",
+ f"https://{account}.snowflakecomputing.com/oauth/token-request",
+ )
+ oauth2_client_info.setdefault("scope", cls.oauth2_scope)
+
return extra
@classmethod
@@ -303,11 +328,9 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
) -> list[SupersetError]:
errors: list[SupersetError] = []
required = {
- "warehouse",
"username",
"database",
"account",
- "role",
"password",
}
parameters = properties.get("parameters", {})
@@ -391,3 +414,90 @@ 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 and not security_manager.is_admin():
+ connect_args.update(
+ {
+ "authenticator": "oauth",
+ "token": access_token,
+ },
+ )
+
+ @classmethod
+ def get_url_for_impersonation(
+ cls,
+ url: URL,
+ impersonate_user: bool,
+ username: str | None,
+ access_token: str | None, # pylint: disable=unused-argument
+ ) -> URL:
+ # force OAuth2
+ if impersonate_user and not security_manager.is_admin():
+ url = url._replace(username="", password="", query="")
+
+ return url
+
+ @classmethod
+ def execute(
+ cls,
+ cursor: Any,
+ query: str,
+ database: Database,
+ **kwargs: Any,
+ ) -> None:
+ try:
+ cursor.execute(query)
+ except ProgrammingError as ex:
+ # refactor into base class method needs_oauth2
+ 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": "refresh_token",
+ },
+ auth=HTTPBasicAuth(config["id"], config["secret"]),
+ timeout=timeout,
+ )
+ return response.json()
diff --git a/superset/models/core.py b/superset/models/core.py
index 79309cdb3d..59c5cf3136 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -117,7 +117,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"
@@ -378,9 +380,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
@@ -1027,7 +1027,8 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
admins to create custom OAuth2 clients from the Superset UI, and assign them to
specific databases.
"""
- oauth2_client_info = self.get_extra().get("oauth2_client_info", {})
+ 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:
@@ -1039,7 +1040,8 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
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"):
+ 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)
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 3f8c1cc737..95719a76a1 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -186,6 +186,15 @@ def get_sql_results( # pylint: disable=too-many-arguments
log_params=log_params,
)
except Exception as ex: # pylint: disable=broad-except
+ query = get_query(query_id)
+ database = query.database
+ print("\n\nBETO 456")
+ print(ex)
+ try:
+ database.db_engine_spec.start_oauth2_dance(database)
+ except OAuth2RedirectError as ex:
+ return handle_query_error(ex, query)
+
logger.debug("Query %d: %s", query_id, ex)
stats_logger.incr("error_sqllab_unhandled")
query = get_query(query_id)
diff --git a/superset/utils/oauth2.py b/superset/utils/oauth2.py
index 7b7440b059..bc4805fd81 100644
--- a/superset/utils/oauth2.py
+++ b/superset/utils/oauth2.py
@@ -188,7 +188,7 @@ class OAuth2ClientConfigSchema(Schema):
scope = fields.String(required=True)
redirect_uri = fields.String(
required=False,
- load_default=url_for("DatabaseRestApi.oauth2", _external=True),
+ load_default=lambda: url_for("DatabaseRestApi.oauth2", _external=True),
)
authorization_request_uri = fields.String(required=True)
token_request_uri = fields.String(required=True)