Fix replacement in attached Postgres catalogs

Signed-off-by: clowrance-proppilot <clowrance@propertypilot.com>

Author: clowrance-proppilot

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

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")