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

test/test_core.py

@@ -228,6 +228,15 @@ class TestTypes(TablesTest):
             Column("float", sa.Float),
             Column("bool", sa.Boolean),
         )
+        Table(
+            "test_all_binary_types",
+            metadata,
+            Column("id", sa.Integer, primary_key=True),
+            Column("bin", sa.BINARY),
+            Column("large_bin", sa.LargeBinary),
+            Column("blob", sa.BLOB),
+            Column("custom_bin", types.Binary),
+        )
         Table(
             "test_datetime_types",
             metadata,
@@ -255,6 +264,21 @@ class TestTypes(TablesTest):
         row = connection.execute(sa.select(table)).fetchone()
         assert row == (42, "Hello World!", 3.5, True)
 
+    def test_all_binary_types(self, connection):
+        table = self.tables.test_all_binary_types
+        data = {
+            "id": 1,
+            "bin": b"binary",
+            "large_bin": b"large_binary",
+            "blob": b"blob",
+            "custom_bin": b"custom_binary",
+        }
+        statement = sa.insert(table).values(**data)
+        connection.execute(statement)
+
+        row = connection.execute(sa.select(table)).fetchone()
+        assert row == (1, b"binary", b"large_binary", b"blob", b"custom_binary")
+
     def test_integer_types(self, connection):
         stmt = sa.select(
             sa.func.FormatType(sa.func.TypeOf(sa.bindparam("p_uint8", 8, types.UInt8))),
@@ -1112,3 +1136,114 @@ class TestTablePathPrefix(TablesTest):
         metadata.reflect(reflection_engine)
 
         assert "nested_dir/table" in metadata.tables
+
+
+class TestAsTable(TablesTest):
+    __backend__ = True
+
+    @classmethod
+    def define_tables(cls, metadata):
+        Table(
+            "test_as_table",
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("val_int", Integer, nullable=True),
+            Column("val_str", String, nullable=True),
+        )
+
+    def test_upsert_as_table(self, connection):
+        table = self.tables.test_as_table
+
+        input_data = [
+            {"id": 1, "val_int": 10, "val_str": "a"},
+            {"id": 2, "val_int": None, "val_str": "b"},
+            {"id": 3, "val_int": 30, "val_str": None},
+        ]
+
+        struct_type = types.StructType(
+            {
+                "id": Integer,
+                "val_int": types.Optional(Integer),
+                "val_str": types.Optional(String),
+            }
+        )
+        list_type = types.ListType(struct_type)
+
+        bind_param = sa.bindparam("data", type_=list_type)
+
+        upsert_stm = ydb_sa.upsert(table).from_select(
+            ["id", "val_int", "val_str"],
+            sa.select(
+                sa.column("id", type_=Integer), sa.column("val_int", type_=Integer), sa.column("val_str", type_=String)
+            ).select_from(sa.func.AS_TABLE(bind_param)),
+        )
+
+        connection.execute(upsert_stm, {"data": input_data})
+
+        rows = connection.execute(sa.select(table).order_by(table.c.id)).fetchall()
+        assert rows == [
+            (1, 10, "a"),
+            (2, None, "b"),
+            (3, 30, None),
+        ]
+
+    def test_insert_as_table(self, connection):
+        table = self.tables.test_as_table
+
+        input_data = [
+            {"id": 4, "val_int": 40, "val_str": "d"},
+            {"id": 5, "val_int": None, "val_str": "e"},
+        ]
+
+        struct_type = types.StructType(
+            {
+                "id": Integer,
+                "val_int": types.Optional(Integer),
+                "val_str": types.Optional(String),
+            }
+        )
+        list_type = types.ListType(struct_type)
+
+        bind_param = sa.bindparam("data", type_=list_type)
+
+        insert_stm = sa.insert(table).from_select(
+            ["id", "val_int", "val_str"],
+            sa.select(
+                sa.column("id", type_=Integer), sa.column("val_int", type_=Integer), sa.column("val_str", type_=String)
+            ).select_from(sa.func.AS_TABLE(bind_param)),
+        )
+
+        connection.execute(insert_stm, {"data": input_data})
+
+        rows = connection.execute(sa.select(table).where(table.c.id >= 4).order_by(table.c.id)).fetchall()
+        assert rows == [
+            (4, 40, "d"),
+            (5, None, "e"),
+        ]
+
+    def test_upsert_from_table_reflection(self, connection):
+        table = self.tables.test_as_table
+
+        input_data = [
+            {"id": 1, "val_int": 10, "val_str": "a"},
+            {"id": 2, "val_int": None, "val_str": "b"},
+        ]
+
+        struct_type = types.StructType.from_table(table)
+        list_type = types.ListType(struct_type)
+
+        bind_param = sa.bindparam("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, {"data": input_data})
+
+        rows = connection.execute(sa.select(table).order_by(table.c.id)).fetchall()
+        assert rows == [
+            (1, 10, "a"),
+            (2, None, "b"),
+        ]

test/test_suite.py

@@ -274,7 +274,6 @@ class NumericTest(_NumericTest):
     pass
 
 
-@pytest.mark.skip("TODO: see issue #13")
 class BinaryTest(_BinaryTest):
     pass
 

ydb_sqlalchemy/_version.py

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

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -157,6 +157,9 @@ class YqlDialect(StrCompileDialect):
         sa.types.DATETIME: types.YqlDateTime,
         sa.types.TIMESTAMP: types.YqlTimestamp,
         sa.types.DECIMAL: types.Decimal,
+        sa.types.BINARY: types.Binary,
+        sa.types.LargeBinary: types.Binary,
+        sa.types.BLOB: types.Binary,
     }
 
     connection_characteristics = util.immutabledict(

ydb_sqlalchemy/sqlalchemy/compiler/base.py

@@ -12,6 +12,7 @@ from sqlalchemy.sql.compiler import (
     StrSQLTypeCompiler,
     selectable,
 )
+from sqlalchemy.sql.type_api import to_instance
 from typing import (
     Any,
     Dict,
@@ -24,6 +25,12 @@ from typing import (
     Union,
 )
 
+try:
+    from sqlalchemy.types import _Binary as _BinaryType
+except ImportError:
+    # For older sqlalchemy versions
+    from sqlalchemy.sql.sqltypes import _Binary as _BinaryType
+
 
 from .. import types
 
@@ -152,11 +159,17 @@ class BaseYqlTypeCompiler(StrSQLTypeCompiler):
         inner = self.process(type_.item_type, **kw)
         return f"List<{inner}>"
 
+    def visit_optional(self, type_: types.Optional, **kw):
+        el = to_instance(type_.element_type)
+        inner = self.process(el, **kw)
+        return f"Optional<{inner}>"
+
     def visit_struct_type(self, type_: types.StructType, **kw):
-        text = "Struct<"
-        for field, field_type in type_.fields_types:
-            text += f"{field}:{self.process(field_type, **kw)}"
-        return text + ">"
+        rendered_types = []
+        for field, field_type in type_.fields_types.items():
+            type_str = self.process(field_type, **kw)
+            rendered_types.append(f"{field}:{type_str}")
+        return f"Struct<{','.join(rendered_types)}>"
 
     def get_ydb_type(
         self, type_: sa.types.TypeEngine, is_optional: bool
@@ -167,6 +180,10 @@ class BaseYqlTypeCompiler(StrSQLTypeCompiler):
         if isinstance(type_, (sa.Text, sa.String)):
             ydb_type = ydb.PrimitiveType.Utf8
 
+        elif isinstance(type_, types.Optional):
+            inner = to_instance(type_.element_type)
+            return self.get_ydb_type(inner, is_optional=True)
+
         # Integers
         elif isinstance(type_, types.UInt64):
             ydb_type = ydb.PrimitiveType.Uint64
@@ -216,7 +233,7 @@ class BaseYqlTypeCompiler(StrSQLTypeCompiler):
             ydb_type = ydb.PrimitiveType.Timestamp
         elif isinstance(type_, sa.Date):
             ydb_type = ydb.PrimitiveType.Date
-        elif isinstance(type_, sa.BINARY):
+        elif isinstance(type_, _BinaryType):
             ydb_type = ydb.PrimitiveType.String
         elif isinstance(type_, sa.Float):
             ydb_type = ydb.PrimitiveType.Float
@@ -235,7 +252,8 @@ class BaseYqlTypeCompiler(StrSQLTypeCompiler):
         elif isinstance(type_, types.StructType):
             ydb_type = ydb.StructType()
             for field, field_type in type_.fields_types.items():
-                ydb_type.add_member(field, self.get_ydb_type(field_type(), is_optional=False))
+                inner_type = to_instance(field_type)
+                ydb_type.add_member(field, self.get_ydb_type(inner_type, is_optional=False))
         else:
             raise NotSupportedError(f"{type_} bind variables not supported")
 

ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py

@@ -35,3 +35,111 @@ def test_ydb_types():
     compiled = query.compile(dialect=dialect, compile_kwargs={"literal_binds": True})
 
     assert str(compiled) == "Date('1996-11-19')"
+
+
+def test_binary_type():
+    dialect = YqlDialect()
+    expr = sa.literal(b"some bytes")
+    compiled = expr.compile(dialect=dialect, compile_kwargs={"literal_binds": True})
+    assert str(compiled) == "'some bytes'"
+
+    expr_binary = sa.cast(expr, sa.BINARY)
+    compiled_binary = expr_binary.compile(dialect=dialect, compile_kwargs={"literal_binds": True})
+    assert str(compiled_binary) == "CAST('some bytes' AS String)"
+
+
+def test_all_binary_types():
+    dialect = YqlDialect()
+    expr = sa.literal(b"some bytes")
+
+    binary_types = [
+        sa.BINARY,
+        sa.LargeBinary,
+        sa.BLOB,
+        types.Binary,
+    ]
+
+    for type_ in binary_types:
+        expr_binary = sa.cast(expr, type_)
+        compiled_binary = expr_binary.compile(dialect=dialect, compile_kwargs={"literal_binds": True})
+        assert str(compiled_binary) == "CAST('some bytes' AS String)"
+
+
+def test_struct_type_generation():
+    dialect = YqlDialect()
+    type_compiler = dialect.type_compiler
+
+    # Test default (non-optional)
+    struct_type = types.StructType(
+        {
+            "id": sa.Integer,
+            "val_int": sa.Integer,
+        }
+    )
+    ydb_type = type_compiler.get_ydb_type(struct_type, is_optional=False)
+    # Keys are sorted
+    assert str(ydb_type) == "Struct<id:Int64,val_int:Int64>"
+
+    # Test optional
+    struct_type_opt = types.StructType(
+        {
+            "id": sa.Integer,
+            "val_int": types.Optional(sa.Integer),
+        }
+    )
+    ydb_type_opt = type_compiler.get_ydb_type(struct_type_opt, is_optional=False)
+    assert str(ydb_type_opt) == "Struct<id:Int64,val_int:Int64?>"
+
+
+def test_types_compilation():
+    dialect = YqlDialect()
+
+    def compile_type(type_):
+        return dialect.type_compiler.process(type_)
+
+    assert compile_type(types.UInt64()) == "UInt64"
+    assert compile_type(types.UInt32()) == "UInt32"
+    assert compile_type(types.UInt16()) == "UInt16"
+    assert compile_type(types.UInt8()) == "UInt8"
+
+    assert compile_type(types.Int64()) == "Int64"
+    assert compile_type(types.Int32()) == "Int32"
+    assert compile_type(types.Int16()) == "Int32"
+    assert compile_type(types.Int8()) == "Int8"
+
+    assert compile_type(types.ListType(types.Int64())) == "List<Int64>"
+
+    struct = types.StructType({"a": types.Int32(), "b": types.ListType(types.Int32())})
+    # Ordered by key: a, b
+    assert compile_type(struct) == "Struct<a:Int32,b:List<Int32>>"
+
+
+def test_optional_type_compilation():
+    dialect = YqlDialect()
+    type_compiler = dialect.type_compiler
+
+    def compile_type(type_):
+        return type_compiler.process(type_)
+
+    # Test Optional(Integer)
+    opt_int = types.Optional(sa.Integer)
+    assert compile_type(opt_int) == "Optional<Int64>"
+
+    # Test Optional(String)
+    opt_str = types.Optional(sa.String)
+    assert compile_type(opt_str) == "Optional<UTF8>"
+
+    # Test Nested Optional
+    opt_opt_int = types.Optional(types.Optional(sa.Integer))
+    assert compile_type(opt_opt_int) == "Optional<Optional<Int64>>"
+
+    # Test get_ydb_type
+    ydb_type = type_compiler.get_ydb_type(opt_int, is_optional=False)
+    import ydb
+
+    assert isinstance(ydb_type, ydb.OptionalType)
+    # Int64 corresponds to PrimitiveType.Int64
+    # Note: ydb.PrimitiveType.Int64 is an enum member, but ydb_type.item is also an instance/enum?
+    # get_ydb_type returns ydb.PrimitiveType.Int64 (enum) wrapped in OptionalType.
+    # OptionalType.item is the inner type.
+    assert ydb_type.item == ydb.PrimitiveType.Int64

ydb_sqlalchemy/sqlalchemy/types.py

@@ -8,7 +8,7 @@ if sa_version.startswith("2."):
 else:
     from sqlalchemy.sql.expression import ColumnElement
 
-from sqlalchemy import ARRAY, exc, types
+from sqlalchemy import ARRAY, exc, Table, types
 from sqlalchemy.sql import type_api
 
 from .datetime_types import YqlDate, YqlDateTime, YqlTimestamp, YqlDate32, YqlTimestamp64, YqlDateTime64  # noqa: F401
@@ -116,12 +116,56 @@ class HashableDict(dict):
         return hash(tuple(self.items()))
 
 
+class Optional(types.TypeEngine):
+    """
+    Wrapper for YDB Optional type.
+
+    Used primarily within StructType to denote nullable fields.
+    """
+
+    __visit_name__ = "optional"
+
+    def __init__(self, element_type: Union[Type[types.TypeEngine], types.TypeEngine]):
+        self.element_type = element_type
+
+
 class StructType(types.TypeEngine[Mapping[str, Any]]):
+    """
+    YDB Struct type.
+
+    Represents a structured data type with named fields, mapped to a Python dictionary.
+    """
+
     __visit_name__ = "struct_type"
 
-    def __init__(self, fields_types: Mapping[str, Union[Type[types.TypeEngine], Type[types.TypeDecorator]]]):
+    def __init__(
+        self,
+        fields_types: Mapping[
+            str,
+            Union[Type[types.TypeEngine], types.TypeEngine, Optional],
+        ],
+    ):
         self.fields_types = HashableDict(dict(sorted(fields_types.items())))
 
+    @classmethod
+    def from_table(cls, table: Table) -> "StructType":
+        """
+        Create a StructType definition from a SQLAlchemy Table.
+
+        Automatically wraps nullable columns in Optional.
+
+        :param table: SQLAlchemy Table object
+        :return: StructType instance
+        """
+        fields = {}
+        for col in table.columns:
+            t = col.type
+            if col.nullable:
+                fields[col.name] = Optional(t)
+            else:
+                fields[col.name] = t
+        return cls(fields)
+
     @property
     def python_type(self):
         return dict
@@ -139,3 +183,10 @@ class Lambda(ColumnElement):
 
         self.type = type_api.NULLTYPE
         self.func = func
+
+
+class Binary(types.LargeBinary):
+    __visit_name__ = "BINARY"
+
+    def bind_processor(self, dialect):
+        return None

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

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

{ydb_sqlalchemy-0.1.14.dist-info → ydb_sqlalchemy-0.1.15.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=XvPJ0MtWK2gqGytps4YMUpHtJWKlEqN1rQBUpeeelAg,42859
+test/test_core.py,sha256=5ijngbM6b_8avees4CqCv0FgzAza2RRPDwgT2saU99k,47254
 test/test_inspect.py,sha256=c4kc3jc48MCOfllO-ciiYf1vO-HOfuv0xVoXYT1Jxro,1106
 test/test_orm.py,sha256=jQVVld50zbUwxwgW9ySIWGaNDEOLzHKXjTkdpsG9TpA,1825
-test/test_suite.py,sha256=3Ht_mEG5zzYUA_7YH3Hax5EZgnGhnPrN0PBuJZvdR9I,31104
+test/test_suite.py,sha256=JYBGZjaRbg_ZiAqTHeCfL7DLnB6N6xkXN82gnooCyd8,31063
 ydb_sqlalchemy/__init__.py,sha256=hX7Gy-KOiHk7B5-0wj3ZmLjk4YDJnSMHIAqxVGn_PJY,181
-ydb_sqlalchemy/_version.py,sha256=7CBWvojWQb3zsrGuOA2awejgwAOfTSzzu4lhva5fpfY,19
-ydb_sqlalchemy/sqlalchemy/__init__.py,sha256=vAeQNsLCuzSzOlvyzMnw_hwOzDY5uqLscdX3flyllj8,17791
+ydb_sqlalchemy/_version.py,sha256=f1tdrTNNxSCrK-kaNyZL5_MXChoOPzxvT7JLn3_KW6k,19
+ydb_sqlalchemy/sqlalchemy/__init__.py,sha256=Z5aprCPByfUmH2DvybA0o7W0lQ8YDCR1XSyyfGmBdoE,17911
 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=QPRgUPsTOqAf0gzWjidsIhTPVkfOILy4SCSgNb1GE7o,1039
-ydb_sqlalchemy/sqlalchemy/types.py,sha256=BVeC8RSa2nFnI3-Q7XZxCnAolHc_k_kMCfWhUmbi5wo,3823
+ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py,sha256=4wyRHmE8YQaMElQPHX6ToEj7A9F8Mvv909aaz_0wRnA,4535
+ydb_sqlalchemy/sqlalchemy/types.py,sha256=_PxK76x6BJZv-7gVdXL-XJtaOFjiYZY4f-_WBRWH8x0,4996
 ydb_sqlalchemy/sqlalchemy/compiler/__init__.py,sha256=QqA6r-_bw1R97nQZy5ZSJN724znXg88l4mi5PpqAOxI,492
-ydb_sqlalchemy/sqlalchemy/compiler/base.py,sha256=zbd9uhV0uxHfMkSTyi4eoZiaiS32yNid4IWfGUo82vA,19433
+ydb_sqlalchemy/sqlalchemy/compiler/base.py,sha256=zKOf4SrcAjMKf1LjRvQA7IzMRic0EIznHuZTHySvc_k,20143
 ydb_sqlalchemy/sqlalchemy/compiler/sa14.py,sha256=LanxAnwOiMnsnrY05B0jpmvGn5NXuOKMcxi_6N3obVM,1186
 ydb_sqlalchemy/sqlalchemy/compiler/sa20.py,sha256=rvVhe-pq5bOyuW4KMMMAD7JIWMzy355eijymBvuPwKw,3421
-ydb_sqlalchemy-0.1.14.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ydb_sqlalchemy-0.1.14.dist-info/METADATA,sha256=Q1Mv9kMxkNqehnueGkrzrICV_HQ3Klgy30wA18YSqyM,5395
-ydb_sqlalchemy-0.1.14.dist-info/WHEEL,sha256=Ll72iyqtt6Rbxp-Q7FSafYA1LeRv98X15xcZWRsFEmY,109
-ydb_sqlalchemy-0.1.14.dist-info/entry_points.txt,sha256=iJxbKYuliWNBmL0iIiw8MxvOXrSEz5xe5fuEBqMRwCE,267
-ydb_sqlalchemy-0.1.14.dist-info/top_level.txt,sha256=iS69Y1GTAcTok0u0oQdxP-Q5iVgUGI71XBsaEUrWhMg,20
-ydb_sqlalchemy-0.1.14.dist-info/RECORD,,
+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,,