mirror of
https://github.com/hanzo-ds/connect.git
synced 2026-10-30 20:31:24 +00:00
cd902b864e
* Implement SQLAlchemy reflection for the cc_sqlalchemy dialect
The cc_sqlalchemy dialect did not support SQLAlchemy reflection
(MetaData.reflect / Inspector multi-table reflection), which broke
sqlacodegen and any tool that calls `dialect.get_columns()` directly.
This change fills in the missing dialect methods so reflection works
end-to-end against a ClickHouse server.
Changes:
1. dialect.py: add `ClickHouseDialect.get_columns()`. Previously only
`ChInspector.get_columns()` existed, but SQLAlchemy's reflection
path goes through `Dialect.get_multi_columns` -> `Dialect.get_columns`
and never touches `Inspector.get_columns` on its own. Without a
dialect implementation, `MetaData.reflect()` raised
`NotImplementedError` from the SQLAlchemy base class.
`get_pk_constraint()` / `get_primary_keys()` now return the
actual primary key columns derived from
`system.columns.is_in_primary_key` (which mirrors MergeTree's
ORDER BY / PRIMARY KEY) instead of empty lists. This lets
sqlacodegen generate declarative classes instead of bare
`Table(...)` definitions for any MergeTree table.
2. inspector.py: promote `get_columns` and `get_pk_constraint` to
module-level functions so the dialect can call the same logic.
`ChInspector.reflect_table()` now applies the PK constraint to
reflected columns (it was building columns with no PK info, so
even direct `Table('asset', md, autoload_with=engine)` reflection
lost the primary key).
3. datatypes/sqltypes.py: replace `python_type = None` on UDT-based
types with concrete Python types. SQLAlchemy's contract for
`TypeEngine.python_type` is that it either returns a class or
raises `NotImplementedError`; returning `None` makes any consumer
that does `python_type.__module__` / `__name__` crash with
`AttributeError: 'NoneType' object has no attribute '__module__'`
(sqlacodegen, and anything else that walks python_type for
annotations or metadata).
- UUID -> uuid.UUID
- IPv4 / IPv6 -> ipaddress.IPv4Address / IPv6Address
- Nothing -> type(None)
- Point -> tuple
- Ring / Polygon -> list
- LineString etc. -> list
- JSON -> dict
- Nested -> list
- (Simple)AggregateFunction -> str
4. datatypes/sqltypes.py: `Array` now subclasses
`sqlalchemy.types.ARRAY` (alongside `ChSqlaType`) and exposes
`item_type` as a regular instance attribute plus `dimensions = 1`.
Two effects:
- `isinstance(col.type, sqlalchemy.types.ARRAY)` now matches CH
arrays, which lets sqlacodegen render `Mapped[list[T]]`
annotations for single-dim arrays without special-casing.
- `item_type` is mutable so sqlacodegen's `fix_column_types`
adaptation pass (which reassigns `new_coltype.item_type`) works.
`dimensions = 1` reflects CH's type system: every Array is
one-dimensional and nested arrays (`Array(Array(String))`) are
represented via the inner item type, not via a dimension count.
Tests:
- tests/integration_tests/test_sqlalchemy/test_reflect.py:
`test_metadata_reflect_and_primary_keys` exercises the
`Dialect.get_columns` reflection path via `MetaData.reflect()`
and asserts composite primary key reflection from a MergeTree
ORDER BY clause, both via `MetaData.reflect()` and via direct
`Table(autoload_with=...)`.
End-to-end effect: `MetaData.reflect()` and
`sqlacodegen <clickhousedb+connect://...>` now produce a complete,
importable Python module with declarative ORM classes, composite
primary keys, and typed `Mapped[...]` annotations against a real
ClickHouse schema. No changes to the runtime client paths.
* Address review: drop PK reflection, set Array.as_tuple
Per review on #766, the dialect no longer reflects a primary key.
ClickHouse PRIMARY KEY / ORDER BY is a sparse index, not a uniqueness
guarantee, so get_primary_keys / get_pk_constraint return empty results
and the identity key is left for application code to declare. Removes the
is_in_primary_key query and the PK-application path in reflect_table.
Also set Array.as_tuple = False. Array bypasses ARRAY.__init__ to allow
nested arrays, but as_tuple has no class-level default and ARRAY.hashable
reads it, so select(arr).unique() raised AttributeError before.
* Add test for user-declared primary key surviving reflection
---------
Signed-off-by: Joe Spadola <[email protected]>
Co-authored-by: Joe Spadola <[email protected]>
125 lines
5.4 KiB
Python
125 lines
5.4 KiB
Python
import sqlalchemy as db
|
|
from sqlalchemy import inspect, text
|
|
from sqlalchemy.engine import Engine
|
|
|
|
from clickhouse_connect import common
|
|
from clickhouse_connect.cc_sqlalchemy.datatypes.sqltypes import Point, SimpleAggregateFunction, UInt32
|
|
|
|
|
|
def test_basic_reflection(test_engine: Engine):
|
|
common.set_setting("invalid_setting_action", "drop")
|
|
with test_engine.begin() as conn:
|
|
metadata = db.MetaData(schema="system")
|
|
table = db.Table("tables", metadata, autoload_with=test_engine)
|
|
query = db.select(table.columns.create_table_query)
|
|
result = conn.execute(query)
|
|
rows = result.fetchmany(100)
|
|
assert rows
|
|
|
|
|
|
def test_full_table_reflection(test_engine: Engine, test_db: str):
|
|
common.set_setting("invalid_setting_action", "drop")
|
|
with test_engine.begin() as conn:
|
|
conn.execute(text(f"DROP TABLE IF EXISTS {test_db}.reflect_test"))
|
|
conn.execute(
|
|
text(
|
|
f"CREATE TABLE {test_db}.reflect_test (key UInt32, value FixedString(20),"
|
|
+ "agg SimpleAggregateFunction(anyLast, String))"
|
|
+ "ENGINE AggregatingMergeTree ORDER BY key"
|
|
)
|
|
)
|
|
metadata = db.MetaData(schema=test_db)
|
|
table = db.Table("reflect_test", metadata, autoload_with=test_engine)
|
|
assert table.columns.key.type.__class__ == UInt32
|
|
assert table.columns.agg.type.__class__ == SimpleAggregateFunction
|
|
assert "MergeTree" in table.engine.name
|
|
|
|
|
|
def test_types_reflection(test_engine: Engine, test_db: str):
|
|
common.set_setting("invalid_setting_action", "drop")
|
|
with test_engine.begin() as conn:
|
|
conn.execute(text(f"DROP TABLE IF EXISTS {test_db}.sqlalchemy_types_test"))
|
|
conn.execute(text(f"CREATE TABLE {test_db}.sqlalchemy_types_test (key UInt32, pt Point) ENGINE MergeTree ORDER BY key"))
|
|
metadata = db.MetaData(schema=test_db)
|
|
table = db.Table("sqlalchemy_types_test", metadata, autoload_with=test_engine)
|
|
assert table.columns.key.type.__class__ == UInt32
|
|
assert table.columns.pt.type.__class__ == Point
|
|
assert "MergeTree" in table.engine.name
|
|
|
|
|
|
def test_table_exists(test_engine: Engine):
|
|
common.set_setting("invalid_setting_action", "drop")
|
|
inspector = inspect(test_engine)
|
|
assert inspector.has_table(table_name="columns", schema="system")
|
|
assert not inspector.has_table(table_name="nope", schema="fake_db")
|
|
|
|
|
|
def test_get_schema_names(test_engine: Engine):
|
|
common.set_setting("invalid_setting_action", "drop")
|
|
inspector = inspect(test_engine)
|
|
schema_names = inspector.get_schema_names()
|
|
assert isinstance(schema_names, list)
|
|
assert "system" in schema_names
|
|
assert "fake_db" not in schema_names
|
|
|
|
|
|
def test_get_table_names(test_engine: Engine, test_db: str):
|
|
common.set_setting("invalid_setting_action", "drop")
|
|
inspector = inspect(test_engine)
|
|
system_tables = inspector.get_table_names(schema="system")
|
|
assert isinstance(system_tables, list)
|
|
assert "columns" in system_tables
|
|
assert "fake_table" not in system_tables
|
|
|
|
|
|
def test_metadata_reflect(test_engine: Engine, test_db: str):
|
|
"""Dialect-level reflection. MetaData.reflect() exercises the
|
|
Dialect.get_multi_columns -> Dialect.get_columns path (not
|
|
Inspector.get_columns), which previously raised NotImplementedError.
|
|
The dialect does not reflect a primary key: ClickHouse PRIMARY KEY /
|
|
ORDER BY is not a uniqueness guarantee, so the identity key is left for
|
|
application code to declare explicitly."""
|
|
common.set_setting("invalid_setting_action", "drop")
|
|
with test_engine.begin() as conn:
|
|
conn.execute(text(f"DROP TABLE IF EXISTS {test_db}.reflect_pk_test"))
|
|
conn.execute(
|
|
text(
|
|
f"CREATE TABLE {test_db}.reflect_pk_test (org_id UInt32, id UInt64, payload String) ENGINE MergeTree ORDER BY (org_id, id)"
|
|
)
|
|
)
|
|
|
|
metadata = db.MetaData(schema=test_db)
|
|
metadata.reflect(bind=test_engine, only=["reflect_pk_test"])
|
|
table = metadata.tables[f"{test_db}.reflect_pk_test"]
|
|
|
|
assert {c.name for c in table.columns} == {"org_id", "id", "payload"}
|
|
assert list(table.primary_key.columns) == []
|
|
|
|
# Direct autoload should also populate columns without a reflected PK.
|
|
table2 = db.Table("reflect_pk_test", db.MetaData(schema=test_db), autoload_with=test_engine)
|
|
assert {c.name for c in table2.columns} == {"org_id", "id", "payload"}
|
|
assert list(table2.primary_key.columns) == []
|
|
|
|
|
|
def test_user_declared_primary_key(test_engine: Engine, test_db: str):
|
|
"""A user-declared primary key on a pre-declared column survives reflection."""
|
|
common.set_setting("invalid_setting_action", "drop")
|
|
with test_engine.begin() as conn:
|
|
conn.execute(text(f"DROP TABLE IF EXISTS {test_db}.reflect_user_pk_test"))
|
|
conn.execute(
|
|
text(
|
|
f"CREATE TABLE {test_db}.reflect_user_pk_test (org_id UInt32, id UInt64, payload String) "
|
|
"ENGINE MergeTree ORDER BY (org_id, id)"
|
|
)
|
|
)
|
|
|
|
table = db.Table(
|
|
"reflect_user_pk_test",
|
|
db.MetaData(schema=test_db),
|
|
db.Column("org_id", UInt32, primary_key=True),
|
|
db.Column("id", db.BigInteger, primary_key=True),
|
|
autoload_with=test_engine,
|
|
)
|
|
assert [c.name for c in table.primary_key.columns] == ["org_id", "id"]
|
|
assert {c.name for c in table.columns} == {"org_id", "id", "payload"}
|