From f25006cb823f9f145caaa3b3e1342c3f8c864f0e Mon Sep 17 00:00:00 2001
From: Oleg Ovcharuk <vgvoleg@gmail.com>
Date: Thu, 3 Jul 2025 10:37:40 +0300
Subject: [PATCH] Fix Date literal processor for YDB

---
 ydb_sqlalchemy/sqlalchemy/__init__.py        |  1 +
 ydb_sqlalchemy/sqlalchemy/datetime_types.py  | 10 ++++++++++
 ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py | 10 ++++++++++
 ydb_sqlalchemy/sqlalchemy/types.py           |  2 +-
 4 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py
index 7bdf8f8..0f271f3 100644
--- a/ydb_sqlalchemy/sqlalchemy/__init__.py
+++ b/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,
diff --git a/ydb_sqlalchemy/sqlalchemy/datetime_types.py b/ydb_sqlalchemy/sqlalchemy/datetime_types.py
index 6cd10cb..7337af6 100644
--- a/ydb_sqlalchemy/sqlalchemy/datetime_types.py
+++ b/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]:
diff --git a/ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py b/ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py
index de8b446..6908193 100644
--- a/ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py
+++ b/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')"
diff --git a/ydb_sqlalchemy/sqlalchemy/types.py b/ydb_sqlalchemy/sqlalchemy/types.py
index 34e26b6..261eb9f 100644
--- a/ydb_sqlalchemy/sqlalchemy/types.py
+++ b/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