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,66 @@ def visit_sequence(self, seq, **kw):
             seq
         )
 
+    def visit_table(self, table, spanner_aliased=False, iscrud=False, **kwargs):
+        """Build the table name.
+
+        Schema names are not allowed in Spanner SELECT statements. When selecting
+        from a schema-qualified table, alias the table to produce SQL like:
+
+        SELECT tbl_1.id, tbl_1.col
+        FROM schema.tbl AS tbl_1
+        """
+        # This closely code 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 iscrud:
+            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):
+        # 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):
+        if (
+            column.table is not None
+            and (not self.isupdate and not self.isdelete 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):
+        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

@@ -31,6 +31,11 @@
 from sqlalchemy.testing import eq_, is_true
 from sqlalchemy.testing.plugin.plugin_base import fixtures
 
+import logging
+
+logging.basicConfig()
+logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+logging.getLogger("sqlalchemy.pool").setLevel(logging.DEBUG)
 
 class TestBasics(fixtures.TablesTest):
     @classmethod
@@ -58,6 +63,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 +154,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 +177,13 @@ 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()
+
+            statement = select(SchemaUser).filter_by(name="SchemaTest")
+            users = session.scalars(statement).all()
+            eq_(1, len(users))
+            is_true(users[0].ID > 0)