feat: OAuth2 client initial work (#29109)

tests/unit_tests/sql_lab_test.py

@@ -16,12 +16,22 @@
 # under the License.
 # pylint: disable=import-outside-toplevel, invalid-name, unused-argument, too-many-locals
 
+import json
+from uuid import UUID
+
 import sqlparse
+from freezegun import freeze_time
 from pytest_mock import MockerFixture
 from sqlalchemy.orm.session import Session
 
 from superset import db
+from superset.common.db_query_status import QueryStatus
+from superset.errors import ErrorLevel, SupersetErrorType
+from superset.exceptions import OAuth2Error
+from superset.models.core import Database
+from superset.sql_lab import get_sql_results
 from superset.utils.core import override_user
+from tests.unit_tests.models.core_test import oauth2_client_info
 
 
 def test_execute_sql_statement(mocker: MockerFixture, app: None) -> None:
@@ -218,3 +228,55 @@ def test_sql_lab_insert_rls_as_subquery(
         query.executed_sql
         == "SELECT c FROM (SELECT * FROM t WHERE (t.c > 5)) AS t\nLIMIT 6"
     )
+
+
+@freeze_time("2021-04-01T00:00:00Z")
+def test_get_sql_results_oauth2(mocker: MockerFixture, app) -> None:
+    """
+    Test that `get_sql_results` works with OAuth2.
+    """
+    app_context = app.test_request_context()
+    app_context.push()
+
+    mocker.patch(
+        "superset.db_engine_specs.base.uuid4",
+        return_value=UUID("fb11f528-6eba-4a8a-837e-6b0d39ee9187"),
+    )
+
+    g = mocker.patch("superset.db_engine_specs.base.g")
+    g.user = mocker.MagicMock()
+    g.user.id = 42
+
+    database = Database(
+        id=1,
+        database_name="my_db",
+        sqlalchemy_uri="sqlite://",
+        encrypted_extra=json.dumps(oauth2_client_info),
+    )
+    database.db_engine_spec.oauth2_exception = OAuth2Error  # type: ignore
+    get_sqla_engine = mocker.patch.object(database, "get_sqla_engine")
+    get_sqla_engine().__enter__().raw_connection.side_effect = OAuth2Error(
+        "OAuth2 required"
+    )
+
+    query = mocker.MagicMock()
+    query.database = database
+    mocker.patch("superset.sql_lab.get_query", return_value=query)
+
+    payload = get_sql_results(query_id=1, rendered_query="SELECT 1")
+    assert payload == {
+        "status": QueryStatus.FAILED,
+        "error": "You don't have permission to access the data.",
+        "errors": [
+            {
+                "message": "You don't have permission to access the data.",
+                "error_type": SupersetErrorType.OAUTH2_REDIRECT,
+                "level": ErrorLevel.WARNING,
+                "extra": {
+                    "url": "https://abcd1234.snowflakecomputing.com/oauth/authorize?scope=refresh_token+session%3Arole%3ASYSADMIN&access_type=offline&include_granted_scopes=false&response_type=code&state=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9%252EeyJleHAiOjE2MTcyMzU1MDAsImRhdGFiYXNlX2lkIjoxLCJ1c2VyX2lkIjo0MiwiZGVmYXVsdF9yZWRpcmVjdF91cmkiOiJodHRwOi8vZXhhbXBsZS5jb20vYXBpL3YxL2RhdGFiYXNlL29hdXRoMi8iLCJ0YWJfaWQiOiJmYjExZjUyOC02ZWJhLTRhOGEtODM3ZS02YjBkMzllZTkxODcifQ%252Ec_m_35xwwSrLgCXwV4aKO1928flOEFQIqqg9ctiXjcM&redirect_uri=http%3A%2F%2Fexample.com%2Fapi%2Fv1%2Fdatabase%2Foauth2%2F&client_id=my_client_id&prompt=consent",
+                    "tab_id": "fb11f528-6eba-4a8a-837e-6b0d39ee9187",
+                    "redirect_uri": "http://example.com/api/v1/database/oauth2/",
+                },
+            }
+        ],
+    }