Basic Support for Sequences

Author: rad-pat

databend_sqlalchemy/databend_dialect.py

@@ -31,7 +31,7 @@
 
 import sqlalchemy.engine.reflection
 import sqlalchemy.types as sqltypes
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Union, List
 from sqlalchemy import util as sa_util
 from sqlalchemy.engine import reflection
 from sqlalchemy.sql import (
@@ -703,7 +703,6 @@ def __init__(self, key_type, value_type):
         super(MAP, self).__init__()
 
 
-
 class DatabendDate(sqltypes.DATE):
     __visit_name__ = "DATE"
 
@@ -857,7 +856,6 @@ class DatabendGeography(GEOGRAPHY):
 }
 
 
-
 # Column spec
 colspecs = {
     sqltypes.Interval: DatabendInterval,
@@ -872,6 +870,12 @@ class DatabendGeography(GEOGRAPHY):
 class DatabendIdentifierPreparer(PGIdentifierPreparer):
     reserved_words = {r.lower() for r in RESERVED_WORDS}
 
+    # overridden to exclude schema from sequence
+    def format_sequence(
+        self, sequence, use_schema: bool = True
+    ) -> str:
+        return super().format_sequence(sequence, use_schema=False)
+
 
 class DatabendCompiler(PGCompiler):
     iscopyintotable: bool = False
@@ -1230,6 +1234,15 @@ def copy_into_table_results(self) -> list[dict]:
     def copy_into_location_results(self) -> dict:
         return self._copy_into_location_results
 
+    def fire_sequence(self, seq, type_):
+        return self._execute_scalar(
+            (
+                "select nextval(%s)"
+                % self.identifier_preparer.format_sequence(seq)
+            ),
+            type_,
+        )
+
 
 class DatabendTypeCompiler(compiler.GenericTypeCompiler):
     def visit_ARRAY(self, type_, **kw):
@@ -1280,7 +1293,6 @@ def visit_GEOGRAPHY(self, type_, **kw):
         return "GEOGRAPHY"
 
 
-
 class DatabendDDLCompiler(compiler.DDLCompiler):
     def visit_primary_key_constraint(self, constraint, **kw):
         return ""
@@ -1394,6 +1406,7 @@ class DatabendDialect(default.DefaultDialect):
     supports_empty_insert = False
     supports_is_distinct_from = True
     supports_multivalues_insert = True
+    supports_sequences = True
 
     supports_statement_cache = False
     supports_server_side_cursors = True
@@ -1621,7 +1634,6 @@ def get_temp_table_names(self, connection, schema=None, **kw):
         result = connection.execute(query, dict(schema_name=schema))
         return [row[0] for row in result]
 
-
     @reflection.cache
     def get_view_names(self, connection, schema=None, **kw):
         view_name_query = """
@@ -1762,7 +1774,6 @@ def get_multi_table_comment(
             schema='system',
         ).alias("a_tab_comments")
 
-
         has_filter_names, params = self._prepare_filter_names(filter_names)
         owner = schema or self.default_schema_name
 
@@ -1804,6 +1815,20 @@ def _check_unicode_description(self, connection):
         # We decode everything as UTF-8
         return True
 
+    @reflection.cache
+    def get_sequence_names(self, connection, schema: Optional[str] = None, **kw: Any) -> List[str]:
+        # N.B. sequences are not defined per schema/database
+        sequence_query = """
+            show sequences
+        """
+        query = text(sequence_query)
+        result = connection.execute(query)
+        return [row[0] for row in result]
+
+    def has_sequence(self, connection, sequence_name: str, schema: Optional[str] = None, **kw: Any) -> bool:
+        # N.B. sequences are not defined per schema/database
+        return sequence_name in self.get_sequence_names(connection, schema, **kw)
+
 
 dialect = DatabendDialect
 

tests/test_sqlalchemy.py

@@ -13,6 +13,7 @@
 from sqlalchemy.testing.suite import LongNameBlowoutTest as _LongNameBlowoutTest
 from sqlalchemy.testing.suite import QuotedNameArgumentTest as _QuotedNameArgumentTest
 from sqlalchemy.testing.suite import JoinTest as _JoinTest
+from sqlalchemy.testing.suite import HasSequenceTest as _HasSequenceTest
 
 from sqlalchemy.testing.suite import ServerSideCursorsTest as _ServerSideCursorsTest
 
@@ -21,7 +22,7 @@
 from sqlalchemy.testing.suite import IntegerTest as _IntegerTest
 
 from sqlalchemy import types as sql_types
-from sqlalchemy.testing import config
+from sqlalchemy.testing import config, skip_test
 from sqlalchemy import testing, Table, Column, Integer
 from sqlalchemy.testing import eq_, fixtures, assertions
 
@@ -586,9 +587,6 @@ def test_geometry_write_and_read(self, connection):
         eq_(result, ('{"type": "GeometryCollection", "geometries": [{"type": "Point", "coordinates": [10,20]},{"type": "LineString", "coordinates": [[10,20],[30,40]]},{"type": "Polygon", "coordinates": [[[10,20],[30,40],[50,60],[10,20]]]}]}'))
 
 
-
-
-
 class GeographyTest(fixtures.TablesTest):
 
     @classmethod
@@ -664,4 +662,74 @@ def test_geography_write_and_read(self, connection):
         result = connection.execute(
             select(geography_table.c.geography_data).where(geography_table.c.id == 7)
         ).scalar()
-        eq_(result, ('{"type": "GeometryCollection", "geometries": [{"type": "Point", "coordinates": [10,20]},{"type": "LineString", "coordinates": [[10,20],[30,40]]},{"type": "Polygon", "coordinates": [[[10,20],[30,40],[50,60],[10,20]]]}]}'))
+        eq_(result, ('{"type": "GeometryCollection", "geometries": [{"type": "Point", "coordinates": [10,20]},{"type": "LineString", "coordinates": [[10,20],[30,40]]},{"type": "Polygon", "coordinates": [[[10,20],[30,40],[50,60],[10,20]]]}]}'))
+
+
+class HasSequenceTest(_HasSequenceTest):
+
+    # ToDo - overridden other_seq definition due to lack of sequence ddl support for nominvalue nomaxvalue
+    @classmethod
+    def define_tables(cls, metadata):
+        normalize_sequence(config, Sequence("user_id_seq", metadata=metadata))
+        normalize_sequence(
+            config,
+            Sequence(
+                "other_seq",
+                metadata=metadata,
+                # nomaxvalue=True,
+                # nominvalue=True,
+            ),
+        )
+        if testing.requires.schemas.enabled:
+            #ToDo - omitted because Databend does not allow schema on sequence
+            # normalize_sequence(
+            #     config,
+            #     Sequence(
+            #         "user_id_seq", schema=config.test_schema, metadata=metadata
+            #     ),
+            # )
+            normalize_sequence(
+                config,
+                Sequence(
+                    "schema_seq", schema=config.test_schema, metadata=metadata
+                ),
+            )
+        Table(
+            "user_id_table",
+            metadata,
+            Column("id", Integer, primary_key=True),
+        )
+
+    @testing.skip("databend")  # ToDo - requires definition of sequences with schema
+    def test_has_sequence_remote_not_in_default(self, connection):
+        eq_(inspect(connection).has_sequence("schema_seq"), False)
+
+    @testing.skip("databend")  # ToDo - requires definition of sequences with schema
+    def test_get_sequence_names(self, connection):
+        exp = {"other_seq", "user_id_seq"}
+
+        res = set(inspect(connection).get_sequence_names())
+        is_true(res.intersection(exp) == exp)
+        is_true("schema_seq" not in res)
+
+    @testing.skip("databend")  # ToDo - requires definition of sequences with schema
+    @testing.requires.schemas
+    def test_get_sequence_names_no_sequence_schema(self, connection):
+        eq_(
+            inspect(connection).get_sequence_names(
+                schema=config.test_schema_2
+            ),
+            [],
+        )
+
+    @testing.skip("databend")  # ToDo - requires definition of sequences with schema
+    @testing.requires.schemas
+    def test_get_sequence_names_sequences_schema(self, connection):
+        eq_(
+            sorted(
+                inspect(connection).get_sequence_names(
+                    schema=config.test_schema
+                )
+            ),
+            ["schema_seq", "user_id_seq"],
+        )