feat: support sequences

google/cloud/sqlalchemy_spanner/requirements.py

@@ -81,7 +81,7 @@ def isolation_level(self):
 
     @property
     def sequences(self):
-        return exclusions.closed()
+        return exclusions.open()
 
     @property
     def temporary_tables(self):

google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

@@ -40,6 +40,7 @@
 )
 from sqlalchemy.sql.default_comparator import operator_lookup
 from sqlalchemy.sql.operators import json_getitem_op
+from sqlalchemy.sql import expression
 
 from google.cloud.spanner_v1.data_types import JsonObject
 from google.cloud import spanner_dbapi
@@ -173,6 +174,16 @@ def pre_exec(self):
         if priority is not None:
             self._dbapi_connection.connection.request_priority = priority
 
+    def fire_sequence(self, seq, type_):
+        """Builds a statement for fetching next value of the sequence."""
+        return self._execute_scalar(
+            (
+                "SELECT GET_NEXT_SEQUENCE_VALUE(SEQUENCE %s)"
+                % self.identifier_preparer.format_sequence(seq)
+            ),
+            type_,
+        )
+
 
 class SpannerIdentifierPreparer(IdentifierPreparer):
     """Identifiers compiler.
@@ -343,6 +354,20 @@ def limit_clause(self, select, **kw):
             text += " OFFSET " + self.process(select._offset_clause, **kw)
         return text
 
+    def returning_clause(self, stmt, returning_cols, **kw):
+        columns = [
+            self._label_select_column(None, c, True, False, {})
+            for c in expression._select_iterables(returning_cols)
+        ]
+
+        return "THEN RETURN " + ", ".join(columns)
+
+    def visit_sequence(self, seq, **kw):
+        """Builds a statement for fetching next value of the sequence."""
+        return " GET_NEXT_SEQUENCE_VALUE(SEQUENCE %s)" % self.preparer.format_sequence(
+            seq
+        )
+
 
 class SpannerDDLCompiler(DDLCompiler):
     """Spanner DDL statements compiler."""
@@ -457,6 +482,24 @@ def post_create_table(self, table):
 
         return post_cmds
 
+    def get_identity_options(self, identity_options):
+        text = ["sequence_kind = 'bit_reversed_positive'"]
+        if identity_options.start is not None:
+            text.append("start_with_counter = %d" % identity_options.start)
+        return ", ".join(text)
+
+    def visit_create_sequence(self, create, prefix=None, **kw):
+        """Builds a ``CREATE SEQUENCE`` statement for the sequence."""
+        text = "CREATE SEQUENCE %s" % self.preparer.format_sequence(create.element)
+        options = self.get_identity_options(create.element)
+        if options:
+            text += " OPTIONS (" + options + ")"
+        return text
+
+    def visit_drop_sequence(self, drop, **kw):
+        """Builds a ``DROP SEQUENCE`` statement for the sequence."""
+        return "DROP SEQUENCE %s" % self.preparer.format_sequence(drop.element)
+
 
 class SpannerTypeCompiler(GenericTypeCompiler):
     """Spanner types compiler.
@@ -531,7 +574,8 @@ class SpannerDialect(DefaultDialect):
     supports_sane_rowcount = False
     supports_sane_multi_rowcount = False
     supports_default_values = False
-    supports_sequences = False
+    supports_sequences = True
+    sequences_optional = False
     supports_native_enum = True
     supports_native_boolean = True
     supports_native_decimal = True
@@ -694,6 +738,36 @@ def get_view_names(self, connection, schema=None, **kw):
 
         return all_views
 
+    @engine_to_connection
+    def get_sequence_names(self, connection, schema=None, **kw):
+        """
+        Return a list of all sequence names available in the database.
+
+        The method is used by SQLAlchemy introspection systems.
+
+        Args:
+            connection (sqlalchemy.engine.base.Connection):
+                SQLAlchemy connection or engine object.
+            schema (str): Optional. Schema name
+
+        Returns:
+            list: List of sequence names.
+        """
+        sql = """
+            SELECT name
+            FROM information_schema.sequences
+            WHERE SCHEMA='{}'
+            """.format(
+            schema or ""
+        )
+        all_sequences = []
+        with connection.connection.database.snapshot() as snap:
+            rows = list(snap.execute_sql(sql))
+            for seq in rows:
+                all_sequences.append(seq[0])
+
+        return all_sequences
+
     @engine_to_connection
     def get_view_definition(self, connection, view_name, schema=None, **kw):
         """
@@ -1294,6 +1368,32 @@ def has_table(self, connection, table_name, schema=None, **kw):
 
         return False
 
+    @engine_to_connection
+    def has_sequence(self, connection, sequence_name, schema=None, **kw):
+        """Check the existence of a particular sequence in the database.
+
+        Given a :class:`_engine.Connection` object and a string
+        `sequence_name`, return True if the given sequence exists in
+        the database, False otherwise.
+        """
+
+        with connection.connection.database.snapshot() as snap:
+            rows = snap.execute_sql(
+                """
+                SELECT true
+                FROM INFORMATION_SCHEMA.SEQUENCES
+                WHERE NAME="{sequence_name}"
+                LIMIT 1
+                """.format(
+                    sequence_name=sequence_name
+                )
+            )
+
+            for _ in rows:
+                return True
+
+        return False
+
     def set_isolation_level(self, conn_proxy, level):
         """Set the connection isolation level.
 

test/test_suite_13.py

@@ -37,6 +37,7 @@
 from sqlalchemy.testing import config
 from sqlalchemy.testing import engines
 from sqlalchemy.testing import eq_
+from sqlalchemy.testing import is_instance_of
 from sqlalchemy.testing import provide_metadata, emits_warning
 from sqlalchemy.testing import fixtures
 from sqlalchemy.testing import is_true
@@ -73,7 +74,10 @@
 from sqlalchemy.testing.suite.test_reflection import *  # noqa: F401, F403
 from sqlalchemy.testing.suite.test_results import *  # noqa: F401, F403
 from sqlalchemy.testing.suite.test_select import *  # noqa: F401, F403
-from sqlalchemy.testing.suite.test_sequence import *  # noqa: F401, F403
+from sqlalchemy.testing.suite.test_sequence import (
+    SequenceTest as _SequenceTest,
+    HasSequenceTest as _HasSequenceTest,
+)  # noqa: F401, F403
 from sqlalchemy.testing.suite.test_update_delete import *  # noqa: F401, F403
 
 from sqlalchemy.testing.suite.test_cte import CTETest as _CTETest
@@ -2059,3 +2063,124 @@ def test_create_engine_wo_database(self):
         engine = create_engine(get_db_url().split("/database")[0])
         with engine.connect() as connection:
             assert connection.connection.database is None
+
+
+@pytest.mark.skipif(
+    bool(os.environ.get("SPANNER_EMULATOR_HOST")), reason="Skipped on emulator"
+)
+class SequenceTest(_SequenceTest):
+    @classmethod
+    def define_tables(cls, metadata):
+        Table(
+            "seq_pk",
+            metadata,
+            Column(
+                "id",
+                Integer,
+                sqlalchemy.Sequence("tab_id_seq"),
+                primary_key=True,
+            ),
+            Column("data", String(50)),
+        )
+
+        Table(
+            "seq_opt_pk",
+            metadata,
+            Column(
+                "id",
+                Integer,
+                sqlalchemy.Sequence("tab_id_seq_opt", data_type=Integer, optional=True),
+                primary_key=True,
+            ),
+            Column("data", String(50)),
+        )
+
+        Table(
+            "seq_no_returning",
+            metadata,
+            Column(
+                "id",
+                Integer,
+                sqlalchemy.Sequence("noret_id_seq"),
+                primary_key=True,
+            ),
+            Column("data", String(50)),
+            implicit_returning=False,
+        )
+
+    def test_insert_lastrowid(self, connection):
+        r = connection.execute(self.tables.seq_pk.insert(), dict(data="some data"))
+        assert len(r.inserted_primary_key) == 1
+        is_instance_of(r.inserted_primary_key[0], int)
+
+    def test_nextval_direct(self, connection):
+        r = connection.execute(self.tables.seq_pk.c.id.default)
+        is_instance_of(r, int)
+
+    def _assert_round_trip(self, table, conn):
+        row = conn.execute(table.select()).first()
+        id, name = row
+        is_instance_of(id, int)
+        eq_(name, "some data")
+
+    @testing.combinations((True,), (False,), argnames="implicit_returning")
+    @testing.requires.schemas
+    @pytest.mark.skip("Not supported by Cloud Spanner")
+    def test_insert_roundtrip_translate(self, connection, implicit_returning):
+        pass
+
+    @testing.requires.schemas
+    @pytest.mark.skip("Not supported by Cloud Spanner")
+    def test_nextval_direct_schema_translate(self, connection):
+        pass
+
+
+@pytest.mark.skipif(
+    bool(os.environ.get("SPANNER_EMULATOR_HOST")), reason="Skipped on emulator"
+)
+class HasSequenceTest(_HasSequenceTest):
+    @classmethod
+    def define_tables(cls, metadata):
+        sqlalchemy.Sequence("user_id_seq", metadata=metadata)
+        sqlalchemy.Sequence(
+            "other_seq", metadata=metadata, nomaxvalue=True, nominvalue=True
+        )
+        Table(
+            "user_id_table",
+            metadata,
+            Column("id", Integer, primary_key=True),
+        )
+
+    @pytest.mark.skip("Not supported by Cloud Spanner")
+    def test_has_sequence_cache(self, connection, metadata):
+        pass
+
+    @testing.requires.schemas
+    @pytest.mark.skip("Not supported by Cloud Spanner")
+    def test_has_sequence_schema(self, connection):
+        pass
+
+    @testing.requires.schemas
+    @pytest.mark.skip("Not supported by Cloud Spanner")
+    def test_has_sequence_schemas_neg(self, connection):
+        pass
+
+    @testing.requires.schemas
+    @pytest.mark.skip("Not supported by Cloud Spanner")
+    def test_has_sequence_default_not_in_remote(self, connection):
+        pass
+
+    @testing.requires.schemas
+    @pytest.mark.skip("Not supported by Cloud Spanner")
+    def test_has_sequence_remote_not_in_default(self, connection):
+        pass
+
+    @testing.requires.schemas
+    @pytest.mark.skip("Not supported by Cloud Spanner")
+    def test_get_sequence_names_no_sequence_schema(self, connection):
+        pass
+
+    @testing.requires.schemas
+    @pytest.mark.skip("Not supported by Cloud Spanner")
+    def test_get_sequence_names_sequences_schema(self, connection):
+        pass