ydb-sqlalchemy 0.1.15__tar.gz → 0.1.17__tar.gz

{ydb_sqlalchemy-0.1.15/ydb_sqlalchemy.egg-info → ydb_sqlalchemy-0.1.17}/PKG-INFO

@@ -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 → ydb_sqlalchemy-0.1.17}/setup.py

@@ -13,7 +13,7 @@ with open("requirements.txt") as f:
 
 setuptools.setup(
     name="ydb-sqlalchemy",
-    version="0.1.15",  # AUTOVERSION
+    version="0.1.17",  # AUTOVERSION
     description="YDB Dialect for SQLAlchemy",
     author="Yandex LLC",
     author_email="ydb@yandex-team.ru",

{ydb_sqlalchemy-0.1.15 → ydb_sqlalchemy-0.1.17}/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-0.1.17/ydb_sqlalchemy/_version.py

@@ -0,0 +1 @@
+VERSION = "0.1.17"

{ydb_sqlalchemy-0.1.15 → ydb_sqlalchemy-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-0.1.15 → ydb_sqlalchemy-0.1.17}/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-0.1.15 → ydb_sqlalchemy-0.1.17}/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-0.1.15 → ydb_sqlalchemy-0.1.17}/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 → ydb_sqlalchemy-0.1.17/ydb_sqlalchemy.egg-info}/PKG-INFO

@@ -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/ydb_sqlalchemy/_version.py

@@ -1 +0,0 @@
-VERSION = "0.1.15"