fix: resolve database and schema qualified temporary tables as Snowflake does

fakesnow/conn.py

@@ -15,6 +15,7 @@
 import fakesnow.macros as macros
 from fakesnow.cursor import FakeSnowflakeCursor
 from fakesnow.dialect import parse
+from fakesnow.temporary_tables import TemporaryTables
 from fakesnow.transforms.stage import TableStages
 from fakesnow.variables import Variables
 
@@ -52,7 +53,8 @@ def __init__(
         self.nop_regexes = nop_regexes
         self._paramstyle = kwargs.get("paramstyle", snowflake.connector.paramstyle)
         self.variables = Variables()
-        self.table_stages = TableStages(duck_conn)
+        self.temporary_tables = TemporaryTables(duck_conn)
+        self.table_stages = TableStages(duck_conn, self.temporary_tables)
         self.results_cache = results_cache
         self._autocommit = kwargs.get("autocommit", True)
         self._in_transaction = False

fakesnow/cursor.py

@@ -301,6 +301,13 @@ 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(
+                self._conn.temporary_tables.transform,
+                current_database=self._conn.database,
+                current_schema=self._conn.schema,
+                in_transaction=self._conn._in_transaction,
+                copy=False,
+            )
             .transform(transforms.alter_session, copy=False)
             .transform(transforms.update_variables, variables=self._conn.variables, copy=False)
             .transform(transforms.current_version, copy=False)
@@ -478,11 +485,7 @@ def _execute(self, transformed: Expr, params: MutableParams | None = None) -> No
                 self._duck_conn.begin()
                 self._conn._in_transaction = True
 
-            # In transaction book keeping
-            elif isinstance(transformed, exp.Transaction):
-                self._conn._in_transaction = True
-            elif isinstance(transformed, (exp.Commit, exp.Rollback)):
-                self._conn._in_transaction = False
+        self._conn.temporary_tables.validate(transformed, self._conn._in_transaction)
 
         affected_count = None
         try:
@@ -534,7 +537,15 @@ def _execute(self, transformed: Expr, params: MutableParams | None = None) -> No
             msg = cast(str, e.args[0]).split("\n")[0]
             raise snowflake.connector.errors.ProgrammingError(msg=msg, errno=100035, sqlstate="22007") from e
 
-        self._conn.table_stages.record_temporary_table(transformed, self._conn.database, self._conn.schema)
+        # In transaction book keeping, including BEGIN in autocommit mode, whose transaction stays open if autocommit
+        # is then turned off. Only once the statement succeeds: a failed BEGIN opens no transaction, and a failed
+        # COMMIT or ROLLBACK leaves it open.
+        if isinstance(transformed, exp.Transaction):
+            self._conn._in_transaction = True
+        elif isinstance(transformed, (exp.Commit, exp.Rollback)):
+            self._conn._in_transaction = False
+
+        self._conn.temporary_tables.record(transformed)
 
         if set_database := transformed.args.get("set_database"):
             self._conn.database = set_database

fakesnow/temporary_tables.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import duckdb
+import snowflake.connector.errors
+from duckdb import DuckDBPyConnection
+from sqlglot import Expr, exp
+
+
+class TemporaryTables:
+    """Track Snowflake schema ownership for DuckDB's session-local temporary tables.
+
+    Table OIDs survive renames and are never reused. Lookups return only live OIDs. Registrations aren't pruned when
+    a table is dropped, because ROLLBACK can restore it under the same OID. Both SQL references and table stages use
+    this registry.
+    """
+
+    def __init__(self, duck_conn: DuckDBPyConnection):
+        self._duck_conn = duck_conn
+        # Temporary table OID -> (catalog, schema, schema OID) it was created in.
+        self._temporary: dict[int, tuple[str | None, str | None, int]] = {}
+        # Names of the registered temporary tables live when this statement started, plus any created since, and
+        # while a transaction is open, any live since it began, which ROLLBACK can restore. A qualified reference to
+        # any other name can't be a temporary table, so it needs no catalog read.
+        self._names: set[str] = set()
+
+    def transform(
+        self, expression: Expr, current_database: str | None, current_schema: str | None, in_transaction: bool
+    ) -> Expr:
+        """Create temporary tables in DuckDB's temp catalog and resolve their qualified references.
+
+        Snowflake temporary tables belong to a database and schema, whereas DuckDB keeps them all in temp.main.
+        See https://docs.snowflake.com/en/user-guide/tables-temp-transient
+        """
+        if expression.parent is None and self._temporary:
+            self._drop_orphans(in_transaction)
+        if (
+            isinstance(expression, exp.Create)
+            and expression.kind == "TABLE"
+            and expression.find(exp.TemporaryProperty)
+            and (table := expression.find(exp.Table))
+        ):
+            catalog, schema = table.catalog or current_database, table.db or current_schema
+            # Removing the qualifiers bypasses DuckDB's schema validation, so check the original owner first.
+            try:
+                schema_oid = self._schema_oid(catalog, schema)
+            except duckdb.TransactionException:
+                # An aborted transaction rejects the lookup. Outside autocommit, DDL ends the transaction before it
+                # runs, so validate() checks the schema then, before the table is created or replaced.
+                schema_oid = None
+            expression.set("temporary_table_owner", (catalog, schema, schema_oid))
+            table.set("catalog", None)
+            table.set("db", None)
+            # Snowflake doesn't enforce foreign keys, and DuckDB won't drop a table that one references, which would
+            # stop a temporary table being dropped with its schema.
+            for node in list(expression.find_all(exp.ForeignKey, exp.Reference)):
+                if isinstance(node.parent, (exp.ColumnConstraint, exp.Constraint)):
+                    node.parent.pop()
+                elif isinstance(node, exp.ForeignKey):
+                    node.pop()
+        elif isinstance(expression, (exp.Table, exp.Column)) and expression.db:
+            parent = expression.parent
+            # A permanent CREATE target can share a temporary table's name.
+            if isinstance(parent, exp.Create) or (
+                isinstance(parent, exp.Schema) and isinstance(parent.parent, exp.Create)
+            ):
+                return expression
+            # Non-table DDL also uses Table nodes, for example DROP SEQUENCE.
+            if isinstance(parent, (exp.Drop, exp.Alter)) and parent.args.get("kind") != "TABLE":
+                return expression
+            # COPY resolves its target by its Snowflake name for table stages and load history.
+            if isinstance(parent, exp.Copy):
+                return expression
+            # Foreign keys left in a CREATE are a permanent table's, which reference permanent tables.
+            if expression.find_ancestor(exp.Reference):
+                return expression
+            name = expression.table if isinstance(expression, exp.Column) else expression.name
+            # Only read the catalog for a name that could be a temporary table's: in an aborted transaction the read
+            # fails, see #470. A reference that does match one still fails then, even in DDL, which outside autocommit
+            # would otherwise end the transaction and run.
+            if (
+                name in self._names
+                and self.oid(expression.catalog or current_database, expression.db, name) is not None
+            ):
+                expression.set("catalog", exp.to_identifier("temp"))
+                expression.set("db", exp.to_identifier("main"))
+        return expression
+
+    def validate(self, expression: Expr, in_transaction: bool) -> None:
+        """Check the schema of a temporary table that transform() couldn't, just before it's created.
+
+        Called once DDL has ended the aborted transaction, so nothing is created or replaced in a missing schema. In
+        autocommit mode the transaction is still aborted, and the lookup fails as the CREATE itself would.
+        """
+        if (owner := expression.args.get("temporary_table_owner")) and owner[2] is None:
+            # transform() couldn't drop orphans either, and one could take the new table's name.
+            self._drop_orphans(in_transaction)
+            catalog, schema, _ = owner
+            expression.set("temporary_table_owner", (catalog, schema, self._schema_oid(catalog, schema)))
+
+    def record(self, expression: Expr) -> None:
+        if not (
+            isinstance(expression, exp.Create)
+            and (owner := expression.args.get("temporary_table_owner"))
+            and (table := expression.find(exp.Table))
+        ):
+            return
+        # Temporary table names are unique within the session, so the name identifies the new table's OID.
+        sql = "SELECT table_name, table_oid FROM duckdb_tables() WHERE temporary"
+        oid = dict(self._duck_conn.execute(sql).fetchall())[table.name]
+        # 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(oid, owner)
+        self._names.add(table.name)
+
+    def _drop_orphans(self, in_transaction: bool) -> None:
+        """Drop temporary tables whose schema has been dropped or replaced since they were created, by any session.
+
+        In Snowflake dropping a schema drops its temporary tables, which DuckDB keeps in temp.main instead. Checking
+        the schema's OID before each statement covers other sessions and ROLLBACK, which restores the original OID.
+        """
+        sql = "SELECT table_oid, table_name FROM duckdb_tables() WHERE temporary"
+        try:
+            live = self._duck_conn.execute(sql).fetchall()
+        except duckdb.TransactionException:
+            # An aborted transaction rejects every query until it ends, so leave them to the next statement. Until it
+            # ends, nothing can rename, drop or restore a temporary table, so the names still hold.
+            return
+        owned = [(name, self._temporary[oid]) for oid, name in live if oid in self._temporary]
+        # Refresh the names, which renames, drops and ROLLBACK change. Ending a transaction, even by DDL once a failed
+        # statement has aborted it, can restore the names live when it began, so until then only add to them.
+        names = {name for name, _ in owned}
+        self._names = self._names | names if in_transaction else names
+        if not owned:
+            return
+        sql = "SELECT database_name, schema_name, oid FROM duckdb_schemas()"
+        schemas = {(catalog, schema): oid for catalog, schema, oid in self._duck_conn.execute(sql).fetchall()}
+        for name, (catalog, schema, schema_oid) in owned:
+            if schemas.get((catalog, schema)) != schema_oid:
+                table = exp.table_(name, db="main", catalog="temp", quoted=True)
+                self._duck_conn.execute(f"DROP TABLE {table.sql(dialect='duckdb')}")
+                if not in_transaction:
+                    self._names.discard(name)
+
+    def _schema_oid(self, catalog: str | None, schema: str | None) -> int:
+        sql = "SELECT oid FROM duckdb_schemas() WHERE database_name = ? AND schema_name = ?"
+        if not (row := self._duck_conn.execute(sql, [catalog, schema]).fetchone()):
+            raise snowflake.connector.errors.ProgrammingError(
+                msg=f"SQL compilation error:\nSchema '{catalog}.{schema}' does not exist or not authorized.",
+                errno=2003,
+                sqlstate="42S02",
+            )
+        return row[0]
+
+    def oid(self, catalog: str | None, schema: str, name: str) -> int | None:
+        if not self._temporary:
+            return None
+        # At most one temporary table has this name in the session. Return its OID only if it was created in the
+        # requested schema, and that schema hasn't since been replaced.
+        sql = """
+            SELECT table_oid, (SELECT oid FROM duckdb_schemas() WHERE database_name = ? AND schema_name = ?)
+            FROM duckdb_tables() WHERE temporary AND table_name = ?
+        """
+        live = self._duck_conn.execute(sql, [catalog, schema, name]).fetchone()
+        return live[0] if live and self._temporary.get(live[0]) == (catalog, schema, live[1]) else None

fakesnow/transforms/stage.py

@@ -20,6 +20,7 @@
 
 from fakesnow.expr import normalise_ident
 from fakesnow.params import MutableParams
+from fakesnow.temporary_tables import TemporaryTables
 from fakesnow.transforms.options import parse_options
 
 # TODO: clean up temp files on exit
@@ -43,37 +44,12 @@ class UploadCommandDict(TypedDict):
 
 
 class TableStages:
-    """Resolve implicit table stages, including session-local temporary tables.
+    """Resolve implicit table stages, including session-local temporary tables."""
 
-    DuckDB keeps temporary tables in one session namespace, so remember the Snowflake schema each was created in.
-    Table OIDs survive renames and are never reused. Lookups embed only the live table's OID, and registrations
-    for dropped tables are pruned when another temporary table is created, so session history stays bounded.
-    """
-
-    def __init__(self, duck_conn: DuckDBPyConnection):
+    def __init__(self, duck_conn: DuckDBPyConnection, temporary_tables: TemporaryTables):
         self._duck_conn = duck_conn
         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]] = {}
-
-    def record_temporary_table(self, expression: Expr, catalog: str | None, schema: str | None) -> None:
-        if not (
-            isinstance(expression, exp.Create)
-            and expression.args.get("kind") == "TABLE"
-            and expression.find(exp.TemporaryProperty)
-            and (table := expression.find(exp.Table))
-        ):
-            return
-        # Temporary table names are unique within the session, so the name identifies the new table's OID.
-        sql = "SELECT table_name, table_oid FROM duckdb_tables() WHERE temporary"
-        live = dict(self._duck_conn.execute(sql).fetchall())
-        # 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))
+        self._temporary_tables = temporary_tables
 
     def resolve_for_copy(
         self, catalog: str, schema: str, stage_name: str, target: tuple[str | None, str | None, str]
@@ -99,18 +75,16 @@ def lookup_sql(self, catalog: str, schema: str, stage_name: str) -> str:
 
         A temporary table shadows a permanent table of the same name, but only in the schema it was created in.
         """
-        # At most one temporary table has this name in the session. Embed its OID only if it was created in the
-        # requested schema; otherwise match nothing (NULL), leaving any permanent table to resolve.
-        sql = "SELECT table_oid FROM duckdb_tables() WHERE temporary AND table_name = ?"
-        live = self._duck_conn.execute(sql, [stage_name[1:]]).fetchone()
-        oid = live[0] if live and self._temporary.get(live[0]) == (catalog, schema) else "NULL"
+        # Embed the temporary table's OID only if it was created in the requested schema; otherwise match nothing
+        # (NULL), leaving any permanent table to resolve.
+        oid = self._temporary_tables.oid(catalog, schema, stage_name[1:])
         # Temporary storage is keyed by session and OID, so it is isolated between sessions and follows renames.
         logical_name = exp.Literal.string(f"{catalog}.{schema}.{stage_name}").sql(dialect="duckdb")
         return f"""
             SELECT CASE WHEN temporary THEN '{self._session}.' || table_oid || '.%' ELSE {logical_name} END
             FROM duckdb_tables()
             WHERE table_name = '{stage_name[1:]}' AND (
-                (temporary AND table_oid = {oid})
+                (temporary AND table_oid = {"NULL" if oid is None else oid})
                 OR (NOT temporary AND database_name = '{catalog}' AND schema_name = '{schema}')
             )
             ORDER BY temporary DESC

fakesnow/transforms/transforms.py

@@ -164,7 +164,7 @@ def create_clone(expression: Expr) -> Expr:
         and str(expression.args.get("kind")).upper() == "TABLE"
         and (clone := expression.find(exp.Clone))
     ):
-        return exp.Create(
+        create = exp.Create(
             this=expression.this,
             kind="TABLE",
             expression=exp.Select(
@@ -174,6 +174,11 @@ def create_clone(expression: Expr) -> Expr:
                 from_=exp.From(this=clone.this),
             ),
         )
+        if expression.find(exp.TemporaryProperty):
+            # keep a temporary clone temporary, in the schema TemporaryTables resolved it to
+            create.set("properties", exp.Properties(expressions=[exp.TemporaryProperty()]))
+            create.set("temporary_table_owner", expression.args.get("temporary_table_owner"))
+        return create
     return expression