Merge pull request #28 Add basic JSON support from LuckySting/json-type

Author: rekby

test/test_suite.py

@@ -1,3 +1,5 @@
+import ctypes
+
 import pytest
 import sqlalchemy as sa
 import sqlalchemy.testing.suite.test_types
@@ -65,6 +67,7 @@
 )
 from sqlalchemy.testing.suite.test_types import DateTimeTest as _DateTimeTest
 from sqlalchemy.testing.suite.test_types import IntegerTest as _IntegerTest
+from sqlalchemy.testing.suite.test_types import JSONTest as _JSONTest
 from sqlalchemy.testing.suite.test_types import NativeUUIDTest as _NativeUUIDTest
 from sqlalchemy.testing.suite.test_types import NumericTest as _NumericTest
 from sqlalchemy.testing.suite.test_types import StringTest as _StringTest
@@ -430,6 +433,25 @@ class TimeTest(_TimeTest):
     pass
 
 
+class JSONTest(_JSONTest):
+    @classmethod
+    def define_tables(cls, metadata):
+        Table(
+            "data_table",
+            metadata,
+            Column("id", Integer, primary_key=True, default=1),
+            Column("name", String(30), primary_key=True, nullable=False),
+            Column("data", cls.datatype, nullable=False),
+            Column("nulldata", cls.datatype(none_as_null=True)),
+        )
+
+    def _json_value_insert(self, connection, datatype, value, data_element):
+        if datatype == "float" and value is not None:
+            # As python's float is stored as C double, it needs to be shrank
+            value = ctypes.c_float(value).value
+        return super()._json_value_insert(connection, datatype, value, data_element)
+
+
 class StringTest(_StringTest):
     @requirements.unbounded_varchar
     def test_nolength_string(self):

ydb_sqlalchemy/dbapi/connection.py

@@ -149,7 +149,7 @@ def _get_table_client_settings(self) -> ydb.TableClientSettings:
             .with_native_datetime_in_result_sets(True)
             .with_native_timestamp_in_result_sets(True)
             .with_native_interval_in_result_sets(True)
-            .with_native_json_in_result_sets(True)
+            .with_native_json_in_result_sets(False)
         )
 
     def _create_driver(self):

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -2,6 +2,7 @@
 Experimental
 Work in progress, breaking changes are possible.
 """
+
 import collections
 import collections.abc
 from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
@@ -63,6 +64,9 @@ def __init__(self, dialect):
 
 
 class YqlTypeCompiler(StrSQLTypeCompiler):
+    def visit_JSON(self, type_: Union[sa.JSON, types.YqlJSON], **kw):
+        return "JSON"
+
     def visit_CHAR(self, type_: sa.CHAR, **kw):
         return "UTF8"
 
@@ -171,8 +175,21 @@ def get_ydb_type(
             ydb_type = ydb.PrimitiveType.Int64
         # Integers
 
+        # Json
         elif isinstance(type_, sa.JSON):
             ydb_type = ydb.PrimitiveType.Json
+        elif isinstance(type_, sa.JSON.JSONStrIndexType):
+            ydb_type = ydb.PrimitiveType.Utf8
+        elif isinstance(type_, sa.JSON.JSONIntIndexType):
+            ydb_type = ydb.PrimitiveType.Int64
+        elif isinstance(type_, sa.JSON.JSONPathType):
+            ydb_type = ydb.PrimitiveType.Utf8
+        elif isinstance(type_, types.YqlJSON):
+            ydb_type = ydb.PrimitiveType.Json
+        elif isinstance(type_, types.YqlJSON.YqlJSONPathType):
+            ydb_type = ydb.PrimitiveType.Utf8
+        # Json
+
         elif isinstance(type_, sa.DateTime):
             ydb_type = ydb.PrimitiveType.Timestamp
         elif isinstance(type_, sa.Date):
@@ -326,6 +343,24 @@ def visit_function(self, func, add_to_result_map=None, **kwargs):
             + [name]
         ) % {"expr": self.function_argspec(func, **kwargs)}
 
+    def _yson_convert_to(self, statement: str, target_type: sa.types.TypeEngine) -> str:
+        type_name = target_type.compile(self.dialect)
+        if isinstance(target_type, sa.Numeric) and not isinstance(target_type, (sa.Float, sa.Double)):
+            # Since Decimal is stored in JSON either as String or as Float
+            string_value = f"Yson::ConvertTo({statement}, Optional<String>, Yson::Options(true AS AutoConvert))"
+            return f"CAST({string_value} AS Optional<{type_name}>)"
+        return f"Yson::ConvertTo({statement}, Optional<{type_name}>)"
+
+    def visit_json_getitem_op_binary(self, binary: sa.BinaryExpression, operator, **kw) -> str:
+        json_field = self.process(binary.left, **kw)
+        index = self.process(binary.right, **kw)
+        return self._yson_convert_to(f"{json_field}[{index}]", binary.type)
+
+    def visit_json_path_getitem_op_binary(self, binary: sa.BinaryExpression, operator, **kw) -> str:
+        json_field = self.process(binary.left, **kw)
+        path = self.process(binary.right, **kw)
+        return self._yson_convert_to(f"Yson::YPath({json_field}, {path})", binary.type)
+
     def visit_regexp_match_op_binary(self, binary, operator, **kw):
         return self._generate_generic_binary(binary, " REGEXP ", **kw)
 
@@ -336,7 +371,7 @@ def _is_bound_to_nullable_column(self, bind_name: str) -> bool:
         if bind_name in self.column_keys and hasattr(self.compile_state, "dml_table"):
             if bind_name in self.compile_state.dml_table.c:
                 column = self.compile_state.dml_table.c[bind_name]
-                return not column.primary_key
+                return column.nullable and not column.primary_key
         return False
 
     def _guess_bound_variable_type_by_parameters(
@@ -503,6 +538,7 @@ class YqlDialect(StrCompileDialect):
     supports_smallserial = False
     supports_schemas = False
     supports_constraint_comments = False
+    supports_json_type = True
 
     insert_returning = False
     update_returning = False
@@ -524,6 +560,10 @@ class YqlDialect(StrCompileDialect):
     statement_compiler = YqlCompiler
     ddl_compiler = YqlDDLCompiler
     type_compiler = YqlTypeCompiler
+    colspecs = {
+        sa.types.JSON: types.YqlJSON,
+        sa.types.JSON.JSONPathType: types.YqlJSON.YqlJSONPathType,
+    }
 
     construct_arguments = [
         (
@@ -544,6 +584,12 @@ class YqlDialect(StrCompileDialect):
     def import_dbapi(cls: Any):
         return dbapi.YdbDBApi()
 
+    def __init__(self, json_serializer=None, json_deserializer=None, **kwargs):
+        super().__init__(**kwargs)
+
+        self._json_deserializer = json_deserializer
+        self._json_serializer = json_serializer
+
     def _describe_table(self, connection, table_name, schema=None):
         if schema is not None:
             raise dbapi.NotSupportedError("unsupported on non empty schema")

ydb_sqlalchemy/sqlalchemy/json.py

@@ -0,0 +1,34 @@
+from typing import Tuple, Union
+
+from sqlalchemy import types as sqltypes
+
+
+class YqlJSON(sqltypes.JSON):
+    class YqlJSONPathType(sqltypes.JSON.JSONPathType):
+        def _format_value(self, value: Tuple[Union[str, int]]) -> str:
+            path = "/"
+            for elem in value:
+                path += f"/{elem}"
+            return path
+
+        def bind_processor(self, dialect):
+            super_proc = self.string_bind_processor(dialect)
+
+            def process(value: Tuple[Union[str, int]]):
+                value = self._format_value(value)
+                if super_proc:
+                    value = super_proc(value)
+                return value
+
+            return process
+
+        def literal_processor(self, dialect):
+            super_proc = self.string_literal_processor(dialect)
+
+            def process(value: Tuple[Union[str, int]]):
+                value = self._format_value(value)
+                if super_proc:
+                    value = super_proc(value)
+                return value
+
+            return process

ydb_sqlalchemy/sqlalchemy/requirements.py

@@ -3,6 +3,10 @@
 
 
 class Requirements(SuiteRequirements):
+    @property
+    def json_type(self):
+        return exclusions.open()
+
     @property
     def array_type(self):
         return exclusions.closed()

ydb_sqlalchemy/sqlalchemy/types.py

@@ -3,6 +3,8 @@
 from sqlalchemy import ARRAY, ColumnElement, exc, types
 from sqlalchemy.sql import type_api
 
+from .json import YqlJSON  # noqa: F401
+
 
 class UInt64(types.Integer):
     __visit_name__ = "uint64"