feat: use personal access token in OAuth2 databases

superset-frontend/src/components/ErrorMessage/OAuth2RedirectMessage.tsx

@@ -0,0 +1,65 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import React from 'react';
+import { t } from '@superset-ui/core';
+
+import { ErrorMessageComponentProps } from './types';
+import ErrorAlert from './ErrorAlert';
+
+interface OAuth2RedirectExtra {
+  url: string;
+}
+
+function OAuth2RedirectMessage({
+  error,
+  source = 'sqllab',
+}: ErrorMessageComponentProps<OAuth2RedirectExtra>) {
+  const { extra, level } = error;
+
+  const body = (
+    <p>
+      This database uses OAuth2 for authentication, and will store your personal
+      access tokens after authentication so that only you can use it. When you
+      click the link above you will be asked to grant access to the data in a
+      new window. After confirming, you can close the window and re-run the
+      query here.
+    </p>
+  );
+  const subtitle = (
+    <>
+      You need to{' '}
+      <a href={extra.url} target="_blank" rel="noreferrer">
+        provide authorization
+      </a>{' '}
+      in order to run this query.
+    </>
+  );
+
+  return (
+    <ErrorAlert
+      title={t('Authorization needed')}
+      subtitle={subtitle}
+      level={level}
+      source={source}
+      body={body}
+    />
+  );
+}
+
+export default OAuth2RedirectMessage;

superset-frontend/src/components/ErrorMessage/types.ts

@@ -56,6 +56,9 @@ export const ErrorTypeEnum = {
   QUERY_SECURITY_ACCESS_ERROR: 'QUERY_SECURITY_ACCESS_ERROR',
   MISSING_OWNERSHIP_ERROR: 'MISSING_OWNERSHIP_ERROR',
 
+  // Redirect user to start OAuth2
+  OAUTH2_REDIRECT: 'OAUTH2_REDIRECT',
+
   // Other errors
   BACKEND_TIMEOUT_ERROR: 'BACKEND_TIMEOUT_ERROR',
   DATABASE_NOT_FOUND_ERROR: 'DATABASE_NOT_FOUND_ERROR',

superset-frontend/src/setup/setupErrorMessages.ts

@@ -22,6 +22,7 @@ import TimeoutErrorMessage from 'src/components/ErrorMessage/TimeoutErrorMessage
 import DatabaseErrorMessage from 'src/components/ErrorMessage/DatabaseErrorMessage';
 import ParameterErrorMessage from 'src/components/ErrorMessage/ParameterErrorMessage';
 import DatasetNotFoundErrorMessage from 'src/components/ErrorMessage/DatasetNotFoundErrorMessage';
+import OAuth2RedirectMessage from 'src/components/ErrorMessage/OAuth2RedirectMessage';
 
 import setupErrorMessagesExtra from './setupErrorMessagesExtra';
 
@@ -144,5 +145,9 @@ export default function setupErrorMessages() {
     ErrorTypeEnum.FAILED_FETCHING_DATASOURCE_INFO_ERROR,
     DatasetNotFoundErrorMessage,
   );
+  errorMessageComponentRegistry.registerValue(
+    ErrorTypeEnum.OAUTH2_REDIRECT,
+    OAuth2RedirectMessage,
+  );
   setupErrorMessagesExtra();
 }

superset/config.py

@@ -333,6 +333,12 @@ def _try_json_readsha(filepath: str, length: int) -> Optional[str]:
 #    { 'name': 'Yahoo', 'url': 'https://open.login.yahoo.com/' },
 #    { 'name': 'Flickr', 'url': 'https://www.flickr.com/<username>' },
 
+# Settings to enable OAuth2 in the Google Sheets DB engine spec.
+# GSHEETS_OAUTH2_CLIENT_ID = "XXX"
+# GSHEETS_OAUTH2_CLIENT_SECRET = "YYY"
+# GSHEETS_OAUTH2_REDIRECT_URI = "http://localhost:8088/api/v1/database/oauth2/"
+
+
 # ---------------------------------------------------
 # Roles config
 # ---------------------------------------------------

superset/connectors/sqla/utils.py

@@ -142,7 +142,7 @@ def get_virtual_table_metadata(dataset: SqlaTable) -> List[ResultSetColumnType]:
             with closing(engine.raw_connection()) as conn:
                 cursor = conn.cursor()
                 query = dataset.database.apply_limit_to_sql(statements[0], limit=1)
-                db_engine_spec.execute(cursor, query)
+                db_engine_spec.execute(cursor, query, dataset.database.id)
                 result = db_engine_spec.fetch_data(cursor, limit=1)
                 result_set = SupersetResultSet(
                     result, cursor.description, db_engine_spec
@@ -164,7 +164,7 @@ def get_columns_description(
                 cursor = conn.cursor()
                 query = database.apply_limit_to_sql(query, limit=1)
                 cursor.execute(query)
-                db_engine_spec.execute(cursor, query)
+                db_engine_spec.execute(cursor, query, database.id)
                 result = db_engine_spec.fetch_data(cursor, limit=1)
                 result_set = SupersetResultSet(
                     result, cursor.description, db_engine_spec

superset/databases/api.py

@@ -17,18 +17,27 @@
 # pylint: disable=too-many-lines
 import json
 import logging
-from datetime import datetime
+from datetime import datetime, timedelta
 from io import BytesIO
 from typing import Any, cast, Dict, List, Optional
 from zipfile import is_zipfile, ZipFile
 
-from flask import request, Response, send_file
+import jwt
+from flask import (
+    current_app,
+    g,
+    make_response,
+    render_template,
+    request,
+    Response,
+    send_file,
+)
 from flask_appbuilder.api import expose, protect, rison, safe
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 from marshmallow import ValidationError
 from sqlalchemy.exc import NoSuchTableError, OperationalError, SQLAlchemyError
 
-from superset import app, event_logger
+from superset import app, db, event_logger
 from superset.commands.importers.exceptions import (
     IncorrectFormatError,
     NoValidFilesFoundError,
@@ -77,7 +86,7 @@
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import SupersetErrorsException, SupersetException
 from superset.extensions import security_manager
-from superset.models.core import Database
+from superset.models.core import Database, DatabaseUserOAuth2Tokens
 from superset.superset_typing import FlaskResponse
 from superset.utils.core import error_msg_from_exception, parse_js_uri_path_item
 from superset.views.base import json_errors_response
@@ -107,6 +116,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi):
         "available",
         "validate_parameters",
         "validate_sql",
+        "oauth2",
     }
     resource_name = "database"
     class_permission_name = "Database"
@@ -855,6 +865,69 @@ def validate_sql(self, pk: int) -> FlaskResponse:
         except DatabaseNotFoundError:
             return self.response_404()
 
+    @expose("/oauth2/", methods=["GET"])
+    @event_logger.log_this_with_context(
+        action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.oauth",
+        log_to_statsd=True,
+    )
+    def oauth2(self) -> FlaskResponse:
+        """
+        ---
+        get:
+          summary: >-
+            Receive user-level authentication tokens from OAuth
+          description: ->
+            Receive and store user-level authentication tokens from OAuth
+          parameters:
+          - in: query
+            name: state
+          - in: query
+            name: code
+          - in: query
+            name: scope
+          - in: query
+            name: error
+          responses:
+            200:
+              description: A dummy self-closing HTML page
+              content:
+                text/html:
+                  schema:
+                    type: string
+            400:
+              $ref: '#/components/responses/400'
+            500:
+              $ref: '#/components/responses/500'
+        """
+        parameters = request.args.to_dict()
+        if "error" in parameters:
+            raise Exception(parameters["error"])  # XXX: raise custom SIP-40 exception
+
+        payload = jwt.decode(
+            parameters["state"].replace("%2E", "."),
+            current_app.config["SECRET_KEY"],
+            algorithms=["HS256"],
+        )
+
+        # exchange code for access/refresh tokens
+        database = db.session.query(Database).filter_by(id=payload["database_id"]).one()
+        token_response = database.db_engine_spec.get_oauth2_token(parameters["code"])
+
+        # store tokens
+        token = DatabaseUserOAuth2Tokens(
+            user_id=payload["user_id"],
+            database_id=database.id,
+            access_token=token_response["access_token"],
+            access_token_expiration=datetime.now()
+            + timedelta(seconds=token_response["expires_in"]),
+            refresh_token=token_response.get("refresh_token"),
+        )
+        db.session.add(token)
+        db.session.commit()
+
+        # return blank page that closes itself
+        return make_response(render_template("superset/close.html"), 200)
+
     @expose("/export/", methods=["GET"])
     @protect()
     @safe

superset/db_engine_specs/base.py

@@ -62,6 +62,8 @@
 from superset import security_manager, sql_parse
 from superset.databases.utils import make_url_safe
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import OAuth2RedirectError
+from superset.models.sql_lab import Query
 from superset.sql_parse import ParsedQuery, Table
 from superset.superset_typing import ResultSetColumnType
 from superset.utils import core as utils
@@ -170,6 +172,30 @@ class MetricType(TypedDict, total=False):
     extra: Optional[str]
 
 
+class OAuth2TokenResponse(TypedDict, total=False):
+    """
+    Type for an OAuth2 response when exchanging or refreshing tokens.
+    """
+
+    access_token: str
+    expires_in: int
+    scope: str
+    token_type: str
+
+    # only present when exchanging code for refresh/access tokens
+    refresh_token: str
+
+
+class OAuth2State(TypedDict):
+    """
+    Type for the state passed during OAuth2.
+    """
+
+    database_id: int
+    user_id: int
+    default_redirect_uri: str
+
+
 class BaseEngineSpec:  # pylint: disable=too-many-public-methods
     """Abstract class for database engine specific configurations
 
@@ -1308,14 +1334,19 @@ def estimate_query_cost(
 
     @classmethod
     def get_url_for_impersonation(
-        cls, url: URL, impersonate_user: bool, username: Optional[str]
+        cls,
+        url: URL,
+        impersonate_user: bool,
+        username: Optional[str],
+        access_token: Optional[str] = None,  # pylint: disable=unused-argument
     ) -> URL:
         """
         Return a modified URL with the username set.
 
         :param url: SQLAlchemy URL object
         :param impersonate_user: Flag indicating if impersonation is enabled
         :param username: Effective username
+        :param access_token: OAuth2 token for the user
         """
         if impersonate_user and username is not None:
             url = url.set(username=username)
@@ -1344,6 +1375,7 @@ def execute(  # pylint: disable=unused-argument
         cls,
         cursor: Any,
         query: str,
+        database_id: int,
         **kwargs: Any,
     ) -> None:
         """
@@ -1361,6 +1393,12 @@ def execute(  # pylint: disable=unused-argument
             cursor.arraysize = cls.arraysize
         try:
             cursor.execute(query)
+        except cls.oauth2_exception as ex:
+            if cls.is_oauth2_enabled():
+                oauth_url = cls.get_oauth2_authorization_uri(database_id)
+                raise OAuth2RedirectError(oauth_url) from ex
+
+            raise cls.get_dbapi_mapped_exception(ex) from ex
         except Exception as ex:
             raise cls.get_dbapi_mapped_exception(ex) from ex
 
@@ -1690,6 +1728,34 @@ def get_public_information(cls) -> Dict[str, Any]:
             "supports_file_upload": cls.supports_file_upload,
         }
 
+    # Driver-specific exception that should be mapped to OAuth2RedirectError
+    oauth2_exception = OAuth2RedirectError
+
+    @staticmethod
+    def is_oauth2_enabled() -> bool:
+        return False
+
+    @staticmethod
+    def get_oauth2_authorization_uri(database_id: int) -> str:
+        """
+        Return URI for initial OAuth2 request.
+        """
+        raise NotImplementedError()
+
+    @staticmethod
+    def get_oauth2_token(code: str) -> OAuth2TokenResponse:
+        """
+        Exchange authorization code for refresh/access tokens.
+        """
+        raise NotImplementedError()
+
+    @staticmethod
+    def get_oauth2_fresh_token(refresh_token: str) -> OAuth2TokenResponse:
+        """
+        Refresh an access token that has expired.
+        """
+        raise NotImplementedError()
+
 
 # schema for adding a database by providing parameters instead of the
 # full SQLAlchemy URI

superset/db_engine_specs/drill.py

@@ -76,7 +76,11 @@ def adjust_database_uri(cls, uri: URL, selected_schema: Optional[str]) -> URL:
 
     @classmethod
     def get_url_for_impersonation(
-        cls, url: URL, impersonate_user: bool, username: Optional[str]
+        cls,
+        url: URL,
+        impersonate_user: bool,
+        username: Optional[str],
+        access_token: Optional[str] = None,
     ) -> URL:
         """
         Return a modified URL with the username set.