You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ru...@apache.org on 2026/09/12 18:00:17 UTC
(superset) branch master updated: fix(dataset): derive virtual dataset schema from SQL Lab query (#43313)
This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new f51ae6a506d fix(dataset): derive virtual dataset schema from SQL Lab query (#43313)
f51ae6a506d is described below
commit f51ae6a506d7020342b2eac36d466112045eef46
Author: Joe Li <jo...@preset.io>
AuthorDate: Sat Sep 12 11:00:10 2026 -0700
fix(dataset): derive virtual dataset schema from SQL Lab query (#43313)
Co-authored-by: Claude Sonnet 5 <no...@anthropic.com>
---
superset/commands/dataset/create.py | 84 ++++++++++-
superset/sql/parse.py | 33 +++++
tests/unit_tests/commands/dataset/test_create.py | 180 +++++++++++++++++++++++
tests/unit_tests/sql/parse_tests.py | 60 ++++++++
4 files changed, 356 insertions(+), 1 deletion(-)
diff --git a/superset/commands/dataset/create.py b/superset/commands/dataset/create.py
index 12084649ec4..5f490c2d36b 100644
--- a/superset/commands/dataset/create.py
+++ b/superset/commands/dataset/create.py
@@ -46,7 +46,8 @@ from superset.exceptions import (
SupersetTimeoutException,
)
from superset.extensions import security_manager
-from superset.sql.parse import Table
+from superset.models.helpers import json_to_dict
+from superset.sql.parse import process_jinja_sql, Table
from superset.utils.decorators import on_error, transaction
logger = logging.getLogger(__name__)
@@ -114,6 +115,87 @@ class CreateDatasetCommand(CreateMixin, BaseCommand):
if not catalog:
catalog = self._properties["catalog"] = database.get_default_catalog()
+ if sql:
+ try:
+ template_params_raw: str | None = self._properties.get(
+ "template_params"
+ )
+ template_params = json_to_dict(template_params_raw or "")
+ parse_result = process_jinja_sql(
+ sql,
+ database,
+ template_params,
+ )
+ except Exception: # pylint: disable=broad-exception-caught
+ logger.debug(
+ "Keeping submitted dataset schema because optional SQL "
+ "template processing failed",
+ exc_info=True,
+ )
+ else:
+ script = parse_result.script
+ tables = {table for table in parse_result.tables if table.table}
+ if len(tables) != len(parse_result.tables):
+ logger.debug(
+ "Ignoring empty table-name parser artifacts while deriving "
+ "the dataset schema"
+ )
+
+ if script.has_unparseable_statement:
+ logger.debug(
+ "Keeping submitted dataset schema because the SQL contains "
+ "an unparseable statement"
+ )
+ elif script.has_show_statement():
+ logger.debug(
+ "Keeping submitted dataset schema because the SQL contains "
+ "a metadata statement"
+ )
+ elif script.has_mutation():
+ logger.debug(
+ "Keeping submitted dataset schema because the SQL mutates "
+ "database state"
+ )
+ elif script.changes_default_schema():
+ logger.debug(
+ "Keeping submitted dataset schema because the SQL changes "
+ "the default schema"
+ )
+ elif script.has_quoted_table_location():
+ logger.debug(
+ "Keeping submitted dataset schema because a quoted table "
+ "location may have case-sensitive identity"
+ )
+ elif not tables:
+ logger.debug(
+ "Keeping submitted dataset schema because the SQL has no "
+ "table references"
+ )
+ elif any(table.schema is None for table in tables):
+ logger.debug(
+ "Keeping submitted dataset schema because the SQL has an "
+ "unqualified table reference"
+ )
+ elif (
+ len(
+ qualified_table_locations := {
+ (table.catalog, table.schema) for table in tables
+ }
+ )
+ != 1
+ ):
+ logger.debug(
+ "Keeping submitted dataset schema because the SQL spans "
+ "multiple catalog/schema locations"
+ )
+ else:
+ derived_catalog, derived_schema = next(
+ iter(qualified_table_locations)
+ )
+ schema = self._properties["schema"] = derived_schema
+ if derived_catalog is not None:
+ catalog = self._properties["catalog"] = derived_catalog
+
table = Table(table_name, schema, catalog)
if not DatasetDAO.validate_uniqueness(database, table):
diff --git a/superset/sql/parse.py b/superset/sql/parse.py
index 0f748a58c6e..d0b08530072 100644
--- a/superset/sql/parse.py
+++ b/superset/sql/parse.py
@@ -633,6 +633,14 @@ class BaseSQLStatement(Generic[InternalRepresentation]):
"""
return False
+ def has_quoted_table_location(self) -> bool:
+ """Check whether a table has a quoted catalog or schema identifier."""
+ return False
+
+ def is_show_statement(self) -> bool:
+ """Check whether this is a SHOW metadata statement."""
+ return False
+
def get_disallowed_tables(
self,
tables: set[str],
@@ -1572,6 +1580,21 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
return inner
return self.changes_search_path()
+ def has_quoted_table_location(self) -> bool:
+ """Return whether a table has a quoted catalog or schema identifier."""
+ for table in self._parsed.find_all(exp.Table):
+ for identifier_name in ("catalog", "db"):
+ identifier = table.args.get(identifier_name)
+ if isinstance(identifier, exp.Identifier) and identifier.args.get(
+ "quoted"
+ ):
+ return True
+ return False
+
+ def is_show_statement(self) -> bool:
+ """Return whether this is a SHOW metadata statement."""
+ return isinstance(self._parsed, exp.Show)
+
def get_disallowed_tables(
self,
tables: set[str],
@@ -2389,6 +2412,16 @@ class SQLScript:
"""
return any(statement.changes_default_schema() for statement in self.statements)
+ def has_quoted_table_location(self) -> bool:
+ """Check if any table has a quoted catalog or schema identifier."""
+ return any(
+ statement.has_quoted_table_location() for statement in self.statements
+ )
+
+ def has_show_statement(self) -> bool:
+ """Check if the script contains a SHOW metadata statement."""
+ return any(statement.is_show_statement() for statement in self.statements)
+
def optimize(self) -> SQLScript:
"""
Return optimized script.
diff --git a/tests/unit_tests/commands/dataset/test_create.py b/tests/unit_tests/commands/dataset/test_create.py
index c631c2711d9..f62953fb3bc 100644
--- a/tests/unit_tests/commands/dataset/test_create.py
+++ b/tests/unit_tests/commands/dataset/test_create.py
@@ -33,10 +33,189 @@ from superset.exceptions import (
from superset.models.core import Database
+def _validate_virtual_dataset(
+ sql: str,
+ engine: str = "postgresql",
+ template_params: str | None = None,
+) -> CreateDatasetCommand:
+ """Validate a virtual dataset while isolating command dependencies."""
+ mock_database = Mock(spec=Database)
+ mock_database.id = 1
+ mock_database.backend = engine
+ mock_database.db_engine_spec.engine = engine
+ mock_database.get_default_catalog.return_value = None
+
+ with (
+ patch(
+ "superset.commands.dataset.create.DatasetDAO.get_database_by_id",
+ return_value=mock_database,
+ ),
+ patch(
+ "superset.commands.dataset.create.DatasetDAO.validate_uniqueness",
+ return_value=True,
+ ),
+ patch("superset.commands.dataset.create.security_manager.raise_for_access"),
+ patch("superset.commands.dataset.create.populate_subjects"),
+ ):
+ properties = {
+ "database": 1,
+ "schema": "information_schema",
+ "table_name": "test_virtual_dataset",
+ "sql": sql,
+ }
+ if template_params is not None:
+ properties["template_params"] = template_params
+ command = CreateDatasetCommand(properties)
+ command.validate()
+
+ return command
+
+
+def test_create_dataset_schema_derived_from_single_schema_query() -> None:
+ """A single explicit query schema replaces the SQL Lab dropdown schema."""
+ command = _validate_virtual_dataset(
+ 'select * from public."Vehicle Sales"',
+ )
+
+ assert command._properties["schema"] == "public"
+
+
+def test_create_dataset_keeps_dropdown_schema_when_query_has_no_schema() -> None:
+ """An unqualified query keeps the SQL Lab dropdown schema."""
+ command = _validate_virtual_dataset("select * from t")
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_keeps_dropdown_schema_for_multi_schema_query() -> None:
+ """A query spanning multiple explicit schemas keeps the dropdown schema."""
+ command = _validate_virtual_dataset(
+ "select * from a.t1 union all select * from b.t2",
+ )
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_keeps_dropdown_schema_for_unqualified_reference() -> None:
+ """A mixed qualified and bare query keeps the dropdown schema."""
+ command = _validate_virtual_dataset(
+ "select * from public.t1 join t2 on public.t1.id = t2.id",
+ )
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_keeps_dropdown_schema_after_schema_change() -> None:
+ """A schema-changing statement prevents static schema derivation."""
+ command = _validate_virtual_dataset("use other; select * from public.t1")
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_keeps_dropdown_schema_when_sql_is_unparseable() -> None:
+ """A parse failure falls back to the SQL Lab dropdown schema."""
+ command = _validate_virtual_dataset("select * from {{ my_schema }}.t")
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_keeps_dropdown_schema_for_unparseable_statement() -> None:
+ """An opaque statement prevents derivation from an incomplete table set."""
+ command = _validate_virtual_dataset(
+ "explain select * from hidden.t; select * from public.t1",
+ )
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_keeps_dropdown_schema_for_mutating_query() -> None:
+ """A mutating statement keeps the dropdown rather than its source schema."""
+ command = _validate_virtual_dataset(
+ "insert into secret.t select * from public.s",
+ )
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_keeps_dropdown_schema_for_multi_catalog_query() -> None:
+ """The same schema in different catalogs is not a single location."""
+ command = _validate_virtual_dataset(
+ "select * from c1.s.t1 union all select * from c2.s.t2",
+ )
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_keeps_dropdown_schema_for_quoted_schema() -> None:
+ """A quoted schema on a case-folding engine prevents unsafe derivation."""
+ command = _validate_virtual_dataset(
+ 'select * from "public".t1 union all select * from public.t2',
+ engine="snowflake",
+ )
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_keeps_dropdown_schema_for_jinja_macro_table() -> None:
+ """Partition-macro tables participate in the single-location check."""
+ command = _validate_virtual_dataset(
+ "select * from public.t1 where ds = "
+ "'{{ presto.latest_partition(\"secret.audit\") }}'",
+ engine="presto",
+ )
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_ignores_empty_table_function_reference() -> None:
+ """An empty parser artifact does not make a qualified query ambiguous."""
+ command = _validate_virtual_dataset(
+ "select * from public.t1 join generate_series(1, 10) on true",
+ )
+
+ assert command._properties["schema"] == "public"
+
+
+def test_create_dataset_keeps_dropdown_schema_for_template_error() -> None:
+ """A Jinja failure in best-effort derivation keeps the dropdown schema."""
+ command = _validate_virtual_dataset("SELECT '{{' AS x FROM public.t")
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_schema_uses_template_params() -> None:
+ """Submitted template parameters determine the derived query schema."""
+ command = _validate_virtual_dataset(
+ "select * from {{ source | default('public.t') }}",
+ template_params='{"source": "secret.t"}',
+ )
+
+ assert command._properties["schema"] == "secret"
+
+
+def test_create_dataset_keeps_dropdown_schema_for_show_statement() -> None:
+ """A metadata statement has no meaningful virtual-dataset schema."""
+ command = _validate_virtual_dataset(
+ "show columns from foo from bar",
+ engine="mysql",
+ )
+
+ assert command._properties["schema"] == "information_schema"
+
+
+def test_create_dataset_derives_catalog_and_schema() -> None:
+ """A single three-part reference updates both catalog and schema."""
+ command = _validate_virtual_dataset("select * from prod.sales.orders")
+
+ assert command._properties["schema"] == "sales"
+ assert command._properties["catalog"] == "prod"
+
+
def test_create_dataset_invalid_sql_parse_error() -> None:
"""Test that invalid SQL returns a 4xx error when caught as SupersetParseError."""
mock_database = Mock(spec=Database)
mock_database.id = 1
+ mock_database.backend = "postgresql"
mock_database.db_engine_spec.engine = "postgresql"
mock_database.get_default_catalog.return_value = None
@@ -88,6 +267,7 @@ def test_create_dataset_valid_sql_with_access_error() -> None:
"""
mock_database = Mock(spec=Database)
mock_database.id = 1
+ mock_database.backend = "postgresql"
mock_database.db_engine_spec.engine = "postgresql"
mock_database.get_default_catalog.return_value = None
diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py
index 1c85ab7adfa..20cf37d5035 100644
--- a/tests/unit_tests/sql/parse_tests.py
+++ b/tests/unit_tests/sql/parse_tests.py
@@ -5868,6 +5868,66 @@ def test_changes_default_schema(sql: str, engine: str, expected: bool) -> None:
assert SQLScript(sql, engine).changes_default_schema() == expected
+@pytest.mark.parametrize(
+ "sql, engine, expected",
+ [
+ # A quoted catalog identifier is case-sensitive and may not match
+ # after the engine's default case folding.
+ ('SELECT * FROM "c1".s.t1', "snowflake", True),
+ # A quoted schema (``db``) identifier is equally unsafe.
+ ('SELECT * FROM c1."s".t1', "snowflake", True),
+ # An unquoted location can be safely folded to the engine's default
+ # case.
+ ("SELECT * FROM c1.s.t1", "snowflake", False),
+ # Quoting the table name itself doesn't affect catalog/schema safety.
+ ('SELECT * FROM "t1"', "snowflake", False),
+ ],
+)
+def test_has_quoted_table_location(sql: str, engine: str, expected: bool) -> None:
+ """
+ `has_quoted_table_location` flags queries whose catalog or schema is
+ quoted, so the SQL Lab dataset-creation flow keeps the dropdown schema
+ instead of deriving a location that may not match after case folding.
+ """
+ assert SQLStatement(sql, engine).has_quoted_table_location() == expected
+ assert SQLScript(sql, engine).has_quoted_table_location() == expected
+
+
+def test_has_quoted_table_location_unsupported_dialect() -> None:
+ """
+ Engines without a sqlglot AST (e.g. Kusto KQL) report no quoted table
+ location instead of raising, matching the ``BaseSQLStatement`` default.
+ """
+ statement = KustoKQLStatement("foo | take 100", "kustokql")
+ assert statement.has_quoted_table_location() is False
+
+
+@pytest.mark.parametrize(
+ "sql, engine, expected",
+ [
+ ("show columns from foo from bar", "mysql", True),
+ ("SELECT * FROM t1", "mysql", False),
+ ],
+)
+def test_is_show_statement(sql: str, engine: str, expected: bool) -> None:
+ """
+ `is_show_statement`/`has_show_statement` identify metadata statements so
+ the SQL Lab dataset-creation flow keeps the dropdown schema rather than
+ deriving one from a query with no meaningful result set.
+ """
+ assert SQLStatement(sql, engine).is_show_statement() == expected
+ assert SQLScript(sql, engine).has_show_statement() == expected
+
+
+def test_is_show_statement_unsupported_dialect() -> None:
+ """
+ Engines without a sqlglot AST are never treated as SHOW statements,
+ matching the ``BaseSQLStatement`` default.
+ """
+ statement = KustoKQLStatement("foo | take 100", "kustokql")
+ assert statement.is_show_statement() is False
+
+
@pytest.mark.parametrize(
"sql, denylist, expected",
[