Fix Date literal processor for YDB

Author: vgvoleg

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -136,6 +136,7 @@ class YqlDialect(StrCompileDialect):
     colspecs = {
         sa.types.JSON: types.YqlJSON,
         sa.types.JSON.JSONPathType: types.YqlJSON.YqlJSONPathType,
+        sa.types.Date: types.YqlDate,
         sa.types.DateTime: types.YqlTimestamp,  # Because YDB's DateTime doesn't store microseconds
         sa.types.DATETIME: types.YqlDateTime,
         sa.types.TIMESTAMP: types.YqlTimestamp,

ydb_sqlalchemy/sqlalchemy/datetime_types.py

@@ -4,6 +4,16 @@
 from sqlalchemy import types as sqltypes
 
 
+class YqlDate(sqltypes.Date):
+    def literal_processor(self, dialect):
+        parent = super().literal_processor(dialect)
+
+        def process(value):
+            return f"Date({parent(value)})"
+
+        return process
+
+
 class YqlTimestamp(sqltypes.TIMESTAMP):
     def result_processor(self, dialect, coltype):
         def process(value: Optional[datetime.datetime]) -> Optional[datetime.datetime]:

ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py

@@ -1,3 +1,4 @@
+from datetime import date
 import sqlalchemy as sa
 
 from . import YqlDialect, types
@@ -25,3 +26,12 @@ def test_casts():
         "CAST(1/2 AS UInt8)",
         "String::JoinFromList(ListMap(TOPFREQ(1/2, 5), ($x) -> { RETURN CAST($x AS UTF8) ;}), ', ')",
     ]
+
+
+def test_ydb_types():
+    dialect = YqlDialect()
+
+    query = sa.literal(date(1996, 11, 19))
+    compiled = query.compile(dialect=dialect, compile_kwargs={"literal_binds": True})
+
+    assert str(compiled) == "Date('1996-11-19')"

ydb_sqlalchemy/sqlalchemy/types.py

@@ -10,7 +10,7 @@
 from sqlalchemy import ARRAY, exc, types
 from sqlalchemy.sql import type_api
 
-from .datetime_types import YqlDateTime, YqlTimestamp  # noqa: F401
+from .datetime_types import YqlDate, YqlDateTime, YqlTimestamp  # noqa: F401
 from .json import YqlJSON  # noqa: F401