aminghadersohi commented on code in PR #43000:
URL: https://github.com/apache/superset/pull/43000#discussion_r3979313399
##########
superset/db_engine_specs/clickhouse.py:
##########
@@ -532,6 +555,342 @@ def get_datatype(cls, type_code: str) -> str:
# keep it lowercase, as ClickHouse types aren't typical SHOUTCASE ANSI
SQL
return type_code
+ @classmethod
+ def get_columns(
+ cls,
+ inspector: Any,
+ table: Table,
+ options: dict[str, Any] | None = None,
+ ) -> list[Any]:
+ # clickhouse-connect's SQLAlchemy inspector runs reflection queries
with
+ # ``Engine.execute()``, which the SQLAlchemy 2.0-style ("future")
engine
+ # Superset builds does not implement — reflecting a table (e.g. the
+ # post-upload ``fetch_metadata`` step) then raises NotImplementedError.
+ # Rebind reflection to an explicit Connection, on which ``execute`` is
+ # supported, and defer to the base implementation from there.
+ # pylint: disable=import-outside-toplevel
+ from sqlalchemy import inspect as sqla_inspect
+ from sqlalchemy.engine import Engine
+
+ bind = inspector.bind
+ engine = bind if isinstance(bind, Engine) else bind.engine
+ with engine.connect() as connection:
+ return super().get_columns(sqla_inspect(connection), table,
options)
+
+ @classmethod
+ def _clickhouse_column_type(cls, series: Any) -> str:
+ """Map a pandas column to a concrete ClickHouse type name.
+
+ We emit clickhouse-connect's native types rather than generic
+ SQLAlchemy ones: in this dialect a generic ``Float`` becomes
+ ``Float32`` (precision loss), a generic ``DateTime`` is
second-precision
+ with a post-1970 range, and ``nullable=True`` does not produce
+ ``Nullable(...)`` at all. Every column is wrapped in ``Nullable`` so
+ missing values round-trip as NULL instead of a coerced default.
+ """
+ # pylint: disable=import-outside-toplevel
+ import pandas as pd
+
+ dtype = series.dtype
+ if pd.api.types.is_bool_dtype(dtype):
+ inner = "Bool"
+ elif pd.api.types.is_unsigned_integer_dtype(dtype):
+ # e.g. 9223372036854775808 is read as uint64 and overflows Int64.
+ inner = "UInt64"
+ elif pd.api.types.is_integer_dtype(dtype):
+ inner = "Int64"
+ elif pd.api.types.is_float_dtype(dtype):
+ inner = "Float64"
+ elif pd.api.types.is_datetime64_any_dtype(dtype):
+ inner = "DateTime64(6)"
+ elif (
+ pd.api.types.infer_dtype(series, skipna=True)
+ in cls._DATE_LIKE_INFERRED_TYPES
+ ):
+ # Object columns that actually hold date/datetime values.
+ inner = "DateTime64(6)"
+ else:
+ # Text, and anything whose exact numeric type can't be inferred
+ # (safer than silently rounding/overflowing).
+ inner = "String"
+ return f"Nullable({inner})"
+
+ @staticmethod
+ def _unwrap_clickhouse_type(ch_type: str) -> str:
+ """Strip ``Nullable(...)`` and ``LowCardinality(...)`` wrappers.
+
+ ``LowCardinality(Nullable(String))`` becomes ``String``; the wrappers
+ change how a column is stored, not what its values have to be.
+ """
+ inner = ch_type.strip()
+ while True:
+ for wrapper in ("Nullable(", "LowCardinality("):
+ if inner.startswith(wrapper) and inner.endswith(")"):
+ inner = inner[len(wrapper) : -1].strip()
+ break
+ else:
+ return inner
+
+ @classmethod
+ def _coerce_to_declared_types(cls, df: Any, column_types: dict[str, str])
-> Any:
+ """Make object columns hold what the table's column types say they
hold.
+
+ The driver's column writers take the declared type at its word: the
+ ``String`` writer calls ``encode()`` on every value and the
+ ``DateTime``/``DateTime64`` writers call ``timestamp()``. An object
+ column reaches them holding whatever pandas parsed out of the file, so
+ it can carry something other than what its column declares, and the
+ insert fails on the first value that doesn't match. Two cases arise
+ from ordinary uploads:
+
+ * A ``String`` column receiving a mix of text and numbers — an ID
column
+ with a single ``N/A`` cell. Values are rendered with ``str()``, which
+ is what the declared type promises. ``Decimal`` and ``UUID`` render
+ as their text form; ``bytes`` render as their Python repr
+ (``"b'ab'"``) rather than being decoded, since decoding could itself
+ fail on data that isn't UTF-8.
+ * A ``DateTime`` or ``DateTime64`` column receiving ``datetime.date``
+ objects, which have no ``timestamp()`` — only ``datetime.datetime``
+ does. ``to_datetime`` normalizes both to real timestamps. Only
columns
+ that already hold date/datetime objects are converted, so text is
+ never reparsed as a date.
+
+ ``column_types`` maps column names to ClickHouse type names: the DDL
+ this spec emits when it creates a table, or the server's own schema
+ when appending to an existing one. Either way the types are honoured,
+ not changed — only the values are made to fit them. Columns the frame
+ doesn't hold, and types other than the two above, are left alone.
+
+ NULLs are preserved, so they still round-trip as NULL.
+ """
+ # pylint: disable=import-outside-toplevel
+ import pandas as pd
+
+ object_columns = {
+ name: cls._unwrap_clickhouse_type(ch_type)
+ for name, ch_type in column_types.items()
+ if name in df.columns and
pd.api.types.is_object_dtype(df[name].dtype)
Review Comment:
Append skips coercion for non-object dtypes, so an all-digits zip/ID column
(int64) into a server `String` column reaches the driver unconverted and raises
`AttributeError: 'numpy.int64' object has no attribute 'encode'` (checked
against the driver's `write_str_col`).
```suggestion
if name in df.columns
and (
pd.api.types.is_object_dtype(df[name].dtype)
or cls._unwrap_clickhouse_type(ch_type) == "String"
)
```
##########
tests/unit_tests/db_engine_specs/test_clickhouse.py:
##########
@@ -765,3 +784,859 @@ def test_multivalue_contains_any_numeric_coercion_sql()
-> None:
# element type) and confirm the emitted array literal is numeric.
expr = spec.array_contains_any(column("scores"), [5, 6])
assert _compile(expr) == "hasAny(scores, array(5, 6))"
+
+
+def test_connect_supports_file_upload() -> None:
+ """
+ File upload is disabled on the legacy clickhouse-sqlalchemy spec but
+ re-enabled on the clickhouse-connect spec, whose driver can insert data.
+ """
+ from superset.db_engine_specs.clickhouse import (
+ ClickHouseConnectEngineSpec,
+ ClickHouseEngineSpec,
+ )
+
+ assert ClickHouseEngineSpec.supports_file_upload is False
+ assert ClickHouseConnectEngineSpec.supports_file_upload is True
+ assert (
+
ClickHouseConnectEngineSpec.get_public_information()["supports_file_upload"]
+ is True
+ )
+
+
+def test_connect_does_not_advertise_multivalues_insert() -> None:
+ """
+ The clickhouse-connect dialect rejects multi-values inserts, so the spec
+ must keep advertising ``supports_multivalues_insert = False``: flipping it
+ on makes ``BaseEngineSpec.df_to_sql`` pass ``method="multi"`` to pandas,
+ which raises before any row reaches ClickHouse.
+ """
+ from superset.db_engine_specs.clickhouse import ClickHouseConnectEngineSpec
+
+ assert ClickHouseConnectEngineSpec.supports_multivalues_insert is False
+
+
+def test_connect_get_columns_reflects_through_a_connection(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Reflection is rebound to an explicit ``Connection``. clickhouse-connect's
+ inspector runs its reflection queries with ``Engine.execute()``, which the
+ 2.0-style engine Superset builds does not implement, so reflecting an
+ ``Engine``-bound inspector (the post-upload ``fetch_metadata`` step) would
+ otherwise raise ``NotImplementedError``.
+ """
+ from sqlalchemy import create_engine, inspect as sqla_inspect
+ from sqlalchemy.engine import Connection
+
+ from superset.db_engine_specs.base import BaseEngineSpec
+ from superset.db_engine_specs.clickhouse import ClickHouseConnectEngineSpec
+
+ columns = [{"column_name": "a"}]
+ base_get_columns = mocker.patch.object(
+ BaseEngineSpec, "get_columns", return_value=columns
+ )
+ engine = create_engine("sqlite://")
+
+ result = ClickHouseConnectEngineSpec.get_columns(
+ sqla_inspect(engine), Table("t"), {"opt": 1}
+ )
+
+ assert result == columns
+ inspector, table, options = base_get_columns.call_args.args
+ assert isinstance(inspector.bind, Connection)
+ # The table and the caller's options are forwarded untouched.
+ assert table.table == "t"
+ assert options == {"opt": 1}
+ # The connection is opened for reflection only, and released afterwards.
+ assert inspector.bind.closed is True
+
+
+def test_connect_get_columns_accepts_a_connection_bound_inspector(
+ mocker: MockerFixture,
+) -> None:
+ """
+ An inspector already bound to a ``Connection`` reaches the engine through
+ ``bind.engine``, and the caller's own connection is left open.
+ """
+ from sqlalchemy import create_engine, inspect as sqla_inspect
+ from sqlalchemy.engine import Connection
+
+ from superset.db_engine_specs.base import BaseEngineSpec
+ from superset.db_engine_specs.clickhouse import ClickHouseConnectEngineSpec
+
+ base_get_columns = mocker.patch.object(
+ BaseEngineSpec, "get_columns", return_value=[]
+ )
+ engine = create_engine("sqlite://")
+
+ with engine.connect() as connection:
+ ClickHouseConnectEngineSpec.get_columns(sqla_inspect(connection),
Table("t"))
+
+ reflection_inspector = base_get_columns.call_args.args[0]
+ assert isinstance(reflection_inspector.bind, Connection)
+ assert reflection_inspector.bind is not connection
+ assert connection.closed is False
+
+
[email protected](
+ "series,expected",
+ [
+ (pd.Series([True, False]), "Nullable(Bool)"),
+ (pd.Series([1, 2]), "Nullable(Int64)"),
+ (pd.Series([-1, 2], dtype="int32"), "Nullable(Int64)"),
+ # 9223372036854775808 overflows Int64 and is read back as uint64.
+ (pd.Series([9223372036854775808], dtype="uint64"), "Nullable(UInt64)"),
+ (pd.Series([1.5, 2.5]), "Nullable(Float64)"),
+ # Integers with a missing value are float64 in pandas.
+ (pd.Series([1, np.nan]), "Nullable(Float64)"),
+ (
+ pd.Series(pd.to_datetime(["2021-01-01", "2021-01-02"])),
+ "Nullable(DateTime64(6))",
+ ),
+ (
+ pd.Series(pd.to_datetime(["2021-01-01T00:00:00Z"])),
+ "Nullable(DateTime64(6))",
+ ),
+ # Object columns that actually hold dates/datetimes.
+ (pd.Series([date(2021, 1, 1), date(2021, 1, 2)]),
"Nullable(DateTime64(6))"),
+ (
+ pd.Series([datetime(2021, 1, 1, 3, 4, 5), None]),
+ "Nullable(DateTime64(6))",
+ ),
+ (pd.Series(["x", "y"]), "Nullable(String)"),
+ (pd.Series([], dtype="object"), "Nullable(String)"),
+ # Mixed object columns fall back to String rather than silently
+ # rounding or overflowing.
+ (pd.Series([1, "x"]), "Nullable(String)"),
+ ],
+)
+def test_clickhouse_column_type(series: pd.Series, expected: str) -> None:
+ """
+ Pandas dtypes map to concrete ClickHouse types, every one of them wrapped
in
+ ``Nullable`` so missing values round-trip as NULL instead of a coerced
+ default.
+ """
+ from superset.db_engine_specs.clickhouse import ClickHouseConnectEngineSpec
+
+ assert ClickHouseConnectEngineSpec._clickhouse_column_type(series) ==
expected
+
+
+# Sentinel for rows that were in a table before the upload started, so a test
+# can tell "the original data survived" from "something new was written here".
+PRE_EXISTING = object()
+
+
+class FakeClickHouseClient:
+ """A stand-in for the native ``clickhouse_connect`` client.
+
+ Tracks which tables exist, what each one holds and what column types each
+ declares, so the upload paths can be asserted on their outcome -- "does the
+ user still have their data?" -- rather than only on the order of the
+ statements issued.
+
+ ``insert_df`` is as strict as the real driver's column writers, which take
+ a declared type at its word: the ``String`` writer calls ``encode()`` on
+ every value and the ``DateTime`` writers call ``timestamp()``. A frame that
+ doesn't satisfy its table's schema fails here the way it fails against a
+ server.
+ """
+
+ _CREATE = re.compile(
+ r"^CREATE TABLE (?P<name>.+?) \((?P<columns>.*)\) "
+ r"ENGINE = MergeTree ORDER BY tuple\(\)$"
+ )
+ _COLUMN = re.compile(r"^`(?P<name>(?:[^`]|``)*)` (?P<type>.+)$")
+ _DROP = re.compile(r"^DROP TABLE (?:IF EXISTS )?(?P<name>.+)$")
+ _EXISTS = re.compile(r"^EXISTS TABLE (?P<name>.+)$")
+ _EXCHANGE = re.compile(r"^EXCHANGE TABLES (?P<left>.+) AND (?P<right>.+)$")
+ _RENAME = re.compile(r"^RENAME TABLE (?P<source>.+) TO (?P<target>.+)$")
+ _DESCRIBE = re.compile(r"^DESCRIBE TABLE (?P<name>.+)$")
+ # Leading storage wrappers; the matching closing parens are at the end, so
+ # a prefix match is enough for the ``startswith``/``==`` checks below.
+ _WRAPPERS = re.compile(r"^(?:LowCardinality\(|Nullable\()+")
+
+ def __init__(self, exists: bool = False) -> None:
+ self.commands: list[str] = []
+ self.queries: list[str] = []
+ self.inserts: list[tuple[str, pd.DataFrame, Optional[str]]] = []
+ # Maps a qualified table name to the rows it holds. Tables that existed
+ # before the upload are seeded with a sentinel so they can be told
+ # apart from anything this upload created.
+ self.tables: dict[str, Any] = {"`t`": PRE_EXISTING} if exists else {}
+ # Maps a qualified table name to its column names and ClickHouse types,
+ # as DESCRIBE TABLE reports them. Filled in by CREATE, or seeded by a
+ # test to stand for a table the user created themselves.
+ self.schemas: dict[str, dict[str, str]] = {}
+ # Failures to simulate.
+ self.insert_error: Optional[Exception] = None
+ self.exchange_error: Optional[Exception] = None
+
+ @property
+ def exists(self) -> bool:
+ return "`t`" in self.tables
+
+ @exists.setter
+ def exists(self, value: bool) -> None:
+ if value:
+ self.tables["`t`"] = PRE_EXISTING
+ else:
+ self.tables.pop("`t`", None)
+
+ def command(self, sql: str) -> Any:
+ self.commands.append(sql)
+ if match := self._EXISTS.match(sql):
+ return 1 if match.group("name") in self.tables else 0
+ if match := self._CREATE.match(sql):
+ self.tables[match.group("name")] = None
+ self.schemas[match.group("name")] = self._parse_columns(
+ match.group("columns")
+ )
+ elif match := self._DROP.match(sql):
+ self.tables.pop(match.group("name"), None)
+ self.schemas.pop(match.group("name"), None)
+ elif match := self._EXCHANGE.match(sql):
+ if self.exchange_error:
+ raise self.exchange_error
+ left, right = match.group("left"), match.group("right")
+ self.tables[left], self.tables[right] = (
+ self.tables[right],
+ self.tables[left],
+ )
+ self.schemas[left], self.schemas[right] = (
+ self.schemas.get(right, {}),
+ self.schemas.get(left, {}),
+ )
+ elif match := self._RENAME.match(sql):
+ source, target = match.group("source"), match.group("target")
+ self.tables[target] = self.tables.pop(source)
+ self.schemas[target] = self.schemas.pop(source, {})
+ return None
+
+ def query(self, sql: str) -> Any:
+ self.queries.append(sql)
+ match = self._DESCRIBE.match(sql)
+ assert match, f"unexpected query: {sql}"
+ # DESCRIBE TABLE reports name, type, default_type, default_expression,
+ # comment, codec_expression and ttl_expression for each column.
+ columns = self.schemas.get(match.group("name"), {})
+ return SimpleNamespace(
+ result_rows=[
+ (name, ch_type, "", "", "", "", "") for name, ch_type in
columns.items()
+ ]
+ )
+
+ def insert_df(
+ self, table: str, df: pd.DataFrame, database: Optional[str] = None
+ ) -> None:
+ if self.insert_error:
+ raise self.insert_error
+ self._check_against_schema(table, df)
+ self.inserts.append((table, df, database))
+ self.tables[table] = df
+
+ @classmethod
+ def _parse_columns(cls, columns_ddl: str) -> dict[str, str]:
+ columns = {}
+ for column in re.split(r", (?=`)", columns_ddl):
+ match = cls._COLUMN.match(column)
+ assert match, f"unparseable column definition: {column}"
+ columns[match.group("name").replace("``", "`")] =
match.group("type")
+ return columns
+
+ def _check_against_schema(self, table: str, df: pd.DataFrame) -> None:
+ for name, ch_type in self.schemas.get(table, {}).items():
+ if name not in df.columns:
+ continue
+ # Unwrapped here rather than with the spec's own helper, so the
+ # double doesn't validate the code under test with that same code.
+ inner = self._WRAPPERS.sub("", ch_type)
+ for value in df[name]:
+ if pd.api.types.is_scalar(value) and pd.isna(value):
+ continue
+ if inner == "String" and not isinstance(value, str):
Review Comment:
`_WRAPPERS.sub` strips only the opening wrappers, so `Nullable(String)`
becomes `String)` and this `==` never matches: the double's String strictness
is inert for every wrapped column, which is the guard that would have caught
the append case above.
```suggestion
if inner.startswith("String") and not isinstance(value, str):
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
