Adding tests

superset/commands/database/update.py

@@ -151,6 +151,9 @@ def _get_schema_names(
                 catalog=catalog,
                 ssh_tunnel=ssh_tunnel,
             )
+        except OAuth2RedirectError:
+            # raise OAuth2 exceptions as-is
+            raise
         except GenericDBException as ex:
             raise DatabaseConnectionFailedError() from ex
 

superset/exceptions.py

@@ -329,6 +329,8 @@ class OAuth2RedirectError(SupersetErrorException):
 
     See the `OAuth2RedirectMessage.tsx` component for more details of how this
     information is handled.
+
+    TODO (betodealmeida): change status to 403.
     """
 
     def __init__(self, url: str, tab_id: str, redirect_uri: str):

tests/unit_tests/commands/databases/create_test.py

@@ -21,6 +21,7 @@
 from pytest_mock import MockerFixture
 
 from superset.commands.database.create import CreateDatabaseCommand
+from superset.exceptions import OAuth2RedirectError
 from superset.extensions import security_manager
 
 
@@ -124,3 +125,33 @@ def test_create_permissions_without_catalog(
         ],
         any_order=True,
     )
+
+
+def test_create_with_oauth2(
+    mocker: MockerFixture,
+    database_without_catalog: MockerFixture,
+) -> None:
+    """
+    Test that the database can be created even if OAuth2 is needed to connect.
+    """
+    TestConnectionDatabaseCommand = mocker.patch(
+        "superset.commands.database.create.TestConnectionDatabaseCommand"
+    )
+    TestConnectionDatabaseCommand().run.side_effect = OAuth2RedirectError(
+        "url",
+        "tab_id",
+        "redirect_uri",
+    )
+    add_permission_view_menu = mocker.patch.object(
+        security_manager,
+        "add_permission_view_menu",
+    )
+
+    CreateDatabaseCommand(
+        {
+            "database_name": "test_database",
+            "sqlalchemy_uri": "sqlite://",
+        }
+    ).run()
+
+    add_permission_view_menu.assert_not_called()

tests/unit_tests/commands/databases/test_connection_test.py

@@ -0,0 +1,91 @@
+# 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 pytest
+from pytest_mock import MockerFixture
+
+from superset.commands.database.test_connection import TestConnectionDatabaseCommand
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import OAuth2RedirectError
+
+
+def test_command(mocker: MockerFixture) -> None:
+    """
+    Test the happy path of the command.
+    """
+    user = mocker.MagicMock()
+    user.email = "[email protected]"
+    mocker.patch("superset.db_engine_specs.gsheets.g", user=user)
+    mocker.patch("superset.db_engine_specs.gsheets.create_engine")
+
+    database = mocker.MagicMock()
+    database.db_engine_spec.__name__ = "GSheetsEngineSpec"
+    with database.get_sqla_engine() as engine:
+        engine.dialect.do_ping.return_value = True
+
+    DatabaseDAO = mocker.patch("superset.commands.database.test_connection.DatabaseDAO")
+    DatabaseDAO.build_db_for_connection_test.return_value = database
+
+    properties = {
+        "sqlalchemy_uri": "gsheets://",
+        "engine": "gsheets",
+        "driver": "gsheets",
+        "catalog": {"test": "https://example.org/"},
+    }
+    command = TestConnectionDatabaseCommand(properties)
+    command.run()
+
+
+def test_command_with_oauth2(mocker: MockerFixture) -> None:
+    """
+    Test the command when OAuth2 is needed.
+    """
+    user = mocker.MagicMock()
+    user.email = "[email protected]"
+    mocker.patch("superset.db_engine_specs.gsheets.g", user=user)
+    mocker.patch("superset.db_engine_specs.gsheets.create_engine")
+
+    database = mocker.MagicMock()
+    database.is_oauth2_enabled.return_value = True
+    database.db_engine_spec.needs_oauth2.return_value = True
+    database.start_oauth2_dance.side_effect = OAuth2RedirectError(
+        "url",
+        "tab_id",
+        "redirect_uri",
+    )
+    database.db_engine_spec.__name__ = "GSheetsEngineSpec"
+    with database.get_sqla_engine() as engine:
+        engine.dialect.do_ping.side_effect = Exception("OAuth2 needed")
+
+    DatabaseDAO = mocker.patch("superset.commands.database.test_connection.DatabaseDAO")
+    DatabaseDAO.build_db_for_connection_test.return_value = database
+
+    properties = {
+        "sqlalchemy_uri": "gsheets://",
+        "engine": "gsheets",
+        "driver": "gsheets",
+        "catalog": {"test": "https://example.org/"},
+    }
+    command = TestConnectionDatabaseCommand(properties)
+    with pytest.raises(OAuth2RedirectError) as excinfo:
+        command.run()
+    assert excinfo.value.error == SupersetError(
+        message="You don't have permission to access the data.",
+        error_type=SupersetErrorType.OAUTH2_REDIRECT,
+        level=ErrorLevel.WARNING,
+        extra={"url": "url", "tab_id": "tab_id", "redirect_uri": "redirect_uri"},
+    )

tests/unit_tests/commands/databases/update_test.py

@@ -21,6 +21,7 @@
 from pytest_mock import MockerFixture
 
 from superset.commands.database.update import UpdateDatabaseCommand
+from superset.exceptions import OAuth2RedirectError
 from superset.extensions import security_manager
 
 
@@ -57,6 +58,24 @@ def database_without_catalog(mocker: MockerFixture) -> MagicMock:
     return database
 
 
+@pytest.fixture()
+def database_needs_oauth2(mocker: MockerFixture) -> MagicMock:
+    """
+    Mock a database without catalogs that needs OAuth2.
+    """
+    database = mocker.MagicMock()
+    database.database_name = "my_db"
+    database.db_engine_spec.__name__ = "test_engine"
+    database.db_engine_spec.supports_catalog = False
+    database.get_all_schema_names.side_effect = OAuth2RedirectError(
+        "url",
+        "tab_id",
+        "redirect_uri",
+    )
+
+    return database
+
+
 def test_update_with_catalog(
     mocker: MockerFixture,
     database_with_catalog: MockerFixture,
@@ -276,3 +295,32 @@ def test_rename_without_catalog(
     )
 
     assert schema2_pvm.view_menu.name == "[my_other_db].[schema2]"
+
+
+def test_update_with_oauth2(
+    mocker: MockerFixture,
+    database_needs_oauth2: MockerFixture,
+) -> None:
+    """
+    Test that the database can be updated even if OAuth2 is needed to connect.
+    """
+    DatabaseDAO = mocker.patch("superset.commands.database.update.DatabaseDAO")
+    DatabaseDAO.find_by_id.return_value = database_needs_oauth2
+    DatabaseDAO.update.return_value = database_needs_oauth2
+
+    find_permission_view_menu = mocker.patch.object(
+        security_manager,
+        "find_permission_view_menu",
+    )
+    find_permission_view_menu.side_effect = [
+        None,  # schema1 has no permissions
+        "[my_db].[schema2]",  # second schema already exists
+    ]
+    add_permission_view_menu = mocker.patch.object(
+        security_manager,
+        "add_permission_view_menu",
+    )
+
+    UpdateDatabaseCommand(1, {}).run()
+
+    add_permission_view_menu.assert_not_called()