Merge pull request #221 from ydb-platform/sqlalchemy-fix

sqlalchemy - fix not null inspection

Author: Valeria1235

CHANGELOG.md

@@ -1,3 +1,6 @@
+## 2.13.4 ##
+* fixed sqlalchemy get_columns method with not null columns
+
 ## 2.13.3 ##
 * fixed use transaction object when commit with flag
 

tests/sqlalchemy/conftest.py

@@ -0,0 +1,16 @@
+import pytest
+import sqlalchemy as sa
+
+from ydb.sqlalchemy import register_dialect
+
+
+@pytest.fixture
+def sa_engine(endpoint, database):
+    register_dialect()
+    engine = sa.create_engine(
+        "yql:///ydb/",
+        connect_args={"database": database, "endpoint": endpoint},
+    )
+
+    yield engine
+    engine.dispose()

tests/sqlalchemy/test_inspect.py

@@ -0,0 +1,24 @@
+import ydb
+
+import sqlalchemy as sa
+
+
+def test_get_columns(driver_sync, sa_engine):
+    session = ydb.retry_operation_sync(
+        lambda: driver_sync.table_client.session().create()
+    )
+    session.execute_scheme(
+        "CREATE TABLE test(id Int64 NOT NULL, value TEXT, num DECIMAL(22, 9), PRIMARY KEY (id))"
+    )
+    inspect = sa.inspect(sa_engine)
+    columns = inspect.get_columns("test")
+    for c in columns:
+        c["type"] = type(c["type"])
+
+    assert columns == [
+        {"name": "id", "type": sa.INTEGER, "nullable": False},
+        {"name": "value", "type": sa.TEXT, "nullable": True},
+        {"name": "num", "type": sa.DECIMAL, "nullable": True},
+    ]
+
+    session.execute_scheme("DROP TABLE test")

ydb/sqlalchemy/__init__.py

@@ -191,11 +191,16 @@ def visit_function(self, func, add_to_result_map=None, **kwargs):
         ydb.PrimitiveType.DyNumber: sa.TEXT,
     }
 
-    def _get_column_type(t):
-        if isinstance(t.item, ydb.DecimalType):
-            return sa.DECIMAL(precision=t.item.precision, scale=t.item.scale)
+    def _get_column_info(t):
+        nullable = False
+        if isinstance(t, ydb.OptionalType):
+            nullable = True
+            t = t.item
 
-        return COLUMN_TYPES[t.item]
+        if isinstance(t, ydb.DecimalType):
+            return sa.DECIMAL(precision=t.precision, scale=t.scale), nullable
+
+        return COLUMN_TYPES[t], nullable
 
     class YqlDialect(DefaultDialect):
         name = "yql"
@@ -250,11 +255,12 @@ def get_columns(self, connection, table_name, schema=None, **kw):
             columns = raw_conn.describe(qt)
             as_compatible = []
             for column in columns:
+                col_type, nullable = _get_column_info(column.type)
                 as_compatible.append(
                     {
                         "name": column.name,
-                        "type": _get_column_type(column.type),
-                        "nullable": True,
+                        "type": col_type,
+                        "nullable": nullable,
                     }
                 )