ydb-sqlalchemy 0.1.15__py2.py3-none-any.whl → 0.1.17__py2.py3-none-any.whl

test/test_core.py

@@ -181,12 +181,19 @@ class TestSimpleSelect(TablesTest):
         rows = connection.execute(stm).fetchall()
         assert set(rows) == {(1,), (2,), (3,), (4,), (6,), (7,)}
 
+        # LIMIT
+        rows = connection.execute(tb.select().order_by(tb.c.id).limit(2)).fetchall()
+        assert rows == [
+            (1, "some text", Decimal("3.141592653")),
+            (2, "test text", Decimal("3.14159265")),
+        ]
+
         # LIMIT/OFFSET
-        # rows = connection.execute(tb.select().order_by(tb.c.id).limit(2)).fetchall()
-        # assert rows == [
-        #     (1, "some text", Decimal("3.141592653")),
-        #     (2, "test text", Decimal("3.14159265")),
-        # ]
+        rows = connection.execute(tb.select().order_by(tb.c.id).limit(2).offset(1)).fetchall()
+        assert rows == [
+            (2, "test text", Decimal("3.14159265")),
+            (3, "test test", Decimal("3.1415926")),
+        ]
 
         # ORDER BY ASC
         rows = connection.execute(sa.select(tb.c.id).order_by(tb.c.id)).fetchall()
@@ -223,7 +230,7 @@ class TestTypes(TablesTest):
             "test_primitive_types",
             metadata,
             Column("int", sa.Integer, primary_key=True),
-            # Column("bin", sa.BINARY),
+            Column("bin", sa.BINARY),
             Column("str", sa.String),
             Column("float", sa.Float),
             Column("bool", sa.Boolean),
@@ -253,7 +260,7 @@ class TestTypes(TablesTest):
 
         statement = sa.insert(table).values(
             int=42,
-            # bin=b"abc",
+            bin=b"abc",
             str="Hello World!",
             float=3.5,
             bool=True,
@@ -262,7 +269,7 @@ class TestTypes(TablesTest):
         connection.execute(statement)
 
         row = connection.execute(sa.select(table)).fetchone()
-        assert row == (42, "Hello World!", 3.5, True)
+        assert row == (42, b"abc", "Hello World!", 3.5, True)
 
     def test_all_binary_types(self, connection):
         table = self.tables.test_all_binary_types
@@ -1150,8 +1157,15 @@ class TestAsTable(TablesTest):
             Column("val_int", Integer, nullable=True),
             Column("val_str", String, nullable=True),
         )
+        Table(
+            "test_as_table_json",
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("data", sa.JSON, nullable=True),
+        )
 
-    def test_upsert_as_table(self, connection):
+    @pytest.mark.parametrize("list_cls", [types.ListType, sa.ARRAY])
+    def test_upsert_as_table(self, connection, list_cls):
         table = self.tables.test_as_table
 
         input_data = [
@@ -1167,7 +1181,7 @@ class TestAsTable(TablesTest):
                 "val_str": types.Optional(String),
             }
         )
-        list_type = types.ListType(struct_type)
+        list_type = list_cls(struct_type)
 
         bind_param = sa.bindparam("data", type_=list_type)
 
@@ -1187,7 +1201,39 @@ class TestAsTable(TablesTest):
             (3, 30, None),
         ]
 
-    def test_insert_as_table(self, connection):
+    @pytest.mark.parametrize("list_cls", [types.ListType, sa.ARRAY])
+    def test_upsert_from_table_json(self, connection, list_cls):
+        table = self.tables.test_as_table_json
+
+        input_data = [
+            {"id": 1, "data": {"a": 1}},
+            {"id": 2, "data": [1, 2, 3]},
+            {"id": 3, "data": None},
+        ]
+
+        struct_type = types.StructType.from_table(table)
+        list_type = list_cls(struct_type)
+
+        bind_param = sa.bindparam("input_data", type_=list_type)
+
+        cols = [sa.column(c.name, type_=c.type) for c in table.columns]
+        upsert_stm = ydb_sa.upsert(table).from_select(
+            [c.name for c in table.columns],
+            sa.select(*cols).select_from(sa.func.AS_TABLE(bind_param)),
+        )
+
+        connection.execute(upsert_stm, {"input_data": input_data})
+
+        rows = connection.execute(sa.select(table).order_by(table.c.id)).fetchall()
+
+        assert rows == [
+            (1, {"a": 1}),
+            (2, [1, 2, 3]),
+            (3, None),
+        ]
+
+    @pytest.mark.parametrize("list_cls", [types.ListType, sa.ARRAY])
+    def test_insert_as_table(self, connection, list_cls):
         table = self.tables.test_as_table
 
         input_data = [
@@ -1202,7 +1248,7 @@ class TestAsTable(TablesTest):
                 "val_str": types.Optional(String),
             }
         )
-        list_type = types.ListType(struct_type)
+        list_type = list_cls(struct_type)
 
         bind_param = sa.bindparam("data", type_=list_type)
 
@@ -1221,7 +1267,8 @@ class TestAsTable(TablesTest):
             (5, None, "e"),
         ]
 
-    def test_upsert_from_table_reflection(self, connection):
+    @pytest.mark.parametrize("list_cls", [types.ListType, sa.ARRAY])
+    def test_upsert_from_table_reflection(self, connection, list_cls):
         table = self.tables.test_as_table
 
         input_data = [
@@ -1230,7 +1277,7 @@ class TestAsTable(TablesTest):
         ]
 
         struct_type = types.StructType.from_table(table)
-        list_type = types.ListType(struct_type)
+        list_type = list_cls(struct_type)
 
         bind_param = sa.bindparam("data", type_=list_type)
 

ydb_sqlalchemy/_version.py

@@ -1 +1 @@
-VERSION = "0.1.15"
+VERSION = "0.1.17"

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -160,6 +160,7 @@ class YqlDialect(StrCompileDialect):
         sa.types.BINARY: types.Binary,
         sa.types.LargeBinary: types.Binary,
         sa.types.BLOB: types.Binary,
+        sa.types.ARRAY: types.ListType,
     }
 
     connection_characteristics = util.immutabledict(
@@ -205,6 +206,7 @@ class YqlDialect(StrCompileDialect):
         json_serializer=None,
         json_deserializer=None,
         _add_declare_for_yql_stmt_vars=False,
+        _statement_prefixes_list=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -214,6 +216,7 @@ class YqlDialect(StrCompileDialect):
         # NOTE: _add_declare_for_yql_stmt_vars is temporary and is soon to be removed.
         # no need in declare in yql statement here since ydb 24-1
         self._add_declare_for_yql_stmt_vars = _add_declare_for_yql_stmt_vars
+        self._statement_prefixes = tuple(_statement_prefixes_list) if _statement_prefixes_list else ()
 
     def _describe_table(self, connection, table_name, schema=None) -> ydb.TableDescription:
         if schema is not None:
@@ -404,6 +407,12 @@ class YqlDialect(StrCompileDialect):
         )
         return f"{declarations}\n{statement}"
 
+    def _apply_statement_prefixes_impl(self, statement: str) -> str:
+        if not self._statement_prefixes:
+            return statement
+        prefixes = "\n".join(self._statement_prefixes) + "\n"
+        return f"{prefixes}{statement}"
+
     def __merge_parameters_values_and_types(
         self, values: Mapping[str, Any], types: Mapping[str, Any], execute_many: bool
     ) -> Sequence[Mapping[str, ydb.TypedValue]]:
@@ -437,9 +446,12 @@ class YqlDialect(StrCompileDialect):
             statement, parameters = self._format_variables(statement, parameters, execute_many)
             if self._add_declare_for_yql_stmt_vars:
                 statement = self._add_declare_for_yql_stmt_vars_impl(statement, parameters_types)
+            statement = self._apply_statement_prefixes_impl(statement)
             return statement, parameters
 
         statement, parameters = self._format_variables(statement, parameters, execute_many)
+        if not is_ddl:
+            statement = self._apply_statement_prefixes_impl(statement)
         return statement, parameters
 
     def do_ping(self, dbapi_connection: ydb_dbapi.Connection) -> bool:

ydb_sqlalchemy/sqlalchemy/compiler/base.py

@@ -44,19 +44,6 @@ else:
     from sqlalchemy import Cast as _cast
 
 
-STR_QUOTE_MAP = {
-    "'": "\\'",
-    "\\": "\\\\",
-    "\0": "\\0",
-    "\b": "\\b",
-    "\f": "\\f",
-    "\r": "\\r",
-    "\n": "\\n",
-    "\t": "\\t",
-    "%": "%%",
-}
-
-
 COMPOUND_KEYWORDS = {
     selectable.CompoundSelect.UNION: "UNION ALL",
     selectable.CompoundSelect.UNION_ALL: "UNION ALL",
@@ -67,6 +54,19 @@ COMPOUND_KEYWORDS = {
 }
 
 
+ESCAPE_RULES = [
+    ("\\", "\\\\"),  # Must be first to avoid double escaping
+    ("'", "\\'"),
+    ("\0", "\\0"),
+    ("\b", "\\b"),
+    ("\f", "\\f"),
+    ("\r", "\\r"),
+    ("\n", "\\n"),
+    ("\t", "\\t"),
+    ("%", "%%"),
+]
+
+
 class BaseYqlTypeCompiler(StrSQLTypeCompiler):
     def visit_JSON(self, type_: Union[sa.JSON, types.YqlJSON], **kw):
         return "JSON"
@@ -293,7 +293,8 @@ class BaseYqlCompiler(StrSQLCompiler):
 
     def render_literal_value(self, value, type_):
         if isinstance(value, str):
-            value = "".join(STR_QUOTE_MAP.get(x, x) for x in value)
+            for pattern, replacement in ESCAPE_RULES:
+                value = value.replace(pattern, replacement)
             return f"'{value}'"
         return super().render_literal_value(value, type_)
 

ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py

@@ -114,6 +114,30 @@ def test_types_compilation():
     assert compile_type(struct) == "Struct<a:Int32,b:List<Int32>>"
 
 
+def test_statement_prefixes_prepended_to_query():
+    dialect = YqlDialect(_statement_prefixes_list=["PRAGMA DistinctOverKeys;"])
+    result = dialect._apply_statement_prefixes_impl("SELECT 1")
+    assert result == "PRAGMA DistinctOverKeys;\nSELECT 1"
+
+
+def test_statement_prefixes_empty_list_unchanged():
+    dialect = YqlDialect(_statement_prefixes_list=[])
+    result = dialect._apply_statement_prefixes_impl("SELECT 1")
+    assert result == "SELECT 1"
+
+
+def test_statement_prefixes_none_unchanged():
+    dialect = YqlDialect()
+    result = dialect._apply_statement_prefixes_impl("SELECT 1")
+    assert result == "SELECT 1"
+
+
+def test_statement_prefixes_multiple():
+    dialect = YqlDialect(_statement_prefixes_list=["PRAGMA Foo;", "PRAGMA Bar;"])
+    result = dialect._apply_statement_prefixes_impl("SELECT 1")
+    assert result == "PRAGMA Foo;\nPRAGMA Bar;\nSELECT 1"
+
+
 def test_optional_type_compilation():
     dialect = YqlDialect()
     type_compiler = dialect.type_compiler

ydb_sqlalchemy/sqlalchemy/types.py

@@ -110,6 +110,18 @@ class Decimal(types.DECIMAL):
 class ListType(ARRAY):
     __visit_name__ = "list_type"
 
+    def bind_processor(self, dialect):
+        item_proc = self.item_type.bind_processor(dialect)
+
+        def process(value):
+            if value is None:
+                return None
+            return [item_proc(v) if v is not None else None for v in value]
+
+        if item_proc:
+            return process
+        return None
+
 
 class HashableDict(dict):
     def __hash__(self):
@@ -173,6 +185,32 @@ class StructType(types.TypeEngine[Mapping[str, Any]]):
     def compare_values(self, x, y):
         return x == y
 
+    def bind_processor(self, dialect):
+        processors = {}
+        for name, type_ in self.fields_types.items():
+            if isinstance(type_, Optional):
+                type_ = type_.element_type
+
+            type_ = type_api.to_instance(type_)
+            proc = type_.bind_processor(dialect)
+            if proc:
+                processors[name] = proc
+
+        if not processors:
+            return None
+
+        def process(value):
+            if value is None:
+                return None
+            new_value = value.copy()
+            for name, proc in processors.items():
+                if name in new_value:
+                    if new_value[name] is not None:
+                        new_value[name] = proc(new_value[name])
+            return new_value
+
+        return process
+
 
 class Lambda(ColumnElement):
     __visit_name__ = "lambda"

{ydb_sqlalchemy-0.1.15.dist-info → ydb_sqlalchemy-0.1.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ydb-sqlalchemy
-Version: 0.1.15
+Version: 0.1.17
 Summary: YDB Dialect for SQLAlchemy
 Home-page: http://github.com/ydb-platform/ydb-sqlalchemy
 Author: Yandex LLC

{ydb_sqlalchemy-0.1.15.dist-info → ydb_sqlalchemy-0.1.17.dist-info}/RECORD

@@ -1,26 +1,26 @@
 test/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 test/conftest.py,sha256=rhWa0EQB9EwO_wAwxPdK17Qi582DdbBE8p5Gv4180Ds,570
-test/test_core.py,sha256=5ijngbM6b_8avees4CqCv0FgzAza2RRPDwgT2saU99k,47254
+test/test_core.py,sha256=cp49GDKgpsxfCBG9EjSwjOlPSchKW68nn8CnMeuxy8k,48929
 test/test_inspect.py,sha256=c4kc3jc48MCOfllO-ciiYf1vO-HOfuv0xVoXYT1Jxro,1106
 test/test_orm.py,sha256=jQVVld50zbUwxwgW9ySIWGaNDEOLzHKXjTkdpsG9TpA,1825
 test/test_suite.py,sha256=JYBGZjaRbg_ZiAqTHeCfL7DLnB6N6xkXN82gnooCyd8,31063
 ydb_sqlalchemy/__init__.py,sha256=hX7Gy-KOiHk7B5-0wj3ZmLjk4YDJnSMHIAqxVGn_PJY,181
-ydb_sqlalchemy/_version.py,sha256=f1tdrTNNxSCrK-kaNyZL5_MXChoOPzxvT7JLn3_KW6k,19
-ydb_sqlalchemy/sqlalchemy/__init__.py,sha256=Z5aprCPByfUmH2DvybA0o7W0lQ8YDCR1XSyyfGmBdoE,17911
+ydb_sqlalchemy/_version.py,sha256=CQKgFdKOd8j2eckBUFbwPNJt1oASUfatjq5eO0eF8cA,19
+ydb_sqlalchemy/sqlalchemy/__init__.py,sha256=8_wPXgqlhyIW0CGmG1vIsjuepGyCOl8rAm7E0Hv90DY,18500
 ydb_sqlalchemy/sqlalchemy/datetime_types.py,sha256=wrI9kpsI_f7Jhbm7Fu0o_S1QoGCLIe6A9jfUwb41aMM,1929
 ydb_sqlalchemy/sqlalchemy/dbapi_adapter.py,sha256=7FDjganh9QStIkoXYPFfcRRhd07YCX63_8OmMnge1FI,3542
 ydb_sqlalchemy/sqlalchemy/dml.py,sha256=k_m6PLOAY7dVzG1gsyo2bB3Lp-o3rhzN0oSX_nfkbFU,310
 ydb_sqlalchemy/sqlalchemy/json.py,sha256=b4ydjlQjBhlhqGP_Sy2uZVKmt__D-9M7-YLGQMdYGME,1043
 ydb_sqlalchemy/sqlalchemy/requirements.py,sha256=zm6fcLormtk3KHnbtrBvxfkbG9ZyzNan38HrRB6vC3c,2505
-ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py,sha256=4wyRHmE8YQaMElQPHX6ToEj7A9F8Mvv909aaz_0wRnA,4535
-ydb_sqlalchemy/sqlalchemy/types.py,sha256=_PxK76x6BJZv-7gVdXL-XJtaOFjiYZY4f-_WBRWH8x0,4996
+ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py,sha256=kSQ2sg90X5QoEAKfJSj3CIqI6j95tLUoUW8Fa197Um0,5410
+ydb_sqlalchemy/sqlalchemy/types.py,sha256=lFopPGW8WLPIrhgXEcnSAqzJMyrV74ryLSaeTnaKBzA,6125
 ydb_sqlalchemy/sqlalchemy/compiler/__init__.py,sha256=QqA6r-_bw1R97nQZy5ZSJN724znXg88l4mi5PpqAOxI,492
-ydb_sqlalchemy/sqlalchemy/compiler/base.py,sha256=zKOf4SrcAjMKf1LjRvQA7IzMRic0EIznHuZTHySvc_k,20143
+ydb_sqlalchemy/sqlalchemy/compiler/base.py,sha256=JyDIEk3CvJggidFLSNoeI4Khjk3a-b0SjS80GrjgSIk,20248
 ydb_sqlalchemy/sqlalchemy/compiler/sa14.py,sha256=LanxAnwOiMnsnrY05B0jpmvGn5NXuOKMcxi_6N3obVM,1186
 ydb_sqlalchemy/sqlalchemy/compiler/sa20.py,sha256=rvVhe-pq5bOyuW4KMMMAD7JIWMzy355eijymBvuPwKw,3421
-ydb_sqlalchemy-0.1.15.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ydb_sqlalchemy-0.1.15.dist-info/METADATA,sha256=iOoVXzaJ3tJ7ct7jssAXLSUwptb9C4e7Yx3dPp8VzdA,5395
-ydb_sqlalchemy-0.1.15.dist-info/WHEEL,sha256=Ll72iyqtt6Rbxp-Q7FSafYA1LeRv98X15xcZWRsFEmY,109
-ydb_sqlalchemy-0.1.15.dist-info/entry_points.txt,sha256=iJxbKYuliWNBmL0iIiw8MxvOXrSEz5xe5fuEBqMRwCE,267
-ydb_sqlalchemy-0.1.15.dist-info/top_level.txt,sha256=iS69Y1GTAcTok0u0oQdxP-Q5iVgUGI71XBsaEUrWhMg,20
-ydb_sqlalchemy-0.1.15.dist-info/RECORD,,
+ydb_sqlalchemy-0.1.17.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ydb_sqlalchemy-0.1.17.dist-info/METADATA,sha256=fhxTYNZXhbJs_hoyerJNTPmdqjVXSDbUa84YmsacZts,5395
+ydb_sqlalchemy-0.1.17.dist-info/WHEEL,sha256=I3glN-nznogni2CWkAvi6vPRkBhTbVeYvJZnIsln6uc,109
+ydb_sqlalchemy-0.1.17.dist-info/entry_points.txt,sha256=iJxbKYuliWNBmL0iIiw8MxvOXrSEz5xe5fuEBqMRwCE,267
+ydb_sqlalchemy-0.1.17.dist-info/top_level.txt,sha256=iS69Y1GTAcTok0u0oQdxP-Q5iVgUGI71XBsaEUrWhMg,20
+ydb_sqlalchemy-0.1.17.dist-info/RECORD,,

{ydb_sqlalchemy-0.1.15.dist-info → ydb_sqlalchemy-0.1.17.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.3.2)
+Generator: setuptools (75.3.3)
 Root-Is-Purelib: true
 Tag: py2-none-any
 Tag: py3-none-any