diff --git a/superset/sql/parse.py b/superset/sql/parse.py index 154a8ae0b19..f2f6b805782 100644 --- a/superset/sql/parse.py +++ b/superset/sql/parse.py @@ -923,9 +923,11 @@ class SQLStatement(BaseSQLStatement[exp.Expression]): :return: A new SQLStatement with the create table statement. """ table_expr = exp.Table( - this=exp.Identifier(this=table.table), - db=exp.Identifier(this=table.schema) if table.schema else None, - catalog=exp.Identifier(this=table.catalog) if table.catalog else None, + this=exp.Identifier(this=table.table, quoted=True), + db=exp.Identifier(this=table.schema, quoted=True) if table.schema else None, + catalog=exp.Identifier(this=table.catalog, quoted=True) + if table.catalog + else None, ) create_table = exp.Create( this=table_expr, diff --git a/superset/sqllab/schemas.py b/superset/sqllab/schemas.py index 3c2413f8d3b..4adb07e810a 100644 --- a/superset/sqllab/schemas.py +++ b/superset/sqllab/schemas.py @@ -14,10 +14,17 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from marshmallow import fields, Schema +from marshmallow import fields, Schema, validate from superset.databases.schemas import ImportV1DatabaseSchema +# Restricts the optional CTAS target name to a bare SQL identifier. Shared by the +# SQL Lab execute payload schemas so both request paths validate it identically. +tmp_table_name_validator = validate.Regexp( + r"^([A-Za-z_][A-Za-z0-9_]*)?\Z", + error="tmp_table_name must contain only letters, digits, and underscores", +) + sql_lab_get_results_schema = { "type": "object", "properties": { @@ -69,7 +76,10 @@ class ExecutePayloadSchema(Schema): tab = fields.String(allow_none=True) ctas_method = fields.String(allow_none=True) templateParams = fields.String(allow_none=True) # noqa: N815 - tmp_table_name = fields.String(allow_none=True) + tmp_table_name = fields.String( + allow_none=True, + validate=tmp_table_name_validator, + ) select_as_cta = fields.Boolean(allow_none=True) runAsync = fields.Boolean(allow_none=True) # noqa: N815 expand_data = fields.Boolean(allow_none=True) diff --git a/superset/views/sql_lab/schemas.py b/superset/views/sql_lab/schemas.py index c9f57d7c023..99f6fed7b89 100644 --- a/superset/views/sql_lab/schemas.py +++ b/superset/views/sql_lab/schemas.py @@ -17,6 +17,8 @@ from marshmallow import fields, Schema +from superset.sqllab.schemas import tmp_table_name_validator + class SqlJsonPayloadSchema(Schema): database_id = fields.Integer(required=True) @@ -28,7 +30,10 @@ class SqlJsonPayloadSchema(Schema): tab = fields.String(allow_none=True) ctas_method = fields.String(allow_none=True) templateParams = fields.String(allow_none=True) # noqa: N815 - tmp_table_name = fields.String(allow_none=True) + tmp_table_name = fields.String( + allow_none=True, + validate=tmp_table_name_validator, + ) select_as_cta = fields.Boolean(allow_none=True) runAsync = fields.Boolean(allow_none=True) # noqa: N815 expand_data = fields.Boolean(allow_none=True) diff --git a/tests/integration_tests/celery_tests.py b/tests/integration_tests/celery_tests.py index 8fc5fe293e0..11f9cdf0707 100644 --- a/tests/integration_tests/celery_tests.py +++ b/tests/integration_tests/celery_tests.py @@ -130,6 +130,18 @@ def quote_f(value: Optional[str]): return inspector.engine.dialect.identifier_preparer.quote_identifier(value) +def expected_cta_sql( + ctas_method: CTASMethod, table: str, schema: Optional[str] = None +) -> str: + target = quote_f(table) + if schema: + target = f"{quote_f(schema)}.{target}" + return ( + f"CREATE {ctas_method.name} {target} AS\n" + "SELECT\n name\nFROM birth_names\nLIMIT 1" + ) + + def cta_result(ctas_method: CTASMethod): if backend() != "presto": return [], [] @@ -225,31 +237,7 @@ def test_run_sync_query_cta_no_data(test_client): @pytest.mark.usefixtures("load_birth_names_data", "login_as_admin") -@pytest.mark.parametrize( - "ctas_method, expected", - [ - ( - CTASMethod.TABLE, - """ -CREATE TABLE sqllab_test_db.test_sync_cta_table AS -SELECT - name -FROM birth_names -LIMIT 1 - """.strip(), - ), - ( - CTASMethod.VIEW, - """ -CREATE VIEW sqllab_test_db.test_sync_cta_view AS -SELECT - name -FROM birth_names -LIMIT 1 - """.strip(), - ), - ], -) +@pytest.mark.parametrize("ctas_method", [CTASMethod.TABLE, CTASMethod.VIEW]) @mock.patch( # noqa: PT008 "superset.sqllab.sqllab_execution_context.get_cta_schema_name", lambda d, u, s, sql: CTAS_SCHEMA_NAME, @@ -257,12 +245,13 @@ LIMIT 1 def test_run_sync_query_cta_config( test_client, ctas_method: CTASMethod, - expected: str, ) -> None: - if backend() == "sqlite": + db_backend = backend() + if db_backend == "sqlite": # sqlite doesn't support schemas return tmp_table_name = f"{TEST_SYNC_CTA}_{ctas_method.name.lower()}" + expected = expected_cta_sql(ctas_method, tmp_table_name, CTAS_SCHEMA_NAME) result = run_sql( test_client, QUERY, cta=True, ctas_method=ctas_method, tmp_table=tmp_table_name ) @@ -281,31 +270,7 @@ def test_run_sync_query_cta_config( @pytest.mark.usefixtures("load_birth_names_data", "login_as_admin") -@pytest.mark.parametrize( - "ctas_method, expected", - [ - ( - CTASMethod.TABLE, - """ -CREATE TABLE sqllab_test_db.test_async_cta_config_table AS -SELECT - name -FROM birth_names -LIMIT 1 - """.strip(), - ), - ( - CTASMethod.VIEW, - """ -CREATE VIEW sqllab_test_db.test_async_cta_config_view AS -SELECT - name -FROM birth_names -LIMIT 1 - """.strip(), - ), - ], -) +@pytest.mark.parametrize("ctas_method", [CTASMethod.TABLE, CTASMethod.VIEW]) @mock.patch( # noqa: PT008 "superset.sqllab.sqllab_execution_context.get_cta_schema_name", lambda d, u, s, sql: CTAS_SCHEMA_NAME, @@ -313,12 +278,13 @@ LIMIT 1 def test_run_async_query_cta_config( test_client, ctas_method: CTASMethod, - expected: str, ) -> None: - if backend() == "sqlite": + db_backend = backend() + if db_backend == "sqlite": # sqlite doesn't support schemas return tmp_table_name = f"{TEST_ASYNC_CTA_CONFIG}_{ctas_method.name.lower()}" + expected = expected_cta_sql(ctas_method, tmp_table_name, CTAS_SCHEMA_NAME) result = run_sql( test_client, QUERY, @@ -341,37 +307,14 @@ def test_run_async_query_cta_config( @pytest.mark.usefixtures("load_birth_names_data", "login_as_admin") -@pytest.mark.parametrize( - "ctas_method, expected", - [ - ( - CTASMethod.TABLE, - """ -CREATE TABLE test_async_cta_table AS -SELECT - name -FROM birth_names -LIMIT 1 - """.strip(), - ), - ( - CTASMethod.VIEW, - """ -CREATE VIEW test_async_cta_view AS -SELECT - name -FROM birth_names -LIMIT 1 - """.strip(), - ), - ], -) +@pytest.mark.parametrize("ctas_method", [CTASMethod.TABLE, CTASMethod.VIEW]) def test_run_async_cta_query( test_client, ctas_method: CTASMethod, - expected: str, ) -> None: + db_backend = backend() table_name = f"{TEST_ASYNC_CTA}_{ctas_method.name.lower()}" + expected = expected_cta_sql(ctas_method, table_name) result = run_sql( test_client, QUERY, @@ -388,7 +331,7 @@ def test_run_async_cta_query( assert query.executed_sql == expected assert QUERY == query.sql - assert query.rows == (1 if backend() == "presto" else 0) + assert query.rows == (1 if db_backend == "presto" else 0) assert query.select_as_cta assert query.select_as_cta_used @@ -396,37 +339,14 @@ def test_run_async_cta_query( @pytest.mark.usefixtures("load_birth_names_data", "login_as_admin") -@pytest.mark.parametrize( - "ctas_method, expected", - [ - ( - CTASMethod.TABLE, - """ -CREATE TABLE test_async_lower_limit_table AS -SELECT - name -FROM birth_names -LIMIT 1 - """.strip(), - ), - ( - CTASMethod.VIEW, - """ -CREATE VIEW test_async_lower_limit_view AS -SELECT - name -FROM birth_names -LIMIT 1 - """.strip(), - ), - ], -) +@pytest.mark.parametrize("ctas_method", [CTASMethod.TABLE, CTASMethod.VIEW]) def test_run_async_cta_query_with_lower_limit( test_client, ctas_method: CTASMethod, - expected: str, ) -> None: + db_backend = backend() tmp_table = f"{TEST_ASYNC_LOWER_LIMIT}_{ctas_method.name.lower()}" + expected = expected_cta_sql(ctas_method, tmp_table) result = run_sql( test_client, QUERY, @@ -441,14 +361,14 @@ def test_run_async_cta_query_with_lower_limit( sqlite_select_sql = f"SELECT\n *\nFROM {tmp_table}\nLIMIT {query.limit}\nOFFSET 0" assert query.select_sql == ( sqlite_select_sql - if backend() == "sqlite" + if db_backend == "sqlite" else get_select_star(tmp_table, query.limit) ) assert query.executed_sql == expected assert QUERY == query.sql - assert query.rows == (1 if backend() == "presto" else 0) + assert query.rows == (1 if db_backend == "presto" else 0) assert query.limit == 50000 assert query.select_as_cta assert query.select_as_cta_used diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py index 27c29322970..7fd7e6a6f3e 100644 --- a/tests/unit_tests/sql/parse_tests.py +++ b/tests/unit_tests/sql/parse_tests.py @@ -2712,7 +2712,7 @@ def test_rls_predicate_transformer( "SELECT * FROM some_table", Table("some_table"), """ -CREATE TABLE some_table AS +CREATE TABLE "some_table" AS SELECT * FROM some_table @@ -2722,7 +2722,7 @@ FROM some_table "SELECT * FROM some_table", Table("some_table", "schema1", "catalog1"), """ -CREATE TABLE catalog1.schema1.some_table AS +CREATE TABLE "catalog1"."schema1"."some_table" AS SELECT * FROM some_table diff --git a/tests/unit_tests/sql_lab_execution_context.py b/tests/unit_tests/sql_lab_execution_context.py index 1591967de48..41374ec293b 100644 --- a/tests/unit_tests/sql_lab_execution_context.py +++ b/tests/unit_tests/sql_lab_execution_context.py @@ -18,6 +18,7 @@ import pytest +from superset.sql.parse import CTASMethod from superset.sqllab.sqllab_execution_context import ( CreateTableAsSelect, SqlJsonExecutionContext, @@ -97,6 +98,6 @@ def test_create_table_as_select(): "tmp_table_name": "temp_table", } ctas = CreateTableAsSelect.create_from(query_params) - assert ctas.ctas_method == "TABLE" + assert ctas.ctas_method == CTASMethod.TABLE assert ctas.target_schema_name == "public" assert ctas.target_table_name == "temp_table"