| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import asyncio |
| import json |
| import math |
| import os |
| import subprocess |
| import tempfile |
| import time |
| from datetime import date, datetime, timezone |
| from datetime import time as dt_time |
| from decimal import Decimal |
| from pathlib import Path |
| |
| import pyarrow as pa |
| import pytest |
| import pytest_asyncio |
| from filelock import FileLock |
| |
| import fluss |
| |
| CLUSTER_NAME = "shared-test" |
| |
| |
| def _find_cli_binary(): |
| env_bin = os.environ.get("FLUSS_TEST_CLUSTER_BIN") |
| if env_bin: |
| if os.path.isfile(env_bin): |
| return env_bin |
| raise FileNotFoundError(f"FLUSS_TEST_CLUSTER_BIN={env_bin!r} does not exist") |
| result = subprocess.run( |
| ["cargo", "locate-project", "--workspace", "--message-format", "plain"], |
| capture_output=True, |
| text=True, |
| ) |
| if result.returncode == 0: |
| root = Path(result.stdout.strip()).parent |
| for profile in ("debug", "release"): |
| bin_path = root / "target" / profile / "fluss-test-cluster" |
| if bin_path.is_file(): |
| return str(bin_path) |
| raise FileNotFoundError( |
| "fluss-test-cluster not found. Run: cargo build -p fluss-test-cluster" |
| ) |
| |
| |
| def _start_cluster(): |
| lock = Path(tempfile.gettempdir()) / f"fluss-{CLUSTER_NAME}.lock" |
| with FileLock(lock): |
| cli = _find_cli_binary() |
| result = subprocess.run( |
| [cli, "start", "--sasl", "--name", CLUSTER_NAME], |
| capture_output=True, |
| text=True, |
| ) |
| if result.returncode != 0: |
| raise RuntimeError( |
| f"fluss-test-cluster start failed:\n{result.stderr}\n{result.stdout}" |
| ) |
| prefix = "CLUSTER_JSON: " |
| for line in result.stdout.strip().split("\n"): |
| if line.startswith(prefix): |
| info = json.loads(line[len(prefix) :]) |
| return info["bootstrap_servers"], info.get("sasl_bootstrap_servers") |
| raise RuntimeError( |
| f"No CLUSTER_JSON token in output:\n{result.stdout}\n{result.stderr}" |
| ) |
| |
| |
| def _stop_cluster(): |
| try: |
| cli = _find_cli_binary() |
| except FileNotFoundError: |
| return |
| subprocess.run([cli, "stop", "--name", CLUSTER_NAME], capture_output=True) |
| |
| |
| async def _connect(bootstrap_servers): |
| config = fluss.Config({"bootstrap.servers": bootstrap_servers}) |
| start = time.time() |
| last_err = None |
| while time.time() - start < 60: |
| try: |
| conn = await fluss.FlussConnection.create(config) |
| admin = conn.get_admin() |
| nodes = await admin.get_server_nodes() |
| if any(n.server_type == "TabletServer" for n in nodes): |
| return conn |
| await conn.close() |
| last_err = RuntimeError("No TabletServer available yet") |
| except Exception as e: |
| last_err = e |
| await asyncio.sleep(1) |
| raise RuntimeError(f"Could not connect after 60s: {last_err}") |
| |
| |
| def pytest_unconfigure(config): |
| if os.environ.get("FLUSS_BOOTSTRAP_SERVERS"): |
| return |
| if hasattr(config, "workerinput"): |
| return |
| if os.environ.get("FLUSS_SKIP_CLUSTER_TEARDOWN"): |
| return |
| _stop_cluster() |
| |
| |
| @pytest.fixture(scope="session") |
| def fluss_cluster(): |
| env = os.environ.get("FLUSS_BOOTSTRAP_SERVERS") |
| if env: |
| sasl_env = os.environ.get("FLUSS_SASL_BOOTSTRAP_SERVERS", env) |
| yield (env, sasl_env) |
| return |
| |
| plaintext_addr, sasl_addr = _start_cluster() |
| yield (plaintext_addr, sasl_addr or plaintext_addr) |
| |
| |
| @pytest_asyncio.fixture(scope="session") |
| async def connection(fluss_cluster): |
| plaintext_addr, _sasl_addr = fluss_cluster |
| conn = await _connect(plaintext_addr) |
| yield conn |
| conn.close() |
| |
| |
| @pytest.fixture(scope="session") |
| def sasl_bootstrap_servers(fluss_cluster): |
| _plaintext_addr, sasl_addr = fluss_cluster |
| return sasl_addr |
| |
| |
| @pytest.fixture(scope="session") |
| def plaintext_bootstrap_servers(fluss_cluster): |
| plaintext_addr, _sasl_addr = fluss_cluster |
| return plaintext_addr |
| |
| |
| @pytest_asyncio.fixture(scope="session") |
| async def admin(connection): |
| return connection.get_admin() |
| |
| |
| @pytest_asyncio.fixture |
| async def wait_for_table_ready(admin): |
| """ |
| Fixture that returns a helper function to wait for a table or partition to be ready. |
| """ |
| |
| async def _wait(table_path, timeout=15, partition_name=None): |
| start_time = time.monotonic() |
| while time.monotonic() - start_time < timeout: |
| try: |
| if partition_name: |
| await admin.list_partition_offsets( |
| table_path, partition_name, [0], fluss.OffsetSpec.earliest() |
| ) |
| else: |
| await admin.list_offsets( |
| table_path, [0], fluss.OffsetSpec.earliest() |
| ) |
| return |
| except (fluss.FlussError, Exception) as e: |
| # Catch "No leader found" or other errors that indicate the table/partition is still initializing |
| err_msg = str(e) |
| if any( |
| msg in err_msg |
| for msg in [ |
| "No leader found", |
| "Table not ready", |
| "Metadata not ready", |
| "not leader or follower", |
| ] |
| ): |
| await asyncio.sleep(1) |
| continue |
| raise |
| raise TimeoutError( |
| f"Table/Partition {table_path} ({partition_name or 'standard'}) " |
| f"did not become ready within {timeout}s" |
| ) |
| |
| return _wait |
| |
| |
| # Complex-type (ARRAY/MAP/ROW) helpers shared by the KV and log |
| # all_complex_datatypes tests: schema sections plus the full/edge/null row matrix. |
| # Read-back shapes: ARRAY -> list, MAP -> list of (key, value) tuples, ROW -> dict. |
| |
| |
| def pa_row_seq_label() -> pa.DataType: |
| return pa.struct([("seq", pa.int32()), ("label", pa.string())]) |
| |
| |
| def pa_row_deep() -> pa.DataType: |
| return pa.struct([("inner", pa.struct([("n", pa.int32())]))]) |
| |
| |
| def pa_row_rich() -> pa.DataType: |
| return pa.struct( |
| [ |
| ("f_bool", pa.bool_()), |
| ("f_int", pa.int32()), |
| ("f_long", pa.int64()), |
| ("f_float", pa.float32()), |
| ("f_double", pa.float64()), |
| ("f_str", pa.string()), |
| ("f_bytes", pa.binary()), |
| ("f_decimal", pa.decimal128(10, 2)), |
| ("f_date", pa.date32()), |
| ("f_time", pa.time32("ms")), |
| ("f_ts_ntz", pa.timestamp("us")), |
| ("f_ts_ltz", pa.timestamp("us", tz="UTC")), |
| ("f_binary", pa.binary(4)), |
| ("f_array_int", pa.list_(pa.int32())), |
| ] |
| ) |
| |
| |
| def array_basics_fields() -> list: |
| return [ |
| pa.field("arr_int", pa.list_(pa.int32())), |
| pa.field("arr_string", pa.list_(pa.string())), |
| pa.field("arr_of_arr", pa.list_(pa.list_(pa.int32()))), |
| pa.field("arr_of_row", pa.list_(pa_row_seq_label())), |
| ] |
| |
| |
| def row_basics_fields() -> list: |
| return [ |
| pa.field("row_basic", pa_row_seq_label()), |
| pa.field("row_deep", pa_row_deep()), |
| pa.field("row_rich", pa_row_rich()), |
| ] |
| |
| |
| def map_basics_fields() -> list: |
| return [ |
| pa.field("map_string_int", pa.map_(pa.string(), pa.int32())), |
| pa.field("map_of_row", pa.map_(pa.string(), pa_row_seq_label())), |
| pa.field("map_of_map", pa.map_(pa.string(), pa.map_(pa.string(), pa.int32()))), |
| pa.field("map_of_array", pa.map_(pa.string(), pa.list_(pa.int32()))), |
| pa.field("array_of_map", pa.list_(pa.map_(pa.string(), pa.int32()))), |
| ] |
| |
| |
| def array_rich_fields() -> list: |
| return [ |
| pa.field("arr_bytes", pa.list_(pa.binary())), |
| pa.field("arr_date", pa.list_(pa.date32())), |
| pa.field("arr_time", pa.list_(pa.time32("ms"))), |
| pa.field("arr_ts", pa.list_(pa.timestamp("us"))), |
| pa.field("arr_ts_ltz", pa.list_(pa.timestamp("us", tz="UTC"))), |
| pa.field("arr_decimal", pa.list_(pa.decimal128(10, 2))), |
| pa.field("arr_decimal_big", pa.list_(pa.decimal128(22, 5))), |
| pa.field("arr_float", pa.list_(pa.float32())), |
| pa.field("arr_double", pa.list_(pa.float64())), |
| pa.field("arr_binary", pa.list_(pa.binary(4))), |
| ] |
| |
| |
| def map_rich_fields() -> list: |
| return [ |
| pa.field("map_bytes", pa.map_(pa.string(), pa.binary())), |
| pa.field("map_decimal", pa.map_(pa.string(), pa.decimal128(10, 2))), |
| pa.field("map_date", pa.map_(pa.string(), pa.date32())), |
| pa.field("map_time", pa.map_(pa.string(), pa.time32("ms"))), |
| pa.field("map_ts", pa.map_(pa.string(), pa.timestamp("us"))), |
| pa.field("map_ts_ltz", pa.map_(pa.string(), pa.timestamp("us", tz="UTC"))), |
| pa.field("map_float", pa.map_(pa.string(), pa.float32())), |
| pa.field("map_double", pa.map_(pa.string(), pa.float64())), |
| pa.field("map_bool", pa.map_(pa.string(), pa.bool_())), |
| pa.field("map_binary", pa.map_(pa.string(), pa.binary(4))), |
| pa.field("map_int_key", pa.map_(pa.int32(), pa.string())), |
| ] |
| |
| |
| def complex_fields() -> list: |
| """`id` + all complex sections, in section order.""" |
| return ( |
| [pa.field("id", pa.int32())] |
| + array_basics_fields() |
| + array_rich_fields() |
| + row_basics_fields() |
| + map_basics_fields() |
| + map_rich_fields() |
| ) |
| |
| |
| def complex_column_names() -> list: |
| """All complex column names (everything except `id`).""" |
| return [f.name for f in complex_fields() if f.name != "id"] |
| |
| |
| def complex_schema(primary_keys=None) -> "fluss.Schema": |
| return fluss.Schema(pa.schema(complex_fields()), primary_keys=primary_keys) |
| |
| |
| _ROW_RICH_FULL = { |
| "f_bool": True, |
| "f_int": 100_000, |
| "f_long": 9_876_543_210, |
| "f_float": float("inf"), |
| "f_double": math.pi, |
| "f_str": "hello world", |
| "f_bytes": b"binary", |
| "f_decimal": Decimal("123.45"), |
| "f_date": date(2026, 1, 23), |
| "f_time": dt_time(10, 13, 47, 123000), |
| "f_ts_ntz": datetime(2026, 1, 23, 10, 13, 47, 123456), |
| "f_ts_ltz": datetime(2026, 1, 23, 10, 13, 47, 123456, tzinfo=timezone.utc), |
| "f_binary": b"\x01\x02\x03\x04", |
| "f_array_int": [7, None, 11], |
| } |
| |
| |
| def complex_full_row(id_: int) -> dict: |
| """Fully-populated row exercising every complex shape (incl. nesting).""" |
| return { |
| "id": id_, |
| "arr_int": [10, 20, 30], |
| "arr_string": ["hello", "world"], |
| "arr_of_arr": [[1, 2], [3, 4]], |
| "arr_of_row": [{"seq": 1, "label": "open"}, {"seq": 2, "label": "close"}], |
| "arr_bytes": [b"\x10\x20\x30", None], |
| "arr_date": [date(2026, 1, 23), None], |
| "arr_time": [dt_time(10, 13, 47, 123000), None], |
| "arr_ts": [datetime(2026, 1, 23, 10, 13, 47, 123456)], |
| "arr_ts_ltz": [datetime(2026, 1, 23, 10, 13, 47, 123456, tzinfo=timezone.utc)], |
| "arr_decimal": [Decimal("123.45"), None], |
| "arr_decimal_big": [ |
| Decimal("12345678901234567.12345"), |
| Decimal("-99999999999999999.99999"), |
| ], |
| "arr_float": [float("nan"), float("inf"), float("-inf")], |
| "arr_double": [float("nan"), float("inf"), float("-inf")], |
| "arr_binary": [b"\xde\xad\xbe\xef", b"\x00\x01\x02\x03"], |
| "row_basic": {"seq": 42, "label": "hello"}, |
| "row_deep": {"inner": {"n": 99}}, |
| "row_rich": dict(_ROW_RICH_FULL), |
| "map_string_int": {"a": 1, "b": None, "c": 3}, |
| "map_of_row": { |
| "e0": {"seq": 1, "label": "open"}, |
| "e1": {"seq": 2, "label": "close"}, |
| }, |
| "map_of_map": {"g1": {"a": 1, "b": 2}, "g2": {"c": 3}}, |
| "map_of_array": {"primes": [2, 3, 5], "squares": [1, 4]}, |
| "array_of_map": [{"x": 1, "y": 2}, {"z": 9}], |
| "map_bytes": {"k": b"\x10\x20\x30"}, |
| "map_decimal": {"p": Decimal("123.45")}, |
| "map_date": {"d": date(2026, 1, 23)}, |
| "map_time": {"t": dt_time(10, 13, 47, 123000)}, |
| "map_ts": {"t": datetime(2026, 1, 23, 10, 13, 47, 123456)}, |
| "map_ts_ltz": { |
| "t": datetime(2026, 1, 23, 10, 13, 47, 123456, tzinfo=timezone.utc) |
| }, |
| "map_float": {"nan": float("nan"), "inf": float("inf"), "ninf": float("-inf")}, |
| "map_double": {"nan": float("nan"), "inf": float("inf"), "ninf": float("-inf")}, |
| "map_bool": {"t": True, "f": False}, |
| "map_binary": {"k": b"\x01\x02\x03\x04"}, |
| "map_int_key": {1: "one", 2: "two"}, |
| } |
| |
| |
| # Rich sections appear only in the full row; edge/null rows leave them NULL. |
| _RICH_COLUMNS = [f.name for f in array_rich_fields() + map_rich_fields()] |
| |
| |
| def complex_edge_row(id_: int) -> dict: |
| """Edge cases: empty collections, null elements, null nested rows.""" |
| return { |
| "id": id_, |
| "arr_int": [], |
| "arr_string": [None], |
| "arr_of_arr": [[5], None, []], |
| "arr_of_row": [{"seq": 7, "label": "x"}, None, {"seq": 8, "label": "y"}], |
| "row_basic": None, |
| "row_deep": None, |
| "row_rich": None, |
| "map_string_int": {}, |
| "map_of_row": {}, |
| "map_of_map": {}, |
| "map_of_array": {}, |
| "array_of_map": [], |
| **{name: None for name in _RICH_COLUMNS}, |
| } |
| |
| |
| def complex_null_row(id_: int) -> dict: |
| """Every complex column set to NULL.""" |
| return {"id": id_, **{name: None for name in complex_column_names()}} |
| |
| |
| def _map(value) -> dict: |
| """A read-back MAP is a list of (key, value) tuples; turn it into a dict for |
| order-independent comparison (test maps use unique scalar keys).""" |
| return dict(value) |
| |
| |
| def _assert_float_triplet(values) -> None: |
| """Assert a 3-element sequence holds [NaN, +Inf, -Inf].""" |
| assert math.isnan(values[0]) |
| assert math.isinf(values[1]) and values[1] > 0 |
| assert math.isinf(values[2]) and values[2] < 0 |
| |
| |
| def assert_complex_full(row: dict) -> None: |
| assert row["arr_int"] == [10, 20, 30] |
| assert row["arr_string"] == ["hello", "world"] |
| assert row["arr_of_arr"] == [[1, 2], [3, 4]] |
| assert row["arr_of_row"] == [ |
| {"seq": 1, "label": "open"}, |
| {"seq": 2, "label": "close"}, |
| ] |
| assert row["row_basic"] == {"seq": 42, "label": "hello"} |
| assert row["row_deep"] == {"inner": {"n": 99}} |
| rr = row["row_rich"] |
| assert rr["f_bool"] is True |
| assert rr["f_int"] == 100_000 |
| assert rr["f_long"] == 9_876_543_210 |
| assert math.isinf(rr["f_float"]) and rr["f_float"] > 0 |
| assert math.isclose(rr["f_double"], math.pi, rel_tol=1e-15) |
| assert rr["f_str"] == "hello world" |
| assert rr["f_bytes"] == b"binary" |
| assert rr["f_decimal"] == Decimal("123.45") |
| assert rr["f_date"] == date(2026, 1, 23) |
| assert rr["f_time"] == dt_time(10, 13, 47, 123000) |
| assert rr["f_ts_ntz"] == datetime(2026, 1, 23, 10, 13, 47, 123456) |
| assert rr["f_ts_ltz"] == datetime( |
| 2026, 1, 23, 10, 13, 47, 123456, tzinfo=timezone.utc |
| ) |
| assert rr["f_binary"] == b"\x01\x02\x03\x04" |
| assert rr["f_array_int"] == [7, None, 11] |
| assert _map(row["map_string_int"]) == {"a": 1, "b": None, "c": 3} |
| assert _map(row["map_of_row"]) == { |
| "e0": {"seq": 1, "label": "open"}, |
| "e1": {"seq": 2, "label": "close"}, |
| } |
| assert {k: _map(v) for k, v in row["map_of_map"]} == { |
| "g1": {"a": 1, "b": 2}, |
| "g2": {"c": 3}, |
| } |
| assert _map(row["map_of_array"]) == {"primes": [2, 3, 5], "squares": [1, 4]} |
| assert [_map(m) for m in row["array_of_map"]] == [{"x": 1, "y": 2}, {"z": 9}] |
| assert row["arr_bytes"] == [b"\x10\x20\x30", None] |
| assert row["arr_date"] == [date(2026, 1, 23), None] |
| assert row["arr_time"] == [dt_time(10, 13, 47, 123000), None] |
| assert row["arr_ts"] == [datetime(2026, 1, 23, 10, 13, 47, 123456)] |
| assert row["arr_ts_ltz"] == [ |
| datetime(2026, 1, 23, 10, 13, 47, 123456, tzinfo=timezone.utc) |
| ] |
| assert row["arr_decimal"] == [Decimal("123.45"), None] |
| assert row["arr_decimal_big"] == [ |
| Decimal("12345678901234567.12345"), |
| Decimal("-99999999999999999.99999"), |
| ] |
| _assert_float_triplet(row["arr_float"]) |
| _assert_float_triplet(row["arr_double"]) |
| assert row["arr_binary"] == [b"\xde\xad\xbe\xef", b"\x00\x01\x02\x03"] |
| assert _map(row["map_bytes"]) == {"k": b"\x10\x20\x30"} |
| assert _map(row["map_decimal"]) == {"p": Decimal("123.45")} |
| assert _map(row["map_date"]) == {"d": date(2026, 1, 23)} |
| assert _map(row["map_time"]) == {"t": dt_time(10, 13, 47, 123000)} |
| assert _map(row["map_ts"]) == {"t": datetime(2026, 1, 23, 10, 13, 47, 123456)} |
| assert _map(row["map_ts_ltz"]) == { |
| "t": datetime(2026, 1, 23, 10, 13, 47, 123456, tzinfo=timezone.utc) |
| } |
| mf = _map(row["map_float"]) |
| _assert_float_triplet([mf["nan"], mf["inf"], mf["ninf"]]) |
| md = _map(row["map_double"]) |
| _assert_float_triplet([md["nan"], md["inf"], md["ninf"]]) |
| assert _map(row["map_bool"]) == {"t": True, "f": False} |
| assert _map(row["map_binary"]) == {"k": b"\x01\x02\x03\x04"} |
| assert _map(row["map_int_key"]) == {1: "one", 2: "two"} |
| |
| |
| def assert_complex_edge(row: dict) -> None: |
| assert row["arr_int"] == [] |
| assert row["arr_string"] == [None] |
| assert row["arr_of_arr"] == [[5], None, []] |
| assert row["arr_of_row"] == [ |
| {"seq": 7, "label": "x"}, |
| None, |
| {"seq": 8, "label": "y"}, |
| ] |
| assert row["row_basic"] is None |
| assert row["row_deep"] is None |
| assert row["row_rich"] is None |
| # empty maps read back as empty lists |
| assert row["map_string_int"] == [] |
| assert row["map_of_row"] == [] |
| assert row["map_of_map"] == [] |
| assert row["map_of_array"] == [] |
| assert row["array_of_map"] == [] |
| for name in _RICH_COLUMNS: |
| assert row[name] is None, f"{name} should be null in the edge row" |
