feat: support named schemas

Add support for SELECT, UPDATE and DELETE statements against tables in
schemas.

Schema names are not allowed in Spanner SELECT statements. We need to
avoid generating SQL like

```sql
SELECT schema.tbl.id
FROM schema.tbl
```

To do so, we alias the table in order to produce SQL like:

```sql
SELECT tbl_1.id, tbl_1.col
FROM schema.tbl AS tbl_1
```

Author: waltaskew

README.rst

@@ -465,7 +465,6 @@ Other limitations
 ~~~~~~~~~~~~~~~~~
 
 -  WITH RECURSIVE statement is not supported.
--  Named schemas are not supported.
 -  Temporary tables are not supported.
 -  Numeric type dimensions (scale and precision) are constant. See the
    `docs <https://cloud.google.com/spanner/docs/data-types#numeric_types>`__.

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,83 @@ 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,35 @@ 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))