Add optional type casting to limit-offset clause

Author: LuckySting

test/test_suite.py

@@ -6,13 +6,12 @@
 
 from sqlalchemy.testing import is_true, is_false
 from sqlalchemy.testing.suite import eq_, testing, inspect, provide_metadata, config, requirements, fixtures
-from sqlalchemy.testing.suite import func, column, literal_column, select, exists
+from sqlalchemy.testing.suite import func, column, literal_column, select, exists, union
 from sqlalchemy.testing.suite import MetaData, Column, Table, Integer, String
 
 from sqlalchemy.testing.suite.test_select import (
     ExistsTest as _ExistsTest,
     LikeFunctionsTest as _LikeFunctionsTest,
-    CompoundSelectTest as _CompoundSelectTest,
 )
 from sqlalchemy.testing.suite.test_reflection import (
     HasTableTest as _HasTableTest,
@@ -49,7 +48,6 @@
 from sqlalchemy.testing.suite.test_insert import InsertBehaviorTest as _InsertBehaviorTest
 from sqlalchemy.testing.suite.test_ddl import LongNameBlowoutTest as _LongNameBlowoutTest
 from sqlalchemy.testing.suite.test_results import RowFetchTest as _RowFetchTest
-from sqlalchemy.testing.suite.test_deprecations import DeprecatedCompoundSelectTest as _DeprecatedCompoundSelectTest
 
 from ydb_sqlalchemy.sqlalchemy import types as ydb_sa_types
 
@@ -294,20 +292,6 @@ def test_not_regexp_match(self):
         self._test(~col.regexp_match("a.cde"), {2, 3, 4, 7, 8, 10, 11})
 
 
-class CompoundSelectTest(_CompoundSelectTest):
-    @pytest.mark.skip("limit don't work")
-    def test_distinct_selectable_in_unions(self):
-        pass
-
-    @pytest.mark.skip("limit don't work")
-    def test_limit_offset_in_unions_from_alias(self):
-        pass
-
-    @pytest.mark.skip("limit don't work")
-    def test_limit_offset_aliased_selectable_in_unions(self):
-        pass
-
-
 class EscapingTest(_EscapingTest):
     @provide_metadata
     def test_percent_sign_round_trip(self):
@@ -364,45 +348,23 @@ def test_group_by_composed(self):
 
 
 class FetchLimitOffsetTest(_FetchLimitOffsetTest):
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
-    def test_bound_limit(self, connection):
-        pass
-
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
-    def test_bound_limit_offset(self, connection):
-        pass
-
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
-    def test_bound_offset(self, connection):
-        pass
-
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
-    def test_expr_limit_simple_offset(self, connection):
-        pass
-
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
     def test_limit_render_multiple_times(self, connection):
-        pass
-
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
-    def test_simple_limit(self, connection):
-        pass
-
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
-    def test_simple_limit_offset(self, connection):
-        pass
-
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
-    def test_simple_offset(self, connection):
-        pass
+        """
+        YQL does not support scalar subquery, so test was refiled with simple subquery
+        """
+        table = self.tables.some_table
+        stmt = select(table.c.id).limit(1).subquery()
 
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
-    def test_simple_offset_zero(self, connection):
-        pass
+        u = union(select(stmt), select(stmt)).subquery().select()
 
-    @pytest.mark.skip("Failed to convert type: Int64 to Uint64")
-    def test_simple_limit_expr_offset(self, connection):
-        pass
+        self._assert_result(
+            connection,
+            u,
+            [
+                (1,),
+                (1,),
+            ],
+        )
 
 
 class InsertBehaviorTest(_InsertBehaviorTest):
@@ -539,8 +501,3 @@ class RowFetchTest(_RowFetchTest):
     @pytest.mark.skip("scalar subquery unsupported")
     def test_row_w_scalar_select(self, connection):
         pass
-
-
-@pytest.mark.skip("TODO: try it after limit/offset tests would fixed")
-class DeprecatedCompoundSelectTest(_DeprecatedCompoundSelectTest):
-    pass

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -24,7 +24,7 @@
 from sqlalchemy.engine.default import StrCompileDialect, DefaultExecutionContext
 from sqlalchemy.util.compat import inspect_getfullargspec
 
-from typing import Any, Union, Mapping, Sequence, Optional, Tuple, List, Dict
+from typing import Any, Union, Mapping, Sequence, Optional, Tuple, List, Dict, Type
 
 from . import types
 
@@ -87,15 +87,30 @@ def visit_FLOAT(self, type_: sa.FLOAT, **kw):
     def visit_BOOLEAN(self, type_: sa.BOOLEAN, **kw):
         return "BOOL"
 
+    def visit_uint64(self, type_: types.UInt64, **kw):
+        return "UInt64"
+
     def visit_uint32(self, type_: types.UInt32, **kw):
         return "UInt32"
 
-    def visit_uint64(self, type_: types.UInt64, **kw):
-        return "UInt64"
+    def visit_uint16(self, type_: types.UInt16, **kw):
+        return "UInt16"
 
     def visit_uint8(self, type_: types.UInt8, **kw):
         return "UInt8"
 
+    def visit_int64(self, type_: types.UInt64, **kw):
+        return "Int64"
+
+    def visit_int32(self, type_: types.UInt32, **kw):
+        return "Int32"
+
+    def visit_int16(self, type_: types.UInt16, **kw):
+        return "Int16"
+
+    def visit_int8(self, type_: types.UInt8, **kw):
+        return "Int8"
+
     def visit_INTEGER(self, type_: sa.INTEGER, **kw):
         return "Int64"
 
@@ -134,8 +149,28 @@ def get_ydb_type(
 
         if isinstance(type_, (sa.Text, sa.String, sa.Uuid)):
             ydb_type = ydb.PrimitiveType.Utf8
+
+        # Integers
+        elif isinstance(type_, types.UInt64):
+            ydb_type = ydb.PrimitiveType.Uint64
+        elif isinstance(type_, types.UInt32):
+            ydb_type = ydb.PrimitiveType.Uint32
+        elif isinstance(type_, types.UInt16):
+            ydb_type = ydb.PrimitiveType.Uint16
+        elif isinstance(type_, types.UInt8):
+            ydb_type = ydb.PrimitiveType.Uint8
+        elif isinstance(type_, types.Int64):
+            ydb_type = ydb.PrimitiveType.Int64
+        elif isinstance(type_, types.Int32):
+            ydb_type = ydb.PrimitiveType.Int32
+        elif isinstance(type_, types.Int16):
+            ydb_type = ydb.PrimitiveType.Int16
+        elif isinstance(type_, types.Int8):
+            ydb_type = ydb.PrimitiveType.Int8
         elif isinstance(type_, sa.Integer):
             ydb_type = ydb.PrimitiveType.Int64
+        # Integers
+
         elif isinstance(type_, sa.JSON):
             ydb_type = ydb.PrimitiveType.Json
         elif isinstance(type_, sa.DateTime):
@@ -188,6 +223,32 @@ def group_by_clause(self, select, **kw):
         kw.update(within_columns_clause=True)
         return super(YqlCompiler, self).group_by_clause(select, **kw)
 
+    def limit_clause(self, select, **kw):
+        text = ""
+        if select._limit_clause is not None:
+            limit_clause = self._maybe_cast(
+                select._limit_clause, types.UInt64, skip_types=(types.UInt64, types.UInt32, types.UInt16, types.UInt8)
+            )
+            text += "\n LIMIT " + self.process(limit_clause, **kw)
+        if select._offset_clause is not None:
+            offset_clause = self._maybe_cast(
+                select._offset_clause, types.UInt64, skip_types=(types.UInt64, types.UInt32, types.UInt16, types.UInt8)
+            )
+            if select._limit_clause is None:
+                text += "\n LIMIT NULL"
+            text += " OFFSET " + self.process(offset_clause, **kw)
+        return text
+
+    def _maybe_cast(
+        self,
+        element: Any,
+        cast_to: Type[sa.types.TypeEngine],
+        skip_types: Optional[Tuple[Type[sa.types.TypeEngine], ...]] = None,
+    ) -> Any:
+        if not hasattr(element, "type") or not isinstance(element.type, skip_types):
+            return sa.Cast(element, cast_to)
+        return element
+
     def render_literal_value(self, value, type_):
         if isinstance(value, str):
             value = "".join(STR_QUOTE_MAP.get(x, x) for x in value)
@@ -277,16 +338,14 @@ def _is_bound_to_nullable_column(self, bind_name: str) -> bool:
     def _guess_bound_variable_type_by_parameters(
         self, bind: sa.BindParameter, post_compile_bind_values: list
     ) -> Optional[sa.types.TypeEngine]:
-        if not bind.expanding:
-            if isinstance(bind.type, sa.types.NullType):
-                return None
-            bind_type = bind.type
-        else:
+        bind_type = bind.type
+        if bind.expanding or (isinstance(bind.type, sa.types.NullType) and post_compile_bind_values):
             not_null_values = [v for v in post_compile_bind_values if v is not None]
             if not_null_values:
                 bind_type = sa.BindParameter("", not_null_values[0]).type
-            else:
-                return None
+
+        if isinstance(bind_type, sa.types.NullType):
+            return None
 
         return bind_type
 

ydb_sqlalchemy/sqlalchemy/types.py

@@ -3,18 +3,38 @@
 from typing import Mapping, Any, Union, Type
 
 
+class UInt64(types.Integer):
+    __visit_name__ = "uint64"
+
+
 class UInt32(types.Integer):
     __visit_name__ = "uint32"
 
 
-class UInt64(types.Integer):
-    __visit_name__ = "uint64"
+class UInt16(types.Integer):
+    __visit_name__ = "uint32"
 
 
 class UInt8(types.Integer):
     __visit_name__ = "uint8"
 
 
+class Int64(types.Integer):
+    __visit_name__ = "int64"
+
+
+class Int32(types.Integer):
+    __visit_name__ = "int32"
+
+
+class Int16(types.Integer):
+    __visit_name__ = "int32"
+
+
+class Int8(types.Integer):
+    __visit_name__ = "int8"
+
+
 class ListType(ARRAY):
     __visit_name__ = "list_type"