diff --git a/sqlmesh/core/engine_adapter/duckdb.py b/sqlmesh/core/engine_adapter/duckdb.py index ebfcaa7901..b8654d269e 100644 --- a/sqlmesh/core/engine_adapter/duckdb.py +++ b/sqlmesh/core/engine_adapter/duckdb.py @@ -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()") @@ -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, @@ -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 diff --git a/tests/core/engine_adapter/test_duckdb.py b/tests/core/engine_adapter/test_duckdb.py index 9fd65a6e66..36fd8c4068 100644 --- a/tests/core/engine_adapter/test_duckdb.py +++ b/tests/core/engine_adapter/test_duckdb.py @@ -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")