import re from unittest import mock from unittest.mock import MagicMock, patch import pyarrow as pa import pytest from clickhouse_connect.driver.summary import QuerySummary from ray.data._internal.datasource.clickhouse_datasink import ( ClickHouseDatasink, ClickHouseTableSettings, SinkMode, ) from ray.data._internal.datasource.clickhouse_datasource import ClickHouseDatasource from ray.data._internal.execution.interfaces.task_context import TaskContext @pytest.fixture(autouse=True) def patch_clickhouse_get_client(): with patch("clickhouse_connect.get_client") as mock_factory: mock_instance = MagicMock() mock_instance.insert_arrow.return_value = QuerySummary({"written_rows": 3}) mock_factory.return_value = mock_instance yield mock_instance @pytest.fixture def mock_clickhouse_client(): client_mock = mock.MagicMock() client_mock.return_value = client_mock return client_mock class TestClickHouseDatasource: """Tests for ClickHouseDatasource.""" @pytest.fixture def datasource(self, mock_clickhouse_client): datasource = ClickHouseDatasource( table="default.table_name", dsn="clickhouse://user:password@localhost:8123/default", columns=["column1", "column2"], order_by=(["column1"], False), client_settings={"setting1": "value1"}, client_kwargs={"client_name": "test-client"}, ) datasource._client = mock_clickhouse_client return datasource def test_init(self, datasource): expected_query = ( "SELECT column1, column2 FROM default.table_name ORDER BY column1" ) assert datasource._query == expected_query @mock.patch.object(ClickHouseDatasource, "_init_client") def test_init_with_filter(self, mock_init_client): mock_client = MagicMock() mock_init_client.return_value = mock_client mock_client.query.return_value = MagicMock() ds_with_filter = ClickHouseDatasource( table="default.table_name", dsn="clickhouse://user:password@localhost:8123/default", columns=["column1", "column2"], filter="label = 2 AND text IS NOT NULL", order_by=(["column1"], False), ) assert ( ds_with_filter._query == "SELECT column1, column2 FROM default.table_name " "WHERE label = 2 AND text IS NOT NULL " "ORDER BY column1" ) def test_estimate_inmemory_data_size(self, datasource): mock_client = mock.MagicMock() datasource._init_client = MagicMock(return_value=mock_client) mock_client.query.return_value.result_rows = [[12345]] size = datasource.estimate_inmemory_data_size() assert size == 12345 mock_client.query.assert_called_once_with( f"SELECT SUM(byteSize(*)) AS estimate FROM ({datasource._query})" ) @pytest.mark.parametrize( "limit_row_count, offset_row_count, expected_query", [ ( 10, 0, """ SELECT column1, column2 FROM default.table_name ORDER BY column1 FETCH FIRST 10 ROWS ONLY """.strip(), ), ( 1, 0, """ SELECT column1, column2 FROM default.table_name ORDER BY column1 FETCH FIRST 1 ROW ONLY """.strip(), ), ( 10, 5, """ SELECT column1, column2 FROM default.table_name ORDER BY column1 OFFSET 5 ROWS FETCH NEXT 10 ROWS ONLY """.strip(), ), ( 1, 1, """ SELECT column1, column2 FROM default.table_name ORDER BY column1 OFFSET 1 ROW FETCH NEXT 1 ROW ONLY """.strip(), ), ], ) def test_build_block_query( self, datasource, limit_row_count, offset_row_count, expected_query ): generated_query = datasource._build_block_query( limit_row_count, offset_row_count ) clean_generated_query = re.sub(r"\s+", " ", generated_query.strip()) clean_expected_query = re.sub(r"\s+", " ", expected_query.strip()) assert clean_generated_query == clean_expected_query @pytest.mark.parametrize( "columns, expected_query_part", [ ( ["field1"], "SELECT field1 FROM default.table_name", ), (["field1", "field2"], "SELECT field1, field2 FROM default.table_name"), (None, "SELECT * FROM default.table_name"), ], ) def test_generate_query_columns(self, datasource, columns, expected_query_part): datasource._columns = columns generated_query = datasource._generate_query() assert expected_query_part in generated_query @pytest.mark.parametrize( "order_by, expected_query_part", [ ((["field1"], False), "ORDER BY field1"), ((["field2"], True), "ORDER BY field2 DESC"), ((["field1", "field2"], False), "ORDER BY (field1, field2)"), ], ) def test_generate_query_with_order_by( self, datasource, order_by, expected_query_part ): datasource._order_by = order_by generated_query = datasource._generate_query() assert expected_query_part in generated_query @mock.patch.object(ClickHouseDatasource, "_init_client") @pytest.mark.parametrize( "query_params, expected_query", [ ( {}, "SELECT * FROM default.table_name", ), ( { "columns": ["field1"], }, "SELECT field1 FROM default.table_name", ), ( { "columns": ["field1"], "order_by": (["field1"], False), }, "SELECT field1 FROM default.table_name ORDER BY field1", ), ( { "columns": ["field1", "field2"], "order_by": (["field1"], True), }, "SELECT field1, field2 FROM default.table_name ORDER BY field1 DESC", ), ( { "columns": ["field1", "field2", "field3"], "order_by": (["field1", "field2"], False), }, "SELECT field1, field2, field3 FROM default.table_name " "ORDER BY (field1, field2)", ), ( { "columns": ["field1", "field2", "field3"], "order_by": (["field1", "field2"], True), }, "SELECT field1, field2, field3 FROM default.table_name " "ORDER BY (field1, field2) DESC", ), ( { "columns": ["field1", "field2", "field3"], "order_by": (["field1", "field2", "field3"], True), }, "SELECT field1, field2, field3 FROM default.table_name " "ORDER BY (field1, field2, field3) DESC", ), ( { "columns": None, "filter": "label = 2", }, "SELECT * FROM default.table_name WHERE label = 2", ), ( { "columns": ["field1", "field2"], "filter": "label = 2 AND text IS NOT NULL", "order_by": (["field1"], False), }, "SELECT field1, field2 FROM default.table_name WHERE label = 2 AND " "text IS NOT NULL ORDER BY field1", ), ], ) def test_generate_query_full( self, mock_init_client, datasource, query_params, expected_query ): mock_client = MagicMock() mock_init_client.return_value = mock_client mock_client.query.return_value = MagicMock() datasource._columns = query_params.get("columns") datasource._filter = query_params.get("filter") datasource._order_by = query_params.get("order_by") generated_query = datasource._generate_query() assert expected_query == generated_query @pytest.mark.parametrize("parallelism", [1, 2, 3, 4]) def test_get_read_tasks_ordered_table(self, datasource, parallelism): batch1 = pa.record_batch([pa.array([1, 2, 3, 4, 5, 6, 7, 8])], names=["field1"]) batch2 = pa.record_batch( [pa.array([9, 10, 11, 12, 13, 14, 15, 16])], names=["field1"] ) mock_stream = MagicMock() mock_client = mock.MagicMock() mock_client.query_arrow_stream.return_value.__enter__.return_value = mock_stream mock_stream.__iter__.return_value = [batch1, batch2] datasource.MIN_ROWS_PER_READ_TASK = 4 datasource._init_client = MagicMock(return_value=mock_client) datasource._get_estimate_count = MagicMock(return_value=16) datasource._get_sampled_estimates = MagicMock(return_value=(100, batch1.schema)) read_tasks = datasource.get_read_tasks(parallelism) expected_num_tasks = parallelism assert len(read_tasks) == expected_num_tasks total_rows = sum(batch.num_rows for batch in [batch1, batch2]) rows_per_task = total_rows // parallelism extra_rows = total_rows % parallelism for i, read_task in enumerate(read_tasks): expected_rows = rows_per_task + (1 if i < extra_rows else 0) assert read_task.metadata.num_rows == expected_rows @pytest.mark.parametrize("parallelism", [1, 4]) def test_get_read_tasks_no_ordering(self, datasource, parallelism): datasource._order_by = None batch1 = pa.record_batch([pa.array([1, 2, 3, 4, 5, 6, 7, 8])], names=["field2"]) batch2 = pa.record_batch( [pa.array([9, 10, 11, 12, 13, 14, 15, 16])], names=["field2"] ) mock_stream = MagicMock() mock_client = mock.MagicMock() mock_client.query_arrow_stream.return_value.__enter__.return_value = mock_stream mock_stream.__iter__.return_value = [batch1, batch2] datasource.MIN_ROWS_PER_READ_TASK = 4 datasource._init_client = MagicMock(return_value=mock_client) datasource._get_estimate_count = MagicMock(return_value=16) datasource._get_sampled_estimates = MagicMock(return_value=(100, batch1.schema)) read_tasks = datasource.get_read_tasks(parallelism) assert len(read_tasks) == 1 for i, read_task in enumerate(read_tasks): assert read_task.metadata.num_rows == 16 def test_get_read_tasks_no_batches(self, datasource, mock_clickhouse_client): mock_reader = mock.MagicMock() mock_reader.__iter__.return_value = iter([]) datasource._init_client = MagicMock(return_value=mock_clickhouse_client) datasource._get_estimate_count = MagicMock(return_value=0) mock_block_accessor = mock.MagicMock() datasource._get_sampled_estimates = MagicMock(return_value=(0, None)) datasource._get_sample_block = MagicMock(return_value=mock_block_accessor) read_tasks = datasource.get_read_tasks(parallelism=2) assert len(read_tasks) == 0 @mock.patch.object(ClickHouseDatasource, "_init_client") @pytest.mark.parametrize( "filter_str, expect_error, expected_error_substring", [ ("label = 2 AND text IS NOT NULL", False, None), ("some_col = 'my;string' AND another_col > 10", False, None), ("AND label = 2", True, "Error: Simulated parse error"), ("some_col =", True, "Error: Simulated parse error"), ("col = 'someval", True, "Error: Simulated parse error"), ("col = NULL", True, "Error: Simulated parse error"), ( "col = 123; DROP TABLE foobar", True, "Invalid characters outside of string literals", ), ], ) def test_filter_validation( self, mock_init_client, filter_str, expect_error, expected_error_substring ): mock_client = MagicMock() mock_init_client.return_value = mock_client if expect_error: if "Invalid characters" not in expected_error_substring: mock_client.query.side_effect = Exception("Simulated parse error") with pytest.raises(ValueError) as exc_info: ClickHouseDatasource( table="default.table_name", dsn="clickhouse://user:password@localhost:8123/default", filter=filter_str, ) assert expected_error_substring in str(exc_info.value), ( f"Expected substring '{expected_error_substring}' " f"not found in: {exc_info.value}" ) else: mock_client.query.return_value = MagicMock() ds = ClickHouseDatasource( table="default.table_name", dsn="clickhouse://user:password@localhost:8123/default", filter=filter_str, ) assert f"WHERE {filter_str}" in ds._query @pytest.mark.parametrize("parallelism", [1, 4]) def test_get_read_tasks_with_filter(self, datasource, parallelism): datasource._filter = "label = 2 AND text IS NOT NULL" batch1 = pa.record_batch([pa.array([1, 2, 3, 4, 5, 6, 7, 8])], names=["field2"]) batch2 = pa.record_batch( [pa.array([9, 10, 11, 12, 13, 14, 15, 16])], names=["field2"] ) mock_stream = MagicMock() mock_client = mock.MagicMock() mock_client.query_arrow_stream.return_value.__enter__.return_value = mock_stream mock_stream.__iter__.return_value = [batch1, batch2] datasource.MIN_ROWS_PER_READ_TASK = 4 datasource._init_client = MagicMock(return_value=mock_client) datasource._get_estimate_count = MagicMock(return_value=16) datasource._get_sampled_estimates = MagicMock(return_value=(100, batch1.schema)) read_tasks = datasource.get_read_tasks(parallelism) assert len(read_tasks) == 1 assert read_tasks[0].metadata.num_rows == 16 def test_filter_none(self): table_name = "default.table_name" dsn = "clickhouse://user:password@localhost:8123/default" with mock.patch.object(ClickHouseDatasource, "_init_client") as mocked_init: mock_client = MagicMock() mocked_init.return_value = mock_client ds = ClickHouseDatasource(table=table_name, dsn=dsn, filter=None) assert "WHERE" not in ds._query assert ds._filter is None @pytest.fixture def mock_clickhouse_sink_client(): client = MagicMock() client.insert_arrow.return_value = QuerySummary({"written_rows": 3}) return client @pytest.fixture(autouse=True) def patch_global_get_client(mock_clickhouse_sink_client): with patch( "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client ): yield @pytest.mark.usefixtures("ray_start_2_cpus_shared") class TestClickHouseDatasink: @pytest.fixture def datasink(self, mock_clickhouse_sink_client): sink = ClickHouseDatasink( table="default.test_table", dsn="clickhouse+http://user:pass@localhost:8123/default", mode=SinkMode.APPEND, table_settings=ClickHouseTableSettings(engine="MergeTree()"), ) return sink @pytest.mark.parametrize( "mode", [ SinkMode.OVERWRITE, SinkMode.APPEND, SinkMode.CREATE, ], ) @pytest.mark.parametrize("table_exists", [True, False]) def test_on_write_start_modes( self, datasink, mock_clickhouse_sink_client, mode, table_exists ): datasink._mode = mode if (mode in [SinkMode.OVERWRITE, SinkMode.CREATE]) or ( mode == SinkMode.APPEND and not table_exists ): datasink._schema = pa.schema([("col1", pa.int32())]) with patch.object( datasink, "_table_exists", return_value=table_exists ) as mock_tbl_exists, patch.object( datasink, "_get_existing_order_by", return_value="(prev_col)" ) as mock_get_order: if mode == SinkMode.CREATE and table_exists: with pytest.raises(ValueError, match="already exists.*CREATE"): datasink.on_write_start() mock_tbl_exists.assert_called_once() mock_get_order.assert_not_called() mock_clickhouse_sink_client.command.assert_not_called() else: datasink.on_write_start() mock_tbl_exists.assert_called_once() if mode == SinkMode.OVERWRITE: drop_cmd = "DROP TABLE IF EXISTS default.test_table" mock_clickhouse_sink_client.command.assert_any_call(drop_cmd) if table_exists: mock_get_order.assert_called_once() else: mock_get_order.assert_not_called() elif mode == SinkMode.APPEND: if table_exists: mock_get_order.assert_called_once() else: mock_get_order.assert_not_called() create_cmds = [ call_args[0][0] for call_args in mock_clickhouse_sink_client.command.call_args_list if "CREATE TABLE" in call_args[0][0] ] assert ( len(create_cmds) == 1 ), "Expected one CREATE TABLE for append + !exists." elif mode == SinkMode.CREATE: if not table_exists: mock_get_order.assert_not_called() create_cmds = [ call_args[0][0] for call_args in mock_clickhouse_sink_client.command.call_args_list if "CREATE TABLE" in call_args[0][0] ] assert ( len(create_cmds) == 1 ), "Expected one CREATE TABLE for create + !exists." @pytest.mark.parametrize("mode", [SinkMode.OVERWRITE, SinkMode.APPEND]) @pytest.mark.parametrize("table_exists", [True, False]) @pytest.mark.parametrize("user_order_by", [None, "user_defined_col", "tuple()"]) def test_write_behavior( self, datasink, mock_clickhouse_sink_client, mode, table_exists, user_order_by, ): datasink._mode = mode if user_order_by is not None: datasink._table_settings.order_by = user_order_by else: datasink._table_settings.order_by = None with patch.object(datasink, "_table_exists", return_value=table_exists), patch( "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client ): if not table_exists or mode == SinkMode.OVERWRITE: datasink._schema = pa.schema([("col1", pa.int32())]) datasink.on_write_start() rb = pa.record_batch([pa.array([1, 2, 3])], names=["col1"]) block_data = pa.Table.from_batches([rb]) ctx = TaskContext(1, "") results = datasink.write([block_data], ctx=ctx) assert results == [3] mock_clickhouse_sink_client.insert_arrow.assert_called() @pytest.mark.parametrize( "schema, expected_order_by", [ (pa.schema([]), "tuple()"), (pa.schema([("ts", pa.timestamp("ns")), ("col2", pa.string())]), "ts"), (pa.schema([("col1", pa.string()), ("val", pa.int64())]), "val"), (pa.schema([("s1", pa.string()), ("s2", pa.large_string())]), "s1"), ], ) def test_pick_best_arrow_field_for_order_by( self, datasink, mock_clickhouse_sink_client, schema, expected_order_by ): datasink._mode = SinkMode.OVERWRITE datasink._table_settings.order_by = None datasink._schema = schema with patch.object(datasink, "_table_exists", return_value=False), patch( "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client ): datasink.on_write_start() # Build an empty table: 0 rows empty_table = pa.Table.from_batches([], schema=schema) datasink.write([empty_table], ctx=None) # Since we're skipping empty inserts now, we expect 0 calls: mock_clickhouse_sink_client.insert_arrow.assert_not_called() @pytest.mark.parametrize( "ddl_str, expected_order_by", [ ( "CREATE TABLE default.test_table (col1 Int32) ENGINE = MergeTree() ORDER BY col1", "col1", ), ("CREATE TABLE default.test_table (col1 Int32) ENGINE = MergeTree()", None), ( "CREATE TABLE default.test_table (col1 Int32) ORDER BY city ENGINE = MergeTree()", "city", ), ( "CREATE TABLE default.test_table (col1 Int32) ENGINE = MergeTree() PARTITION BY toYYYYMMDD(date_col)", None, ), ], ) def test_get_existing_order_by( self, datasink, mock_clickhouse_sink_client, ddl_str, expected_order_by ): mock_clickhouse_sink_client.command.return_value = ddl_str result = datasink._get_existing_order_by(mock_clickhouse_sink_client) assert result == expected_order_by @pytest.mark.parametrize( "table_settings, schema, expected_engine, expected_order_by_part, expected_clauses", [ ( ClickHouseTableSettings(), pa.schema([("col1", pa.int32())]), "MergeTree()", "ORDER BY col1", [], ), ( ClickHouseTableSettings(engine="ReplacingMergeTree()"), pa.schema([("col1", pa.int32())]), "ReplacingMergeTree()", "ORDER BY col1", [], ), ( ClickHouseTableSettings(order_by="user_col"), pa.schema([("col1", pa.int32())]), "MergeTree()", "ORDER BY user_col", [], ), ( ClickHouseTableSettings(partition_by="toYYYYMMDD(ts)"), pa.schema([("ts", pa.timestamp("ns"))]), "MergeTree()", "ORDER BY ts", ["PARTITION BY toYYYYMMDD(ts)"], ), ( ClickHouseTableSettings(primary_key="id"), pa.schema([("id", pa.int64()), ("val", pa.string())]), "MergeTree()", "ORDER BY id", ["PRIMARY KEY (id)"], ), ( ClickHouseTableSettings(settings="index_granularity=8192"), pa.schema([("id", pa.int64())]), "MergeTree()", "ORDER BY id", ["SETTINGS index_granularity=8192"], ), ( ClickHouseTableSettings( engine="SummingMergeTree()", order_by="col2", partition_by="toYYYYMMDD(ts)", primary_key="id", settings="index_granularity=8192", ), pa.schema( [ ("id", pa.int64()), ("col2", pa.float64()), ("ts", pa.timestamp("ns")), ] ), "SummingMergeTree()", "ORDER BY col2", [ "PARTITION BY toYYYYMMDD(ts)", "PRIMARY KEY (id)", "SETTINGS index_granularity=8192", ], ), ], ) def test_generate_create_table_sql( self, datasink, mock_clickhouse_sink_client, table_settings, schema, expected_engine, expected_order_by_part, expected_clauses, ): datasink._mode = SinkMode.OVERWRITE datasink._table_settings = table_settings datasink._schema = schema with patch.object(datasink, "_table_exists", return_value=False), patch( "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client ): datasink.on_write_start() arrays = [] for field in schema: if pa.types.is_integer(field.type): arrays.append(pa.array([1, 2, 3], type=field.type)) elif pa.types.is_floating(field.type): arrays.append(pa.array([1.1, 2.2, 3.3], type=field.type)) elif pa.types.is_timestamp(field.type): arrays.append(pa.array([1, 2, 3], type=field.type)) else: arrays.append(pa.array(["a", "b", "c"], type=field.type)) block_data = pa.Table.from_arrays(arrays, names=[f.name for f in schema]) datasink.write([block_data], ctx=TaskContext(1, "")) create_sql = None for call_arg in mock_clickhouse_sink_client.command.call_args_list: sql_arg = call_arg[0][0] if "CREATE TABLE" in sql_arg: create_sql = sql_arg break assert create_sql is not None, "No CREATE TABLE statement was generated!" assert f"ENGINE = {expected_engine}" in create_sql assert expected_order_by_part in create_sql for clause in expected_clauses: assert clause in create_sql @pytest.mark.parametrize( "provided_schema,block_fields,expected_create_columns", [ ( pa.schema([("my_col", pa.float64()), ("ts", pa.timestamp("ns"))]), [("my_col", pa.int32()), ("ts", pa.int64())], ["`my_col` Float64", "`ts` DateTime64(3)"], ), ( pa.schema([("my_col", pa.float64()), ("col2", pa.string())]), [("my_col", pa.int64()), ("col2", pa.large_string())], [ "`my_col` Float64", "`col2` String", ], ), ( pa.schema([("id", pa.int32()), ("val", pa.string())]), [("id", pa.int64()), ("val", pa.large_string())], [ "`id` Int32", "`val` String", ], ), ( pa.schema([("f1", pa.int32()), ("f2", pa.float64())]), [("f1", pa.int32()), ("f2", pa.int32())], [ "`f1` Int32", "`f2` Float64", ], ), ], ) def test_write_schema_override( self, datasink, mock_clickhouse_sink_client, provided_schema, block_fields, expected_create_columns, ): datasink._mode = SinkMode.CREATE datasink._table_settings.order_by = None with patch.object(datasink, "_table_exists", return_value=False), patch( "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client ): datasink._schema = provided_schema datasink.on_write_start() arrays = [] for name, typ in block_fields: if pa.types.is_integer(typ): arrays.append(pa.array([1, 2, 3], type=typ)) elif pa.types.is_string(typ) or pa.types.is_large_string(typ): arrays.append(pa.array(["a", "b", "c"], type=typ)) elif pa.types.is_timestamp(typ): arrays.append(pa.array([1, 2, 3], type=typ)) else: arrays.append(pa.array([1.0, 2.0, 3.0], type=typ)) block_data = pa.Table.from_arrays( arrays, names=[n for (n, _) in block_fields] ) datasink.write([block_data], ctx=TaskContext(1, "")) create_sql = None for call_arg in mock_clickhouse_sink_client.command.call_args_list: sql_arg = call_arg[0][0] if "CREATE TABLE" in sql_arg: create_sql = sql_arg break assert create_sql is not None, "Expected CREATE TABLE to be issued." for expected_col_def in expected_create_columns: assert expected_col_def in create_sql @pytest.mark.parametrize( "max_insert_block_rows,block_sizes,expected_insert_calls", [ (2, [6], [3]), (2, [6, 3], [3, 2]), (None, [6, 3], [1, 1]), (3, [3, 5, 2], [1, 2, 1]), ], ) def test_chunked_inserts( self, datasink, mock_clickhouse_sink_client, max_insert_block_rows, block_sizes, expected_insert_calls, ): datasink._mode = SinkMode.CREATE datasink._schema = pa.schema([("col1", pa.int32())]) datasink._max_insert_block_rows = max_insert_block_rows with patch.object(datasink, "_table_exists", return_value=False), patch( "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client ): datasink.on_write_start() blocks = [] for size in block_sizes: arr = pa.array(range(size), type=pa.int32()) block_table = pa.Table.from_arrays([arr], names=["col1"]) blocks.append(block_table) datasink.write(blocks, ctx=TaskContext(1, "")) insert_calls = [ call_args[0][1] for call_args in mock_clickhouse_sink_client.insert_arrow.call_args_list ] actual_inserts = len(insert_calls) assert actual_inserts == sum(expected_insert_calls), ( f"Expected total insert calls {sum(expected_insert_calls)}, " f"got {actual_inserts}." ) offset = 0 for block_idx, size in enumerate(block_sizes): calls_for_block = expected_insert_calls[block_idx] chunk_tables = insert_calls[offset : offset + calls_for_block] offset += calls_for_block total_rows = sum(tbl.num_rows for tbl in chunk_tables) assert total_rows == size, ( f"Block of size {size} was split incorrectly. " f"Sum of chunk sizes is {total_rows}." ) @pytest.mark.parametrize( "table_exists,mode,user_schema,block_fields,expected_error_regex", [ ( False, SinkMode.CREATE, pa.schema([("id", pa.int32())]), [("id", pa.int32()), ("extra_col", pa.int32())], r"(ArrowInvalid|Could not convert|field names are not matching|columns not in target schema.*)", ), ( True, SinkMode.OVERWRITE, pa.schema([("id", pa.timestamp("ns"))]), [("id", pa.int32())], r"(ArrowInvalid|Could not convert|field names are not matching|columns not in target schema|Unsupported cast.*)", ), ], ) def test_user_schema_block_mismatch( self, datasink, mock_clickhouse_sink_client, table_exists, mode, user_schema, block_fields, expected_error_regex, ): datasink._mode = mode datasink._schema = user_schema with patch.object(datasink, "_table_exists", return_value=table_exists), patch( "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client ): try: datasink.on_write_start() except ValueError: pass arrays = [] for name, typ in block_fields: arrays.append(pa.array([1, 2, 3], type=typ)) block_data = pa.Table.from_arrays( arrays, names=[n for (n, _) in block_fields] ) with pytest.raises( (ValueError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError), match=expected_error_regex, ): datasink.write([block_data], ctx=TaskContext(1, "")) if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", __file__]))