googleapis · olavloite · May 2, 2025 · Apr 22, 2025 · May 1, 2025 · May 2, 2025
diff --git a/README.rst b/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>`__.

diff --git a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/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,11 @@ def limit_clause(self, select, **kw):
        return text

    def returning_clause(self, stmt, returning_cols, **kw):
+        # Set the spanner_is_returning flag which is passed to visit_column.
        columns = [
-            self._label_select_column(None, c, True, False, {})
+            self._label_select_column(
+                None, c, True, False, {"spanner_is_returning": True}
+            )
            for c in expression._select_iterables(returning_cols)
        ]

@@ -391,6 +398,98 @@ 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.
+
+        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)
+
+        # Add an alias for schema-qualified tables.
+        # Tables in the default schema are not aliased and follow the
+        # standard SQLAlchemy code path.
+        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, spanner_is_returning=False, **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)
+        if spanner_is_returning:
+            # Set include_table=False because although table names are
+            # allowed in RETURNING clauses, schema names are not.  We
+            # can't use the same aliasing trick above that we use with
+            # other statements, because INSERT statements don't result
+            # in visit_table calls and INSERT table names can't be
+            # aliased.  Statements like:
+            #
+            # INSERT INTO table (id, name)
+            # SELECT id, name FROM another_table
+            # THEN RETURN another_table.id
+            #
+            # aren't legal, so the columns remain unambiguous when not
+            # qualified by table name.
+            kw["include_table"] = False
+
+        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."""

diff --git a/test/mockserver_tests/test_auto_increment.py b/test/mockserver_tests/test_auto_increment.py
@@ -125,9 +125,7 @@ def test_create_table_with_specific_sequence_kind(self):
    def test_insert_row(self):
        from test.mockserver_tests.auto_increment_model import Singer

-        self.add_insert_result(
-            "INSERT INTO singers (name) VALUES (@a0) THEN RETURN singers.id"
-        )
+        self.add_insert_result("INSERT INTO singers (name) VALUES (@a0) THEN RETURN id")
        engine = create_engine(
            "spanner:///projects/p/instances/i/databases/d",
            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},

diff --git a/test/mockserver_tests/test_basics.py b/test/mockserver_tests/test_basics.py
@@ -262,3 +262,53 @@ class Singer(Base):
            singer.name = "New Name"
            session.add(singer)
            session.commit()
+
+    def test_select_table_in_named_schema(self):
+        class Base(DeclarativeBase):
+            pass
+
+        class Singer(Base):
+            __tablename__ = "singers"
+            __table_args__ = {"schema": "my_schema"}
+            id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
+            name: Mapped[str] = mapped_column(String)
+
+        query = (
+            "SELECT"
+            " singers_1.id AS my_schema_singers_id,"
+            " singers_1.name AS my_schema_singers_name\n"
+            "FROM my_schema.singers AS singers_1\n"
+            "WHERE singers_1.id = @a0\n"
+            " LIMIT @a1"
+        )
+        add_singer_query_result(query)
+        engine = create_engine(
+            "spanner:///projects/p/instances/i/databases/d",
+            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+        )
+
+        insert = "INSERT INTO my_schema.singers (name) VALUES (@a0) THEN RETURN id"
+        add_single_result(insert, "id", TypeCode.INT64, [("1",)])
+        with Session(engine) as session:
+            singer = Singer(name="New Name")
+            session.add(singer)
+            session.commit()
+
+        update = (
+            "UPDATE my_schema.singers AS singers_1 "
+            "SET name=@a0 "
+            "WHERE singers_1.id = @a1"
+        )
+        add_update_count(update, 1)
+        with Session(engine) as session:
+            singer = session.query(Singer).filter(Singer.id == 1).first()
+            singer.name = "New Name"
+            session.add(singer)
+            session.commit()
+
+        delete = "DELETE FROM my_schema.singers AS singers_1 WHERE singers_1.id = @a0"
+        add_update_count(delete, 1)
+        with Session(engine) as session:
+            singer = session.query(Singer).filter(Singer.id == 1).first()
+            session.delete(singer)
+            session.commit()
diff --git a/test/mockserver_tests/test_bit_reversed_sequence.py b/test/mockserver_tests/test_bit_reversed_sequence.py
@@ -110,7 +110,7 @@ def test_insert_row(self):
        add_result(
            "INSERT INTO singers (id, name) "
            "VALUES ( GET_NEXT_SEQUENCE_VALUE(SEQUENCE singer_id), @a0) "
-            "THEN RETURN singers.id",
+            "THEN RETURN id",
            result,
        )
        engine = create_engine(

diff --git a/test/system/test_basics.py b/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))