Skip to content

Fix replacement in attached Postgres catalogs #6078

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
clowrance-proppilot wants to merge 1 commit into
base: main
Choose a base branch
from
+32 −10
Open

Fix replacement in attached Postgres catalogs #6078

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
20 changes: 10 additions & 10 deletions sqlmesh/core/engine_adapter/duckdb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,12 @@ def _create_table(
track_rows_processed: bool = True,
**kwargs: t.Any,
) -> None:
catalog = self.get_current_catalog()
table = (
table_name_or_schema.this
if isinstance(table_name_or_schema, exp.Schema)
else exp.to_table(table_name_or_schema)
)
catalog = table.catalog or self.get_current_catalog()
catalog_type_tuple = self.fetchone(
exp.select("type")
.from_("duckdb_databases()")
Expand All @@ -184,6 +189,9 @@ def _create_table(
partitioned_by_exps = None
if catalog_type == "ducklake":
partitioned_by_exps = kwargs.pop("partitioned_by", None)
elif catalog_type == "postgres" and replace:
self.execute(exp.Drop(this=table, kind="TABLE", exists=True, cascade=True))
replace = False

super()._create_table(
table_name_or_schema,
Expand All @@ -199,16 +207,8 @@ def _create_table(
)

if partitioned_by_exps:
# Schema object contains column definitions, so we extract Table
table_name = (
table_name_or_schema.this
if isinstance(table_name_or_schema, exp.Schema)
else table_name_or_schema
)
table_name_str = (
table_name.sql(dialect=self.dialect)
if isinstance(table_name, exp.Table)
else table_name
table.sql(dialect=self.dialect) if isinstance(table, exp.Table) else table
)
partitioned_by_str = ", ".join(
expr.sql(dialect=self.dialect) for expr in partitioned_by_exps
Expand Down
22 changes: 22 additions & 0 deletions tests/core/engine_adapter/test_duckdb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,28 @@ def test_replace_query_pandas(adapter: EngineAdapter, duck_conn):
pd.testing.assert_frame_equal(adapter.fetchdf("SELECT * FROM test_table"), df)


def test_replace_query_attached_postgres(
make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
) -> None:
adapter = make_mocked_engine_adapter(DuckDBEngineAdapter)
fetchone = mocker.patch.object(adapter, "fetchone", return_value=("postgres",))

adapter.replace_query(
"attached_postgres.test_schema.test_table",
parse_one("SELECT 1 AS a"),
)

assert fetchone.call_count == 1
assert (
fetchone.call_args.args[0].sql(dialect=adapter.dialect)
== "SELECT type FROM DUCKDB_DATABASES() WHERE database_name = 'attached_postgres'"
)
assert to_sql_calls(adapter) == [
'DROP TABLE IF EXISTS "attached_postgres"."test_schema"."test_table" CASCADE',
'CREATE TABLE IF NOT EXISTS "attached_postgres"."test_schema"."test_table" AS SELECT 1 AS "a"',
]


def test_set_current_catalog(make_mocked_engine_adapter: t.Callable, duck_conn):
adapter = make_mocked_engine_adapter(DuckDBEngineAdapter)
adapter.set_current_catalog("test_catalog")
Expand Down
Loading