Support Named Schemas

Author: waltaskew

google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

@@ -233,6 +233,10 @@ class SpannerSQLCompiler(SQLCompiler):
 
     compound_keywords = _compound_keywords
 
+    def __init__(self, *args, **kwargs):
+        self.tablealiases = {}
+        super().__init__(*args, **kwargs)
+
     def get_from_hint_text(self, _, text):
         """Return a hint text.
 
@@ -378,8 +382,10 @@ def limit_clause(self, select, **kw):
         return text
 
     def returning_clause(self, stmt, returning_cols, **kw):
+        # Set include_table=False because although table names are allowed in
+        # RETURNING clauses, schema names are not.
         columns = [
-            self._label_select_column(None, c, True, False, {})
+            self._label_select_column(None, c, True, False, {}, include_table=False)
             for c in expression._select_iterables(returning_cols)
         ]
 
@@ -391,6 +397,87 @@ def visit_sequence(self, seq, **kw):
             seq
         )
 
+    def visit_table(self, table, spanner_aliased=False, iscrud=False, **kwargs):
+        """Produces the table name.
+
+        Schema names are not allowed in Spanner SELECT statements. We
+        need to avoid generating SQL like
+
+        SELECT schema.tbl.id
+        FROM schema.tbl
+
+        To do so, we alias the table in order to produce SQL like:
+
+        SELECT tbl_1.id, tbl_1.col
+        FROM schema.tbl AS tbl_1
+
+        And do similar for UPDATE and DELETE statements.
+
+        We don't need to correct INSERT statements, which is fortunate
+        because INSERT statements actually do not currently result in
+        calls to `visit_table`.
+
+        This closely mirrors the mssql dialect which also avoids
+        schema-qualified columns in SELECTs, although the behaviour is
+        currently behind a deprecated 'legacy_schema_aliasing' flag.
+        """
+        if spanner_aliased is table or self.isinsert:
+            return super().visit_table(table, **kwargs)
+
+        # alias schema-qualified tables
+        alias = self._schema_aliased_table(table)
+        if alias is not None:
+            return self.process(alias, spanner_aliased=table, **kwargs)
+        else:
+            return super().visit_table(table, **kwargs)
+
+    def visit_alias(self, alias, **kw):
+        """Produces alias statements."""
+        # translate for schema-qualified table aliases
+        kw["spanner_aliased"] = alias.element
+        return super().visit_alias(alias, **kw)
+
+    def visit_column(self, column, add_to_result_map=None, **kw):
+        """Produces column expressions.
+
+        In tandem with visit_table, replaces schema-qualified column
+        names with column names qualified against an alias.
+        """
+        if (
+            column.table is not None
+            and not self.isinsert
+            or self.is_subquery()
+        ):
+            # translate for schema-qualified table aliases
+            t = self._schema_aliased_table(column.table)
+            if t is not None:
+                converted = elements._corresponding_column_or_error(t, column)
+                if add_to_result_map is not None:
+                    add_to_result_map(
+                        column.name,
+                        column.name,
+                        (column, column.name, column.key),
+                        column.type,
+                    )
+
+                return super().visit_column(converted, **kw)
+
+        return super().visit_column(column, add_to_result_map=add_to_result_map, **kw)
+
+    def _schema_aliased_table(self, table):
+        """Creates an alias for the table if it is schema-qualified.
+
+        If the table is schema-qualified, returns an alias for the
+        table and caches the alias for future references to the
+        table. If the table is not schema-qualified, returns None.
+        """
+        if getattr(table, "schema", None) is not None:
+            if table not in self.tablealiases:
+                self.tablealiases[table] = table.alias()
+            return self.tablealiases[table]
+        else:
+            return None
+
 
 class SpannerDDLCompiler(DDLCompiler):
     """Spanner DDL statements compiler."""

test/system/test_basics.py

@@ -25,6 +25,8 @@
     Boolean,
     BIGINT,
     select,
+    update,
+    delete,
 )
 from sqlalchemy.orm import Session, DeclarativeBase, Mapped, mapped_column
 from sqlalchemy.types import REAL
@@ -58,6 +60,16 @@ def define_tables(cls, metadata):
             Column("name", String(20)),
         )
 
+        with cls.bind.begin() as conn:
+            conn.execute(text("CREATE SCHEMA IF NOT EXISTS schema"))
+        Table(
+            "users",
+            metadata,
+            Column("ID", Integer, primary_key=True),
+            Column("name", String(20)),
+            schema="schema",
+        )
+
     def test_hello_world(self, connection):
         greeting = connection.execute(text("select 'Hello World'"))
         eq_("Hello World", greeting.fetchone()[0])
@@ -139,6 +151,12 @@ class User(Base):
             ID: Mapped[int] = mapped_column(primary_key=True)
             name: Mapped[str] = mapped_column(String(20))
 
+        class SchemaUser(Base):
+            __tablename__ = "users"
+            __table_args__ = {"schema": "schema"}
+            ID: Mapped[int] = mapped_column(primary_key=True)
+            name: Mapped[str] = mapped_column(String(20))
+
         engine = connection.engine
         with Session(engine) as session:
             number = Number(
@@ -156,3 +174,25 @@ class User(Base):
             users = session.scalars(statement).all()
             eq_(1, len(users))
             is_true(users[0].ID > 0)
+
+        with Session(engine) as session:
+            user = SchemaUser(name="SchemaTest")
+            session.add(user)
+            session.commit()
+
+            users = session.scalars(select(SchemaUser).where(SchemaUser.name == "SchemaTest")).all()
+            eq_(1, len(users))
+            is_true(users[0].ID > 0)
+
+            session.execute(update(SchemaUser).where(SchemaUser.name == "SchemaTest").values(name="NewName"))
+            session.commit()
+
+            users = session.scalars(select(SchemaUser).where(SchemaUser.name == "NewName")).all()
+            eq_(1, len(users))
+            is_true(users[0].ID > 0)
+
+            session.execute(delete(SchemaUser).where(SchemaUser.name=="NewName"))
+            session.commit()
+
+            users = session.scalars(select(SchemaUser).where(SchemaUser.name=="NewName")).all()
+            eq_(0, len(users))