8000 [BUG] Reject SQL statements with invalid namespaces by eddiethedean · Pull Request #61 · dmux/minilake · GitHub
[go: up one dir, main page]

Skip to content
< 8000 link rel="preload" as="image" href="https://github.githubassets.com/assets/changelog-dark-b1c25ae971bd641b.png"/>
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
205 changes: 136 additions & 69 deletions src/minilake/services/sql_statements.py
8000
Original file line number Diff line number Diff line change
Expand Up @@ -884,7 +884,11 @@ def _ensure_delta_tree_writable(storage_location: str) -> None:


async def _execute_delta_write(
kind: str, storage_location: str, match: "re.Match"
kind: str,
storage_location: str,
match: "re.Match",
catalog: Optional[str] = None,
schema_name: Optional[str] = None,
) -> tuple[List[str], List[List[Any]]]:
"""Execute a real INSERT/UPDATE/DELETE against an EXTERNAL Delta table via
a generated Spark job — DuckDB's delta extension can't write, so this is
Expand Down Expand Up @@ -914,69 +918,75 @@ async def _execute_delta_write(
parquet_path = scratch_dir / f"{tmp_table}.parquet"
conn = await pool.get_uc_connection()
async with await pool.get_uc_lock():
previous = await _apply_default_namespace(conn, catalog, schema_name)
try:
# Read the existing Delta table's real schema (name + type) via
# delta_scan (read-only, fine for inspection) so the appended
# data's types are CAST to match exactly — DuckDB's literal type
# inference (e.g. a bare integer -> INTEGER) commonly doesn't
# match the real table's type (e.g. pandas/Spark int64 -> BIGINT),
# and Delta's writer rejects mismatched types rather than
# silently truncating. A first-ever write (no existing data) has
# no schema to conform to, so falls back to no casting.
try:
schema_rows = conn.execute(f"DESCRIBE SELECT * FROM delta_scan('{storage_location}')").fetchall()
target_schema = [(row[0], row[1]) for row in schema_rows]
except Exception:
target_schema = None

if payload.strip().upper().startswith("VALUES"):
# Delta's append mode checks column *names*, not just position,
# so a plain `VALUES (...)` insert (no explicit column list)
# needs real column names too.
if explicit_columns:
col_names = [c.strip() for c in explicit_columns.split(",")]
elif target_schema:
col_names = [name for name, _ in target_schema]
else:
raise DatabricksError(
error_code="INVALID_REQUEST",
message=(
"Cannot infer column names for INSERT INTO ... VALUES without an "
"explicit column list: the Delta table has no existing data/schema "
"to read yet. Either write the table's first version with Spark "
"directly, or use INSERT INTO table (col1, col2) VALUES (...)."
),
status_code=400,
)
col_list = ", ".join(col_names)
if target_schema:
type_by_name = dict(target_schema)
select_exprs = ", ".join(
f"CAST({name} AS {type_by_name[name]}) AS {name}" if name in type_by_name else name
for name in col_names
# Read the existing Delta table's real schema (name + type) via
# delta_scan (read-only, fine for inspection) so the appended
# data's types are CAST to match exactly — DuckDB's literal type
# inference (e.g. a bare integer -> INTEGER) commonly doesn't
# match the real table's type (e.g. pandas/Spark int64 -> BIGINT),
# and Delta's writer rejects mismatched types rather than
# silently truncating. A first-ever write (no existing data) has
# no schema to conform to, so falls back to no casting.
try:
schema_rows = conn.execute(
f"DESCRIBE SELECT * FROM delta_scan('{storage_location}')"
).fetchall()
target_schema = [(row[0], row[1]) for row in schema_rows]
except Exception:
target_schema = None

if payload.strip().upper().startswith("VALUES"):
# Delta's append mode checks column *names*, not just position,
# so a plain `VALUES (...)` insert (no explicit column list)
# needs real column names too.
if explicit_columns:
col_names = [c.strip() for c in explicit_columns.split(",")]
elif target_schema:
col_names = [name for name, _ in target_schema]
else:
raise DatabricksError(
error_code="INVALID_REQUEST",
message=(
"Cannot infer column names for INSERT INTO ... VALUES without an "
"explicit column list: the Delta table has no existing data/schema "
"to read yet. Either write the table's first version with Spark "
"directly, or use INSERT INTO table (col1, col2) VALUES (...)."
8000 ),
status_code=400,
)
col_list = ", ".join(col_names)
if target_schema:
type_by_name = dict(target_schema)
select_exprs = ", ".join(
f"CAST({name} AS {type_by_name[name]}) AS {name}" if name in type_by_name else name
for name in col_names
)
else:
select_exprs = "*"
conn.execute(
f"CREATE TEMP TABLE {tmp_table} AS SELECT {select_exprs} FROM ({payload}) AS t({col_list})"
)
else:
select_exprs = "*"
conn.execute(
f"CREATE TEMP TABLE {tmp_table} AS SELECT {select_exprs} FROM ({payload}) AS t({col_list})"
if target_schema:
cast_exprs = ", ".join(
f"CAST({name} AS {sql_type}) AS {name}" for name, sql_type in target_schema
)
conn.execute(f"CREATE TEMP TABLE {tmp_table} AS SELECT {cast_exprs} FROM ({payload})")
else:
conn.execute(f"CREATE TEMP TABLE {tmp_table} AS {payload}")
row_count = conn.execute(f"SELECT COUNT(*) FROM {tmp_table}").fetchone()[0]
conn.execute(f"COPY {tmp_table} TO '{parquet_path}' (FORMAT PARQUET)")
conn.execute(f"DROP TABLE {tmp_table}")
except DatabricksError:
raise
except Exception as e:
raise DatabricksError(
error_code="INVALID_REQUEST", message=f"Failed to prepare INSERT data: {e}", status_code=400
)
else:
if target_schema:
cast_exprs = ", ".join(
f"CAST({name} AS {sql_type}) AS {name}" for name, sql_type in target_schema
)
conn.execute(f"CREATE TEMP TABLE {tmp_table} AS SELECT {cast_exprs} FROM ({payload})")
else:
conn.execute(f"CREATE TEMP TABLE {tmp_table} AS {payload}")
row_count = conn.execute(f"SELECT COUNT(*) FROM {tmp_table}").fetchone()[0]
conn.execute(f"COPY {tmp_table} TO '{parquet_path}' (FORMAT PARQUET)")
conn.execute(f"DROP TABLE {tmp_table}")
except DatabricksError:
raise
except Exception as e:
raise DatabricksError(
error_code="INVALID_REQUEST", message=f"Failed to prepare INSERT data: {e}", status_code=400
)
finally:
_restore_default_namespace(conn, previous)

script = (
"from pyspark.sql import SparkSession\n"
Expand Down Expand Up @@ -1058,20 +1068,53 @@ async def _execute_delta_write(

async def _apply_default_namespace(conn, catalog: Optional[str], schema_name: Optional[str]):
"""Point the connection at `catalog[.schema]`, returning what to restore afterwards."""
if not catalog and not schema_name:
if catalog is None and schema_name is None:
return None

if catalog == "":
raise DatabricksError(
error_code="INVALID_PARAMETER_VALUE",
message="Catalog cannot be empty",
status_code=400,
)
if schema_name == "":
raise DatabricksError(
error_code="INVALID_PARAMETER_VALUE",
message="Schema cannot be empty",
status_code=400,
)

if catalog:
from minilake.services import unity_catalog

# DuckDB accepts `USE <name>` when the name is a schema in the current
# catalog, even if it is not a catalog. Check the UC registry first so a
# catalog-only request cannot silently select a same-named schema.
if catalog != "unity_catalog" and catalog not in unity_catalog._state["catalogs"]:
raise DatabricksError(
error_code="INVALID_PARAMETER_VALUE",
message=f"Catalog '{catalog}' does not exist",
status_code=400,
)

current = conn.execute("SELECT current_catalog(), current_schema()").fetchone()
target_catalog = catalog if catalog is not None else current[0]
# A catalog-only request keeps DuckDB's established `main` schema default,
# but names it explicitly so DuckDB cannot mistake the catalog for a schema
# in the connection's previous catalog.
target_schema = schema_name if schema_name is not None else ("main" if catalog is not None else current[1])
target = f'"{target_catalog}"."{target_schema}"'
try:
current = conn.execute("SELECT current_catalog(), current_schema()").fetchone()
target = f'"{catalog}"' if catalog else f'"{current[0]}"'
if schema_name:
target = f'{target}."{schema_name}"'
conn.execute(f"USE {target}")
return current
except Exception as e:
# A bad default is the caller's problem to see in the statement error, not a
# reason to fail before the statement even runs.
logger.warning(f"Could not set default namespace {catalog}.{schema_name}: {e}")
return None
_restore_default_namespace(conn, current)
raise DatabricksError(
error_code="INVALID_PARAMETER_VALUE",
message=f"Could not set default namespace {catalog}.{schema_name}: {e}",
status_code=400,
) from e
return current


def _restore_default_namespace(conn, previous) -> None:
Expand Down Expand Up @@ -1105,6 +1148,28 @@ async def _register_external_delta(
return [], []


async def _validate_special_statement_namespace(catalog: Optional[str], schema_name: Optional[str]) -> None:
"""Validate an explicit statement namespace before special SQL paths write data."""
if catalog is None and schema_name is None:
return

pool = get_duckdb_pool()
if not pool:
raise DatabricksError(error_code="INVALID_STATE", message="DuckDB pool not initialized", status_code=500)

try:
conn = await pool.get_uc_connection()
lock = await pool.get_uc_lock()
except Exception as e:
raise DatabricksError(
error_code="INVALID_REQUEST", message=f"Failed to get UC connection: {e}", status_code=400
) from e

async with lock:
previous = await _apply_default_namespace(conn, catalog, schema_name)
_restore_default_namespace(conn, previous)


async def _execute_sql_real(
warehouse_id: str,
sql: str,
Expand Down Expand Up @@ -1136,15 +1201,17 @@ async def _execute_sql_real_typed(

create_external = _classify_create_external(sql)
if create_external is not None:
await _validate_special_statement_namespace(catalog, schema_name)
columns, rows = await _register_external_delta(*create_external)
return columns, [None] * len(columns), rows, None

sql = _strip_databricks_ddl_clauses(sql)

delta_write = _classify_delta_write(sql)
if delta_write is not None:
await _validate_special_statement_namespace(catalog, schema_name)
kind, storage_location, match = delta_write
columns, rows = await _execute_delta_write(kind, storage_location, match)
columns, rows = await _execute_delta_write(kind, storage_location, match, catalog, schema_name)
return columns, [None] * len(columns), rows, None

pool = get_duckdb_pool()
Expand Down
114 changes: 114 additions & 0 deletions tests/test_sql_statements.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import pytest
import requests
from databricks.sdk import WorkspaceClient
from databricks.sdk.errors import DatabricksError
from databricks.sdk.service.catalog import DataSourceFormat, TableType
from databricks.sdk.service.sql import ColumnInfo, Disposition, Format

Expand Down Expand Up @@ -385,6 +386,119 @@ def test_execute_insert_and_select(warehouse_with_table, workspace_client):
print("✓ INSERT and SELECT executed successfully")


@pytest.mark.parametrize(
("catalog", "schema"),
[
pytest.param("missing_namespace_catalog", "missing_schema", id="missing-catalog"),
pytest.param("namespace_valid_catalog", "missing_schema", id="missing-schema"),
pytest.param("main", None, id="missing-catalog-matches-default-schema"),
pytest.param("", None, id="empty-catalog"),
pytest.param(None, "", id="empty-schema"),
pytest.param("", "", id="both-empty"),
],
)
@pytest.mark.parametrize(
"statement_template",
[
pytest.param("DELETE FROM {table}", id="delete"),
pytest.param("UPDATE {table} SET id = id + 1", id="update"),
pytest.param("INSERT INTO {table} VALUES (30)", id="insert"),
],
)
@pytest.mark.workflow
def test_invalid_namespace_rejects_statement_without_fallback(
workspace_client: WorkspaceClient,
catalog: str,
schema: str | None,
statement_template: str,
):
"""Invalid explicit catalog/schema must not run SQL in the connection's old namespace."""
warehouse = workspace_client.warehouses.create(name="invalid_namespace_wh")
if catalog == "namespace_valid_catalog":
workspace_client.catalogs.create(name=catalog)

table_name = f"namespace_guard_{uuid.uuid4().hex[:8]}"
workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"CREATE TABLE {table_name} AS SELECT 10 AS id UNION ALL SELECT 20",
)

try:
with pytest.raises(DatabricksError):
workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=statement_template.format(table=table_name),
catalog=catalog,
schema=schema,
)

result = workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"SELECT id FROM {table_name} ORDER BY id",
)
assert result.result.data_array == [[10], [20]]
finally:
workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"DROP TABLE IF EXISTS {table_name}",
)


@pytest.mark.serial
@pytest.mark.workflow
def test_catalog_only_namespace_is_not_confused_with_same_named_schema(workspace_client: WorkspaceClient):
"""A catalog-only context must select that catalog when its name is also a schema."""
warehouse = workspace_client.warehouses.create(name="namespace_collision_wh")
catalog = f"namespace_catalog_{uuid.uuid4().hex[:8]}"
table_name = f"namespace_collision_{uuid.uuid4().hex[:8]}"
workspace_client.catalogs.create(name=catalog)

workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f'CREATE SCHEMA unity_catalog."{catalog}"',
)
workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"CREATE TABLE unity_catalog.{catalog}.{table_name} AS SELECT 99 AS id",
)
workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"CREATE TABLE {catalog}.main.{table_name} AS SELECT 1 AS id",
)

try:
result = workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"DELETE FROM {table_name}",
catalog=catalog,
)
assert result.result.data_array == [[1]]

intended_catalog = workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"SELECT id FROM {catalog}.main.{table_name}",
)
same_named_schema = workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"SELECT id FROM unity_catalog.{catalog}.{table_name}",
)
assert intended_catalog.result.data_array == []
assert same_named_schema.result.data_array == [[99]]
finally:
workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"DROP TABLE IF EXISTS {catalog}.main.{table_name}",
)
workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f"DROP TABLE IF EXISTS unity_catalog.{catalog}.{table_name}",
)
workspace_client.statement_execution.execute_statement(
warehouse_id=warehouse.id,
statement=f'DROP SCHEMA IF EXISTS unity_catalog."{catalog}" CASCADE',
)


@pytest.mark.crud
@pytest.mark.parametrize(
"mutation, expected_rows",
Expand Down
Loading
Loading
0