Fix type hints and tests behaviour

Author: tretyak-rd

test/test_core.py

@@ -23,7 +23,8 @@ def test_sa_text(self, connection):
                 DECLARE :data AS List<Struct<x:Int64, y:Int64>>;
                 SELECT x, y FROM AS_TABLE(:data)
                 """
-            ), [{"data": [{"x": 2, "y": 1}, {"x": 3, "y": 2}]}]
+            ),
+            [{"data": [{"x": 2, "y": 1}, {"x": 3, "y": 2}]}],
         )
         assert set(rs.fetchall()) == {(2, 1), (3, 2)}
 

test/test_suite.py

@@ -5,7 +5,7 @@
 from sqlalchemy.testing.suite import *  # noqa: F401, F403
 
 from sqlalchemy.testing import is_true, is_false
-from sqlalchemy.testing.suite import eq_, testing, inspect, provide_metadata, config, requirements
+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 MetaData, Column, Table, Integer, String
 
@@ -21,7 +21,6 @@
     CompositeKeyReflectionTest as _CompositeKeyReflectionTest,
     ComponentReflectionTestExtra as _ComponentReflectionTestExtra,
     QuotedNameArgumentTest as _QuotedNameArgumentTest,
-    BizarroCharacterFKResolutionTest as _BizarroCharacterFKResolutionTest,
 )
 from sqlalchemy.testing.suite.test_types import (
     IntegerTest as _IntegerTest,
@@ -33,6 +32,10 @@
     NativeUUIDTest as _NativeUUIDTest,
     TimeMicrosecondsTest as _TimeMicrosecondsTest,
     DateTimeCoercedToDateTimeTest as _DateTimeCoercedToDateTimeTest,
+    DateTest as _DateTest,
+    DateTimeMicrosecondsTest as _DateTimeMicrosecondsTest,
+    DateTimeTest as _DateTimeTest,
+    TimestampMicrosecondsTest as _TimestampMicrosecondsTest,
 )
 from sqlalchemy.testing.suite.test_dialect import (
     EscapingTest as _EscapingTest,
@@ -64,11 +67,6 @@ def column_getter(*args, **kwargs):
 test_types_suite.Column = column_getter
 
 
-@pytest.mark.skip("foreign keys unsupported")
-class BizarroCharacterFKResolutionTest(_BizarroCharacterFKResolutionTest):
-    pass
-
-
 class ComponentReflectionTest(_ComponentReflectionTest):
     def _check_list(self, result, exp, req_keys=None, msg=None):
         try:
@@ -121,7 +119,6 @@ def define_reflected_tables(cls, metadata, schema):
             Column("id", sa.Integer, primary_key=True, comment="id comment"),
             Column("data", sa.String(20), comment="data % comment"),
             Column("d2", sa.String(20), comment=r"""Comment types type speedily ' " \ '' Fun!"""),
-            Column("d3", sa.String(42), comment="Comment\nwith\rescapes"),
             schema=schema,
             comment=r"""the test % ' " \ table comment""",
         )
@@ -422,6 +419,22 @@ def test_no_results_for_non_returning_insert(self, connection):
         pass
 
 
+class DateTest(_DateTest):
+    run_dispose_bind = "once"
+
+
+class DateTimeMicrosecondsTest(_DateTimeMicrosecondsTest):
+    run_dispose_bind = "once"
+
+
+class DateTimeTest(_DateTimeTest):
+    run_dispose_bind = "once"
+
+
+class TimestampMicrosecondsTest(_TimestampMicrosecondsTest):
+    run_dispose_bind = "once"
+
+
 @pytest.mark.skip("unsupported Time data type")
 class TimeTest(_TimeTest):
     pass

test_dbapi/test_dbapi.py

@@ -12,21 +12,21 @@ def test_connection(connection):
 
     cur = connection.cursor()
     with suppress(dbapi.DatabaseError):
-        cur.execute(dbapi.YdbOperation("DROP TABLE foo", is_ddl=True))
+        cur.execute(dbapi.YdbQuery("DROP TABLE foo", is_ddl=True))
 
     assert not connection.check_exists("/local/foo")
     with pytest.raises(dbapi.ProgrammingError):
         connection.describe("/local/foo")
 
-    cur.execute(dbapi.YdbOperation("CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))", is_ddl=True))
+    cur.execute(dbapi.YdbQuery("CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))", is_ddl=True))
 
     assert connection.check_exists("/local/foo")
 
     col = connection.describe("/local/foo").columns[0]
     assert col.name == "id"
     assert col.type == ydb.PrimitiveType.Int64
 
-    cur.execute(dbapi.YdbOperation("DROP TABLE foo", is_ddl=True))
+    cur.execute(dbapi.YdbQuery("DROP TABLE foo", is_ddl=True))
     cur.close()
 
 
@@ -35,18 +35,17 @@ def test_cursor_raw_query(connection):
     assert cur
 
     with suppress(dbapi.DatabaseError):
-        cur.execute(dbapi.YdbOperation("DROP TABLE test", is_ddl=True))
+        cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
 
-    cur.execute(dbapi.YdbOperation("CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))", is_ddl=True))
+    cur.execute(dbapi.YdbQuery("CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))", is_ddl=True))
 
     cur.execute(
-        dbapi.YdbOperation(
+        dbapi.YdbQuery(
             """
             DECLARE $data AS List<Struct<id:Int64, text: Utf8>>;
 
             INSERT INTO test SELECT id, text FROM AS_TABLE($data);
             """,
-            is_ddl=False,
             parameters_types={
                 "$data": ydb.ListType(
                     ydb.StructType()
@@ -63,7 +62,7 @@ def test_cursor_raw_query(connection):
         },
     )
 
-    cur.execute(dbapi.YdbOperation("DROP TABLE test", is_ddl=True))
+    cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
 
     cur.close()
 
@@ -75,25 +74,25 @@ def test_errors(connection):
     cur = connection.cursor()
 
     with suppress(dbapi.DatabaseError):
-        cur.execute(dbapi.YdbOperation("DROP TABLE test", is_ddl=True))
+        cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
 
     with pytest.raises(dbapi.DataError):
-        cur.execute(dbapi.YdbOperation("SELECT 18446744073709551616", is_ddl=False))
+        cur.execute(dbapi.YdbQuery("SELECT 18446744073709551616"))
 
     with pytest.raises(dbapi.DataError):
-        cur.execute(dbapi.YdbOperation("SELECT * FROM 拉屎", is_ddl=False))
+        cur.execute(dbapi.YdbQuery("SELECT * FROM 拉屎"))
 
     with pytest.raises(dbapi.DataError):
-        cur.execute(dbapi.YdbOperation("SELECT floor(5 / 2)", is_ddl=False))
+        cur.execute(dbapi.YdbQuery("SELECT floor(5 / 2)"))
 
     with pytest.raises(dbapi.ProgrammingError):
-        cur.execute(dbapi.YdbOperation("SELECT * FROM test", is_ddl=False))
+        cur.execute(dbapi.YdbQuery("SELECT * FROM test"))
 
-    cur.execute(dbapi.YdbOperation("CREATE TABLE test(id Int64, PRIMARY KEY (id))", is_ddl=True))
+    cur.execute(dbapi.YdbQuery("CREATE TABLE test(id Int64, PRIMARY KEY (id))", is_ddl=True))
 
-    cur.execute(dbapi.YdbOperation("INSERT INTO test(id) VALUES(1)", is_ddl=False))
+    cur.execute(dbapi.YdbQuery("INSERT INTO test(id) VALUES(1)"))
     with pytest.raises(dbapi.IntegrityError):
-        cur.execute(dbapi.YdbOperation("INSERT INTO test(id) VALUES(1)", is_ddl=False))
+        cur.execute(dbapi.YdbQuery("INSERT INTO test(id) VALUES(1)"))
 
-    cur.execute(dbapi.YdbOperation("DROP TABLE test", is_ddl=True))
+    cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
     cur.close()

ydb_sqlalchemy/dbapi/__init__.py

@@ -1,5 +1,5 @@
 from .connection import Connection
-from .cursor import YdbOperation, Cursor
+from .cursor import Cursor, YdbQuery  # noqa: F401
 from .errors import (
     Warning,
     Error,

ydb_sqlalchemy/dbapi/cursor.py

@@ -2,7 +2,7 @@
 import itertools
 import logging
 
-from typing import Any, Mapping, Optional, Sequence, Union
+from typing import Any, Mapping, Optional, Sequence, Union, Dict
 
 import ydb
 from .errors import (
@@ -24,10 +24,12 @@ def get_column_type(type_obj: Any) -> str:
 
 
 @dataclasses.dataclass
-class YdbOperation:
+class YdbQuery:
     yql_text: str
-    is_ddl: bool
-    parameters_types: dict[str, Union[ydb.PrimitiveType, ydb.AbstractTypeBuilder]] = dataclasses.field(default_factory=dict)
+    parameters_types: Dict[str, Union[ydb.PrimitiveType, ydb.AbstractTypeBuilder]] = dataclasses.field(
+        default_factory=dict
+    )
+    is_ddl: bool = False
 
 
 class Cursor(object):
@@ -38,7 +40,7 @@ def __init__(self, connection):
         self.rows = None
         self._rows_prefetched = None
 
-    def execute(self, operation: YdbOperation, parameters: Optional[Mapping[str, Any]] = None):
+    def execute(self, operation: YdbQuery, parameters: Optional[Mapping[str, Any]] = None):
         self.description = None
 
         if operation.is_ddl or not operation.parameters_types:
@@ -129,7 +131,7 @@ def _ensure_prefetched(self):
             self.rows = iter(self._rows_prefetched)
         return self._rows_prefetched
 
-    def executemany(self, operation: YdbOperation, seq_of_parameters: Optional[Sequence[Mapping[str, Any]]]):
+    def executemany(self, operation: YdbQuery, seq_of_parameters: Optional[Sequence[Mapping[str, Any]]]):
         for parameters in seq_of_parameters:
             self.execute(operation, parameters)
 

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -22,9 +22,9 @@
 from sqlalchemy.engine.default import StrCompileDialect, DefaultExecutionContext
 from sqlalchemy.util.compat import inspect_getfullargspec
 
-from typing import Any, Union, Mapping, Sequence, Optional, Tuple
+from typing import Any, Union, Mapping, Sequence, Optional, Tuple, List, Dict
 
-from .types import UInt32, UInt64
+from . import types
 
 STR_QUOTE_MAP = {
     "'": "\\'",
@@ -136,7 +136,7 @@ def get_ydb_type(
             ydb_type = ydb.PrimitiveType.Int64
         elif isinstance(type_, sa.JSON):
             ydb_type = ydb.PrimitiveType.Json
-        elif isinstance(type_, (sa.DateTime, sa.TIMESTAMP)):
+        elif isinstance(type_, sa.DateTime):
             ydb_type = ydb.PrimitiveType.Timestamp
         elif isinstance(type_, sa.Date):
             ydb_type = ydb.PrimitiveType.Date
@@ -288,7 +288,7 @@ def _guess_bound_variable_type_by_parameters(
 
         return bind_type
 
-    def _get_expanding_bind_names(self, bind_name: str, parameters_values: Mapping[str, list[Any]]) -> list[Any]:
+    def _get_expanding_bind_names(self, bind_name: str, parameters_values: Mapping[str, List[Any]]) -> List[Any]:
         expanding_bind_names = []
         for parameter_name in parameters_values:
             parameter_bind_name = "_".join(parameter_name.split("_")[:-1])
@@ -298,7 +298,7 @@ def _get_expanding_bind_names(self, bind_name: str, parameters_values: Mapping[s
 
     def get_bind_types(
         self, post_compile_parameters: Optional[Union[Sequence[Mapping[str, Any]], Mapping[str, Any]]]
-    ) -> dict[str, Union[ydb.PrimitiveType, ydb.AbstractTypeBuilder]]:
+    ) -> Dict[str, Union[ydb.PrimitiveType, ydb.AbstractTypeBuilder]]:
         """
         This method extracts information about bound variables from the table definition and parameters.
         """
@@ -357,8 +357,8 @@ def upsert(table):
     ydb.PrimitiveType.Int64: sa.INTEGER,
     ydb.PrimitiveType.Uint8: sa.INTEGER,
     ydb.PrimitiveType.Uint16: sa.INTEGER,
-    ydb.PrimitiveType.Uint32: UInt32,
-    ydb.PrimitiveType.Uint64: UInt64,
+    ydb.PrimitiveType.Uint32: types.UInt32,
+    ydb.PrimitiveType.Uint64: types.UInt64,
     ydb.PrimitiveType.Float: sa.FLOAT,
     ydb.PrimitiveType.Double: sa.FLOAT,
     ydb.PrimitiveType.String: sa.BINARY,
@@ -528,17 +528,17 @@ def _make_ydb_operation(
         context: Optional[DefaultExecutionContext] = None,
         parameters: Optional[Union[Sequence[Mapping[str, Any]], Mapping[str, Any]]] = None,
         execute_many: bool = False,
-    ) -> Tuple[dbapi.YdbOperation, Optional[Union[Sequence[Mapping[str, Any]], Mapping[str, Any]]]]:
+    ) -> Tuple[dbapi.YdbQuery, Optional[Union[Sequence[Mapping[str, Any]], Mapping[str, Any]]]]:
         is_ddl = context.isddl if context is not None else False
 
         if not is_ddl and parameters:
             parameters_types = context.compiled.get_bind_types(parameters)
             parameters_types = {f"${k}": v for k, v in parameters_types.items()}
             statement, parameters = self._format_variables(statement, parameters, execute_many)
-            return dbapi.YdbOperation(is_ddl=is_ddl, yql_text=statement, parameters_types=parameters_types), parameters
+            return dbapi.YdbQuery(yql_text=statement, parameters_types=parameters_types, is_ddl=is_ddl), parameters
 
         statement, parameters = self._format_variables(statement, parameters, execute_many)
-        return dbapi.YdbOperation(is_ddl=is_ddl, yql_text=statement), parameters
+        return dbapi.YdbQuery(yql_text=statement, is_ddl=is_ddl), parameters
 
     def do_ping(self, dbapi_connection: dbapi.Connection) -> bool:
         cursor = dbapi_connection.cursor()