feat: support schema-qualified temporary tables

fakesnow/copy_into.py

@@ -66,7 +66,7 @@ def copy_into(
         duck_conn.execute(sql)
         return sql
 
-    schema = table.db or current_schema
+    _, schema = table_stages.logical_schema(table, current_database, current_schema)
     assert schema
 
     parquet_info = None
@@ -343,7 +343,7 @@ def stage_url_from_var(
     fqname = f"{database_name}.{schema_name}.{name}"
 
     if stage.is_table_stage(name):
-        target = (target_table.catalog or current_database, target_table.db or current_schema, target_table.name)
+        target = (*table_stages.logical_schema(target_table, current_database, current_schema), target_table.name)
         url = stage.internal_dir(table_stages.resolve_for_copy(database_name, schema_name, name, target))
     else:
         duck_conn.execute(

fakesnow/cursor.py

@@ -301,6 +301,12 @@ def _transform(self, expression: Expr, params: MutableParams | None) -> Expr:
         return (
             expression.transform(lambda e: transforms.identifier(e, params), copy=False)
             .transform(transforms.upper_case_unquoted_identifiers, copy=False)
+            .transform(
+                lambda e: transforms.temporary_tables(
+                    e, self._duck_conn, self._conn.database, self._conn.schema, self._conn.table_stages
+                ),
+                copy=False,
+            )
             .transform(transforms.alter_session, copy=False)
             .transform(transforms.update_variables, variables=self._conn.variables, copy=False)
             .transform(transforms.current_version, copy=False)
@@ -630,8 +636,9 @@ def _execute(self, transformed: Expr, params: MutableParams | None = None) -> No
 
                 # a newly created table object has no load history
                 if table := transformed.find(exp.Table):
-                    catalog = table.catalog or self._conn.database
-                    schema = table.db or self._conn.schema
+                    catalog, schema = self._conn.table_stages.logical_schema(
+                        table, self._conn.database, self._conn.schema
+                    )
                     assert catalog and schema
                     self._duck_conn.execute(
                         f"DELETE FROM {catalog}._fs_information_schema._fs_load_history "
@@ -668,17 +675,15 @@ def _execute(self, transformed: Expr, params: MutableParams | None = None) -> No
         if table_comment := cast(tuple[exp.Table, str], transformed.args.get("table_comment")):
             # record table comment
             table, comment = table_comment
-            catalog = table.catalog or self._conn.database
-            schema = table.db or self._conn.schema
+            catalog, schema = self._conn.table_stages.metadata_owner(table, self._conn.database, self._conn.schema)
             assert catalog and schema
             self._duck_conn.execute(info_schema.insert_table_comment_sql(catalog, schema, table.name, comment))
 
         if (text_lengths := cast(list[tuple[str, int]], transformed.args.get("text_lengths"))) and (
             table := transformed.find(exp.Table)
         ):
             # record text lengths
-            catalog = table.catalog or self._conn.database
-            schema = table.db or self._conn.schema
+            catalog, schema = self._conn.table_stages.metadata_owner(table, self._conn.database, self._conn.schema)
             assert catalog and schema
             self._duck_conn.execute(info_schema.insert_text_lengths_sql(catalog, schema, table.name, text_lengths))
 

fakesnow/info_schema.py

@@ -114,16 +114,18 @@
     null::VARCHAR as IDENTITY_ORDERED,
     null::VARCHAR as DATA_TYPE_ALIAS
 from system.information_schema.columns columns
-left join _fs_global._fs_information_schema._fs_columns_ext ext
-  on ext_table_catalog = columns.table_catalog
- AND ext_table_schema = columns.table_schema
- AND ext_table_name = columns.table_name
- AND ext_column_name = columns.column_name
 LEFT JOIN duckdb_columns ddb_columns
   ON ddb_columns.database_name = columns.table_catalog
  AND ddb_columns.schema_name = columns.table_schema
  AND ddb_columns.table_name = columns.table_name
  AND ddb_columns.column_name = columns.column_name
+-- every session's temporary tables share temp.main, so their metadata is keyed by OID, which is unique across sessions
+LEFT JOIN _fs_global._fs_information_schema._fs_columns_ext ext
+  ON ext_table_catalog = columns.table_catalog
+ AND ext_table_schema = CASE WHEN columns.table_catalog = 'temp' THEN ddb_columns.table_oid::VARCHAR
+                             ELSE columns.table_schema END
+ AND ext_table_name = columns.table_name
+ AND ext_column_name = columns.column_name
 where schema_name != '_fs_information_schema'
 """
 

fakesnow/transforms/__init__.py

@@ -28,6 +28,7 @@
     list_stage as list_stage,
     put_stage as put_stage,
 )
+from fakesnow.transforms.temporary import temporary_tables as temporary_tables
 from fakesnow.transforms.transforms import (
     SERVER_VERSION as SERVER_VERSION,
     SUCCESS_NOP as SUCCESS_NOP,

fakesnow/transforms/merge.py

@@ -89,7 +89,7 @@ def _create_merge_candidates(merge_expr: exp.Merge) -> Expr:
             )
 
     sql = f"""
-    CREATE OR REPLACE TEMPORARY TABLE merge_candidates AS
+    CREATE OR REPLACE TEMPORARY TABLE "temp"."main".merge_candidates AS
     SELECT
         {", ".join(sorted(values))},
         CASE
@@ -128,7 +128,7 @@ def _mutations(merge_expr: exp.Merge) -> list[Expr]:
             if isinstance(then, exp.Var) and then.name.upper() == "DELETE":
                 delete_sql = f"""
                     DELETE FROM {target_tbl}
-                    USING merge_candidates AS {source_tbl}
+                    USING "temp"."main".merge_candidates AS {source_tbl}
                     WHERE {join_expr}
                     AND {source_tbl}.merge_op = {w_idx}
                 """
@@ -142,7 +142,7 @@ def _mutations(merge_expr: exp.Merge) -> list[Expr]:
                 update_sql = f"""
                     UPDATE {target_tbl}
                     SET {set_clauses}
-                    FROM merge_candidates AS {source_tbl}
+                    FROM "temp"."main".merge_candidates AS {source_tbl}
                     WHERE {join_expr}
                     AND {source_tbl}.merge_op = {w_idx}
                 """
@@ -157,7 +157,7 @@ def _mutations(merge_expr: exp.Merge) -> list[Expr]:
             insert_sql = f"""
                 INSERT INTO {target_tbl} {columns}
                 SELECT {values}
-                FROM merge_candidates AS {source_tbl}
+                FROM "temp"."main".merge_candidates AS {source_tbl}
                 WHERE {source_tbl}.merge_op = {w_idx}
             """
             statements.append(sqlglot.parse_one(insert_sql))
@@ -219,7 +219,7 @@ def _counts(merge_expr: exp.Merge) -> Expr:
     ]
     sql = f"""
     SELECT {", ".join(count_statements)}
-    FROM merge_candidates
+    FROM "temp"."main".merge_candidates
     """
 
     counts = sqlglot.parse_one(sql)

fakesnow/transforms/stage.py

@@ -55,6 +55,8 @@ def __init__(self, duck_conn: DuckDBPyConnection):
         self._session = uuid.uuid4().hex
         # Temporary table OID -> (catalog, schema) it was created in, which DuckDB does not record.
         self._temporary: dict[int, tuple[str | None, str | None]] = {}
+        # OIDs of this session's temporary tables whose metadata is recorded, see metadata_owner.
+        self._metadata_oids: set[int] = set()
 
     def record_temporary_table(self, expression: Expr, catalog: str | None, schema: str | None) -> None:
         if not (
@@ -70,10 +72,50 @@ def record_temporary_table(self, expression: Expr, catalog: str | None, schema:
         # Forget dropped tables here rather than tracking DROP, so the registry is bounded by live tables.
         live_oids = set(live.values())
         self._temporary = {oid: owner for oid, owner in self._temporary.items() if oid in live_oids}
-        # CREATE TEMP TABLE IF NOT EXISTS succeeds without creating anything when the name is already taken, even
-        # from another schema, so the OID may already be registered. Keep the schema it was created in; overwriting
-        # it would move the existing table's stage to the current schema.
-        self._temporary.setdefault(live[table.name], (table.catalog or catalog, table.db or schema))
+        if dropped := self._metadata_oids - live_oids:
+            # Metadata is keyed by OID, so a dropped or replaced table's rows would otherwise never be removed.
+            oids = [str(oid) for oid in dropped]
+            for ext in ("_fs_columns_ext", "_fs_tables_ext"):
+                self._duck_conn.execute(
+                    f"DELETE FROM _fs_global._fs_information_schema.{ext} "
+                    "WHERE ext_table_catalog = 'temp' AND list_contains(?, ext_table_schema)",
+                    [oids],
+                )
+            self._metadata_oids &= live_oids
+        # CREATE TEMP TABLE IF NOT EXISTS succeeds without creating anything when the name is already taken, so
+        # the OID may already be registered. Keep the schema it was created in; overwriting it would move the
+        # existing table's stage to the current schema.
+        if not (owner := expression.args.get("temporary_owner")):
+            # Created directly in DuckDB's temp catalog, so it has no Snowflake schema to remember.
+            return
+        self._temporary.setdefault(live[table.name], owner)
+
+    def temporary_owners(self) -> dict[str, tuple[str | None, str | None]]:
+        """The (catalog, schema) each live temporary table was created in, by table name."""
+        if not self._temporary:
+            return {}
+        sql = "SELECT table_name, table_oid FROM duckdb_tables() WHERE temporary"
+        live = self._duck_conn.execute(sql).fetchall()
+        return {name: owner for name, oid in live if (owner := self._temporary.get(oid))}
+
+    def logical_schema(
+        self, table: exp.Table, catalog: str | None, schema: str | None
+    ) -> tuple[str | None, str | None]:
+        """The Snowflake (catalog, schema) of a table, including one rewritten to DuckDB's temp catalog."""
+        if table.catalog == "temp":
+            return self.temporary_owners().get(table.name) or (catalog, schema)
+        return table.catalog or catalog, table.db or schema
+
+    def metadata_owner(
+        self, table: exp.Table, catalog: str | None, schema: str | None
+    ) -> tuple[str | None, str | None]:
+        """The (catalog, schema) a table's Snowflake-only metadata is keyed by; temporary tables use their OID."""
+        if table.catalog == "temp":
+            sql = "SELECT table_oid FROM duckdb_tables() WHERE temporary AND table_name = ?"
+            if live := self._duck_conn.execute(sql, [table.name]).fetchone():
+                self._metadata_oids.add(live[0])
+                return "temp", str(live[0])
+        return table.catalog or catalog, table.db or schema
 
     def resolve_for_copy(
         self, catalog: str, schema: str, stage_name: str, target: tuple[str | None, str | None, str]

fakesnow/transforms/temporary.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import snowflake.connector.errors
+from duckdb import DuckDBPyConnection
+from sqlglot import Expr, exp
+from sqlglot.optimizer.scope import traverse_scope
+
+from fakesnow.transforms.stage import TableStages
+
+
+def temporary_tables(
+    expression: Expr,
+    duck_conn: DuckDBPyConnection,
+    current_database: str | None,
+    current_schema: str | None,
+    table_stages: TableStages,
+) -> Expr:
+    """Resolve the Snowflake names of temporary tables to DuckDB's temp catalog.
+
+    A Snowflake temporary table belongs to the schema it was created in. DuckDB keeps every temporary table in
+    temp.main and only accepts that catalog when creating one, so temporary tables are rewritten to temp.main
+    where they're created or referenced in their own schema, and unqualified references from another schema are
+    qualified so they don't reach the temporary table.
+    """
+    if expression.parent is not None:
+        return expression
+
+    owners = table_stages.temporary_owners()
+    target = _create_table_target(expression)
+
+    # a table created directly in DuckDB's temp catalog, like MERGE's scratch table, has no Snowflake schema
+    if target and target.catalog != "temp" and expression.find(exp.TemporaryProperty):
+        owner = (target.catalog or current_database, target.db or current_schema)
+        # DuckDB holds one temporary table per name, ignoring case, so only the same table can be replaced
+        existing = next(((n, o) for n, o in owners.items() if n.casefold() == target.name.casefold()), None)
+        if existing and existing != (target.name, owner):
+            name, (catalog, schema) = existing
+            raise snowflake.connector.errors.ProgrammingError(
+                msg=(
+                    f"SQL compilation error:\nfakesnow supports one temporary table named {target.name} per "
+                    f"session, ignoring case, and {name} already exists in {catalog}.{schema}."
+                ),
+                errno=2002,
+                sqlstate="42710",
+            )
+        sql = "SELECT 1 FROM duckdb_schemas() WHERE database_name = ? AND schema_name = ?"
+        if not duck_conn.execute(sql, owner).fetchone():
+            raise snowflake.connector.errors.ProgrammingError(
+                msg=f"SQL compilation error:\nSchema '{owner[0]}.{owner[1]}' does not exist or not authorized.",
+                errno=2003,
+                sqlstate="02000",
+            )
+        _to_temp(target)
+        # the cursor records this as the table's schema once the CREATE succeeds
+        expression.args["temporary_owner"] = owner
+
+    if not owners:
+        return expression
+
+    for table in _referenced_tables(expression, target):
+        if not (owner := owners.get(table.name)):
+            continue
+        if (table.catalog or current_database, table.db or current_schema) == owner:
+            _to_temp(table)
+        elif not table.db and current_database and current_schema:
+            # DuckDB finds a temporary table by its unqualified name from any schema, Snowflake only from its own
+            table.set("catalog", exp.to_identifier(current_database))
+            table.set("db", exp.to_identifier(current_schema))
+
+    return expression
+
+
+def _create_table_target(expression: Expr) -> exp.Table | None:
+    if isinstance(expression, exp.Create) and expression.args.get("kind") == "TABLE":
+        this = expression.this
+        table = this.this if isinstance(this, exp.Schema) else this
+        if isinstance(table, exp.Table):
+            return table
+    return None
+
+
+def _referenced_tables(expression: Expr, create_target: exp.Table | None) -> list[exp.Table]:
+    """Tables a statement reads or writes, other than the one it creates."""
+    tables = [
+        source
+        for scope in traverse_scope(expression)
+        for source in scope.sources.values()
+        if isinstance(source, exp.Table) and isinstance(source.this, exp.Identifier)
+    ]
+
+    # DESCRIBE's kind is lower case, DROP keeps its tables in a list
+    is_table_statement = str(expression.args.get("kind")).upper() == "TABLE"
+    if isinstance(expression, (exp.Insert, exp.Copy, exp.Update, exp.Delete)) or (
+        isinstance(expression, exp.Describe) and is_table_statement
+    ):
+        this = expression.this
+        table = this.this if isinstance(this, exp.Schema) else this
+        if isinstance(table, exp.Table):
+            tables.append(table)
+    elif isinstance(expression, exp.Drop) and is_table_statement:
+        tables.extend(t for t in expression.args.get("tables") or [] if isinstance(t, exp.Table))
+
+    # the same name can be referenced more than once, and each reference is rewritten
+    unique = {id(table): table for table in tables if table is not create_target}
+    return list(unique.values())
+
+
+def _to_temp(table: exp.Table) -> None:
+    table.set("catalog", exp.to_identifier("temp"))
+    table.set("db", exp.to_identifier("main"))

fakesnow/transforms/transforms.py

@@ -164,16 +164,8 @@ def create_clone(expression: Expr) -> Expr:
         and str(expression.args.get("kind")).upper() == "TABLE"
         and (clone := expression.find(exp.Clone))
     ):
-        return exp.Create(
-            this=expression.this,
-            kind="TABLE",
-            expression=exp.Select(
-                expressions=[
-                    exp.Star(),
-                ],
-                from_=exp.From(this=clone.this),
-            ),
-        )
+        expression.set("clone", None)
+        expression.set("expression", exp.Select(expressions=[exp.Star()], from_=exp.From(this=clone.this)))
     return expression
 
 
@@ -242,7 +234,7 @@ def create_database(expression: Expr, db_path: Path | None = None) -> Expr:
     NULL::VARCHAR AS "comment",
     NULL::VARCHAR AS "policy name",
     NULL::JSON AS "privacy domain",
-FROM _fs_information_schema._fs_columns
+FROM _fs_global._fs_information_schema._fs_columns
 WHERE table_catalog = '${catalog}' AND table_schema = '${schema}' AND table_name = '${table}'
 ORDER BY ordinal_position
 """

tests/test_copy_into.py

@@ -568,6 +568,35 @@ def test_copy_temp_table_stage_after_schema_switch(dcur: snowflake.connector.cur
     assert dcur.fetchall() == [{"status": "Copy executed with 0 files processed."}]
 
 
+def test_copy_into_qualified_temporary_table_from_its_stage(
+    dcur: snowflake.connector.cursor.DictCursor, tmp_path: Path
+) -> None:
+    dcur.execute("CREATE TEMP TABLE db1.schema1.temporary_stage_table (a INT, b INT)")
+    path = tmp_path / "data.csv"
+    path.write_text("1,2\n")
+    dcur.execute(f"PUT 'file://{path}' @%temporary_stage_table AUTO_COMPRESS=FALSE")
+
+    dcur.execute("COPY INTO db1.schema1.temporary_stage_table FROM @%temporary_stage_table")
+
+    dcur.execute("SELECT * FROM temporary_stage_table")
+    assert dcur.fetchall() == [{"A": 1, "B": 2}]
+
+
+def test_copy_same_file_into_replaced_temporary_table(
+    dcur: snowflake.connector.cursor.DictCursor, tmp_path: Path
+) -> None:
+    # a replaced table is a new table, so it has no load history
+    path = tmp_path / "data.csv"
+    path.write_text("1,2\n")
+    dcur.execute("CREATE STAGE reload_stage")
+    dcur.execute(f"PUT 'file://{path}' @reload_stage AUTO_COMPRESS=FALSE")
+
+    for _ in range(2):
+        dcur.execute("CREATE OR REPLACE TEMP TABLE db1.schema1.reloaded_table (a INT, b INT)")
+        dcur.execute("COPY INTO db1.schema1.reloaded_table FROM @reload_stage")
+        assert [r["status"] for r in dcur.fetchall()] == ["LOADED"]
+
+
 @pytest.fixture
 def temporary_table_stage_sessions(
     _fakesnow: None, tmp_path: Path