chore: improve schema security (#23385)

tests/unit_tests/db_engine_specs/test_drill.py

@@ -20,6 +20,7 @@ from datetime import datetime
 from typing import Optional
 
 import pytest
+from sqlalchemy.engine.url import make_url
 
 from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
 from tests.unit_tests.fixtures.common import dttm
@@ -106,3 +107,18 @@ def test_convert_dttm(
     from superset.db_engine_specs.drill import DrillEngineSpec as spec
 
     assert_convert_dttm(spec, target_type, expected_result, dttm)
+
+
+def test_get_schema_from_engine_params() -> None:
+    """
+    Test ``get_schema_from_engine_params``.
+    """
+    from superset.db_engine_specs.drill import DrillEngineSpec
+
+    assert (
+        DrillEngineSpec.get_schema_from_engine_params(
+            make_url("drill+sadrill://localhost:8047/dfs/test?use_ssl=False"),
+            {},
+        )
+        == "dfs.test"
+    )

tests/unit_tests/db_engine_specs/test_hive.py

@@ -20,6 +20,7 @@ from datetime import datetime
 from typing import Optional
 
 import pytest
+from sqlalchemy.engine.url import make_url
 
 from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
 from tests.unit_tests.fixtures.common import dttm
@@ -42,3 +43,17 @@ def test_convert_dttm(
     from superset.db_engine_specs.hive import HiveEngineSpec as spec
 
     assert_convert_dttm(spec, target_type, expected_result, dttm)
+
+
+def test_get_schema_from_engine_params() -> None:
+    """
+    Test the ``get_schema_from_engine_params`` method.
+    """
+    from superset.db_engine_specs.hive import HiveEngineSpec
+
+    assert (
+        HiveEngineSpec.get_schema_from_engine_params(
+            make_url("hive://localhost:10000/default"), {}
+        )
+        == "default"
+    )

tests/unit_tests/db_engine_specs/test_mysql.py

@@ -148,3 +148,17 @@ def test_cancel_query_failed(engine_mock: Mock) -> None:
     query = Query()
     cursor_mock = engine_mock.raiseError.side_effect = Exception()
     assert MySQLEngineSpec.cancel_query(cursor_mock, query, "123") is False
+
+
+def test_get_schema_from_engine_params() -> None:
+    """
+    Test the ``get_schema_from_engine_params`` method.
+    """
+    from superset.db_engine_specs.mysql import MySQLEngineSpec
+
+    assert (
+        MySQLEngineSpec.get_schema_from_engine_params(
+            make_url("mysql://user:password@host/db1"), {}
+        )
+        == "db1"
+    )

tests/unit_tests/db_engine_specs/test_postgres.py

@@ -21,6 +21,7 @@ from typing import Any, Dict, Optional, Type
 import pytest
 from sqlalchemy import types
 from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION, ENUM, JSON
+from sqlalchemy.engine.url import make_url
 
 from superset.utils.core import GenericDataType
 from tests.unit_tests.db_engine_specs.utils import (
@@ -89,3 +90,44 @@ def test_get_column_spec(
     from superset.db_engine_specs.postgres import PostgresEngineSpec as spec
 
     assert_column_spec(spec, native_type, sqla_type, attrs, generic_type, is_dttm)
+
+
+def test_get_schema_from_engine_params() -> None:
+    """
+    Test the ``get_schema_from_engine_params`` method.
+    """
+    from superset.db_engine_specs.postgres import PostgresEngineSpec
+
+    assert (
+        PostgresEngineSpec.get_schema_from_engine_params(
+            make_url("postgresql://user:password@host/db1"), {}
+        )
+        is None
+    )
+
+    assert (
+        PostgresEngineSpec.get_schema_from_engine_params(
+            make_url("postgresql://user:password@host/db1"),
+            {"options": "-csearch_path=secret"},
+        )
+        == "secret"
+    )
+
+    assert (
+        PostgresEngineSpec.get_schema_from_engine_params(
+            make_url("postgresql://user:password@host/db1"),
+            {"options": "-c search_path = secret -cfoo=bar -c debug"},
+        )
+        == "secret"
+    )
+
+    with pytest.raises(Exception) as excinfo:
+        PostgresEngineSpec.get_schema_from_engine_params(
+            make_url("postgresql://user:password@host/db1"),
+            {"options": "-csearch_path=secret,public"},
+        )
+    assert str(excinfo.value) == (
+        "Multiple schemas are configured in the search path, which means "
+        "Superset is unable to determine the schema of unqualified table "
+        "names and enforce permissions."
+    )

tests/unit_tests/db_engine_specs/test_presto.py

@@ -20,6 +20,7 @@ from typing import Any, Dict, Optional, Type
 import pytest
 import pytz
 from sqlalchemy import types
+from sqlalchemy.engine.url import make_url
 
 from superset.utils.core import GenericDataType
 from tests.unit_tests.db_engine_specs.utils import (
@@ -82,3 +83,26 @@ def test_get_column_spec(
     from superset.db_engine_specs.presto import PrestoEngineSpec as spec
 
     assert_column_spec(spec, native_type, sqla_type, attrs, generic_type, is_dttm)
+
+
+def test_get_schema_from_engine_params() -> None:
+    """
+    Test the ``get_schema_from_engine_params`` method.
+    """
+    from superset.db_engine_specs.presto import PrestoEngineSpec
+
+    assert (
+        PrestoEngineSpec.get_schema_from_engine_params(
+            make_url("presto://localhost:8080/hive/default"),
+            {},
+        )
+        == "default"
+    )
+
+    assert (
+        PrestoEngineSpec.get_schema_from_engine_params(
+            make_url("presto://localhost:8080/hive"),
+            {},
+        )
+        is None
+    )

tests/unit_tests/db_engine_specs/test_snowflake.py

@@ -24,6 +24,7 @@ from unittest import mock
 
 import pytest
 from pytest_mock import MockerFixture
+from sqlalchemy.engine.url import make_url
 
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
@@ -169,3 +170,34 @@ def test_get_extra_params(mocker: MockerFixture) -> None:
             "connect_args": {"application": "Custom user agent", "foo": "bar"}
         }
     }
+
+
+def test_get_schema_from_engine_params() -> None:
+    """
+    Test the ``get_schema_from_engine_params`` method.
+    """
+    from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
+
+    assert (
+        SnowflakeEngineSpec.get_schema_from_engine_params(
+            make_url("snowflake://user:pass@account/database_name/default"),
+            {},
+        )
+        == "default"
+    )
+
+    assert (
+        SnowflakeEngineSpec.get_schema_from_engine_params(
+            make_url("snowflake://user:pass@account/database_name"),
+            {},
+        )
+        is None
+    )
+
+    assert (
+        SnowflakeEngineSpec.get_schema_from_engine_params(
+            make_url("snowflake://user:pass@account/"),
+            {},
+        )
+        is None
+    )

tests/unit_tests/explore/utils_test.py

@@ -274,10 +274,8 @@ def test_query_no_access(mocker: MockFixture, client) -> None:
     from superset.models.core import Database
     from superset.models.sql_lab import Query
 
-    inspect = mocker.patch("superset.security.manager.inspect")
-    inspect().default_schema_name = "public"
-
     database = mocker.MagicMock()
+    database.get_default_schema_for_query.return_value = "public"
     mocker.patch(
         query_find_by_id,
         return_value=Query(database=database, sql="select * from foo"),

tests/unit_tests/security/manager_test.py

@@ -52,8 +52,7 @@ def test_raise_for_access_query_default_schema(
     SqlaTable.query_datasources_by_name.return_value = []
 
     database = mocker.MagicMock()
-    database.db_engine_spec.dynamic_schema = False
-    database.get_inspector_with_context().__enter__().default_schema_name = "public"
+    database.get_default_schema_for_query.return_value = "public"
     query = mocker.MagicMock()
     query.database = database
     query.sql = "SELECT * FROM ab_user"