datastore/tests/integration/test_storage_iceberg_disks/test.py

import logging
import os

import pyspark
import pytest
import uuid

from helpers.cluster import DatastoreCluster
from helpers.s3_tools import (
    AzureUploader,
    LocalUploader,
    S3Uploader,
    LocalDownloader,
    prepare_s3_bucket,
)
from helpers.spark_tools import ResilientSparkSession, write_spark_log_config

from helpers.iceberg_utils import (
    default_upload_directory,
    write_iceberg_from_df,
    generate_data,
)

SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))


def get_spark(log_dir=None):
    builder = (
        pyspark.sql.SparkSession.builder.appName("test_storage_iceberg_disks")
        .config(
            "spark.sql.catalog.spark_catalog",
            "org.apache.iceberg.spark.SparkSessionCatalog",
        )
        .config("spark.sql.catalog.local", "org.apache.iceberg.spark.SparkCatalog")
        .config("spark.sql.catalog.spark_catalog.type", "hadoop")
        .config(
            "spark.sql.catalog.spark_catalog.warehouse",
            "/var/lib/datastore/user_files/iceberg_data",
        )
        .config(
            "spark.sql.extensions",
            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions",
        )
        .master("local")
    )

    if log_dir:
        props_path = write_spark_log_config(log_dir)
        builder = builder.config(
            "spark.driver.extraJavaOptions",
            f"-Dlog4j2.configurationFile=file:{props_path}",
        )

    return builder.getOrCreate()


def generate_cluster_def(common_path, port, azure_container):
    worker_id = os.environ.get("PYTEST_XDIST_WORKER", "")
    suffix = f"_{worker_id}" if worker_id else ""
    path = os.path.join(
        os.path.dirname(os.path.realpath(__file__)),
        f"./_gen/named_collections{suffix}.xml",
    )
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w") as f:
        f.write(
            f"""
<datastore>
    <storage_configuration>
        <disks>
            <disk_local_common>
                <type>local</type>
                <path>/var/lib/datastore/user_files/iceberg_data/default/</path>
            </disk_local_common>
            <disk_s3_common>
                <type>s3</type>
                <endpoint>http://minio1:9001/root/var/lib/datastore/user_files/iceberg_data/default/</endpoint>
                <access_key_id>minio</access_key_id>
                <secret_access_key>ClickHouse_Minio_P@ssw0rd</secret_access_key>
                <skip_access_check>true</skip_access_check>
            </disk_s3_common>
            <disk_azure_common>
                <type>object_storage</type>
                <object_storage_type>azure_blob_storage</object_storage_type>
                <storage_account_url>http://azurite1:{port}/devstoreaccount1</storage_account_url>
                <container_name>{azure_container}</container_name>
                <skip_access_check>true</skip_access_check>
                <account_name>devstoreaccount1</account_name>
                <account_key>Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==</account_key>
            </disk_azure_common>
            <disk_s3_with_cache>
                <type>cache</type>
                <disk>disk_s3_common</disk>
                <path>/tmp/s3_cache/</path>
                <max_size>1000000000</max_size>
                <skip_access_check>true</skip_access_check>
            </disk_s3_with_cache>
            <disk_azure_with_cache>
                <type>cache</type>
                <disk>disk_azure_common</disk>
                <path>/tmp/azure_cache/</path>
                <max_size>1000000000</max_size>
                <skip_access_check>true</skip_access_check>
            </disk_azure_with_cache>
        </disks>
    </storage_configuration>
    <allowed_disks_for_table_engines>disk_local_common,disk_s3_common,disk_azure_common,disk_s3_with_cache,disk_azure_with_cache</allowed_disks_for_table_engines>
</datastore>
"""
        )
    return path


@pytest.fixture(scope="module")
def started_cluster():
    try:
        cluster = DatastoreCluster(__file__)
        port = cluster.azurite_port
        user_files_path = os.path.join(
            SCRIPT_DIR, f"{cluster.instances_dir_name}/node1/database/user_files"
        )
        conf_path = generate_cluster_def(user_files_path + "/", port, "mycontainer")

        cluster.add_instance(
            "node1",
            main_configs=[conf_path, "configs/cluster.xml"],
            with_minio=True,
            with_azurite=True,
            stay_alive=True,
        )
        cluster.add_instance(
            "node2",
            main_configs=[conf_path, "configs/cluster.xml"],
            with_minio=True,
            with_azurite=True,
            stay_alive=True,
        )
        cluster.add_instance(
            "node3",
            main_configs=[conf_path, "configs/cluster.xml"],
            with_minio=True,
            with_azurite=True,
            stay_alive=True,
        )

        logging.info("Starting cluster...")
        cluster.start()

        prepare_s3_bucket(cluster)
        logging.info("S3 bucket created")

        cluster.spark_session = ResilientSparkSession(
            lambda: get_spark(cluster.instances_dir)
        )
        cluster.default_s3_uploader = S3Uploader(
            cluster.minio_client, cluster.minio_bucket
        )

        cluster.azure_container_name = "mycontainer"

        cluster.blob_service_client = cluster.blob_service_client

        container_client = cluster.blob_service_client.get_container_client(
            cluster.azure_container_name
        )
        if not container_client.exists():
            container_client.create_container()

        cluster.default_azure_uploader = AzureUploader(
            cluster.blob_service_client, cluster.azure_container_name
        )

        cluster.default_local_uploader = LocalUploader(cluster.instances["node1"])
        cluster.default_local_downloader = LocalDownloader(cluster.instances["node1"])

        yield cluster
    finally:
        cluster.shutdown()


def get_uuid_str():
    return str(uuid.uuid4()).replace("-", "_")


@pytest.mark.parametrize("format_version", ["2"])
@pytest.mark.parametrize("storage_type", ["local", "s3", "azure"])
@pytest.mark.parametrize("with_cache", [False, True])
def test_single_iceberg_file(started_cluster, format_version, storage_type, with_cache):
    if with_cache and storage_type == "local":
        pytest.skip("Cache is not supported with local disk")

    instance = started_cluster.instances["node1"]
    spark = started_cluster.spark_session
    TABLE_NAME = f"test_single_iceberg_file_{get_uuid_str()}"
    disk_name = f"disk_{storage_type}_with_cache" if with_cache else f"disk_{storage_type}_common"

    write_iceberg_from_df(spark, generate_data(spark, 0, 100), TABLE_NAME)
    default_upload_directory(
        started_cluster,
        storage_type,
        f"/iceberg_data/default/{TABLE_NAME}/",
        f"/iceberg_data/default/{TABLE_NAME}/",
    )

    suffix = f"{storage_type}_cache" if with_cache else storage_type
    table_name_1 = f"{TABLE_NAME}_{suffix}_1"
    table_name_2 = f"{TABLE_NAME}_{suffix}_2"
    table_name_3 = f"{TABLE_NAME}_{suffix}_3"
    table_name_4 = f"{TABLE_NAME}_{suffix}_4"
    table_name_5 = f"{TABLE_NAME}_{suffix}_5"

    storage_path = (
        f"{TABLE_NAME}"
        if storage_type != "azure"
        else f"var/lib/datastore/user_files/iceberg_data/default/{TABLE_NAME}"
    )
    assert "Path suffixes" in instance.query_and_get_error(
        f"CREATE TABLE {table_name_1} ENGINE=Iceberg('../', 'Parquet') SETTINGS disk = '{disk_name}'"
    )
    assert "Path suffixes" in instance.query_and_get_error(
        f"CREATE TABLE {table_name_1} ENGINE=Iceberg('/var/lib/datastore/user_files/default', 'Parquet') SETTINGS disk = '{disk_name}'"
    )

    instance.query(
        f"CREATE TABLE {table_name_2} ENGINE=Iceberg('{storage_path}', 'Parquet') SETTINGS disk = '{disk_name}'"
    )
    assert instance.query(f"SELECT * FROM {table_name_2}") == instance.query(
        "SELECT number, toString(number + 1) FROM numbers(100)"
    )

    instance.query(
        f"CREATE TABLE {table_name_3} ENGINE=Iceberg(path = '{storage_path}', format = Parquet) SETTINGS disk = '{disk_name}'"
    )
    assert instance.query(f"SELECT * FROM {table_name_3}") == instance.query(
        "SELECT number, toString(number + 1) FROM numbers(100)"
    )

    instance.query(
        f"CREATE TABLE {table_name_4} ENGINE=Iceberg(path = '{storage_path}', format = Parquet, compression_method = 'auto') SETTINGS disk = '{disk_name}'"
    )
    assert instance.query(f"SELECT * FROM {table_name_4}") == instance.query(
        "SELECT number, toString(number + 1) FROM numbers(100)"
    )

    instance.query(
        f"CREATE TABLE {table_name_5} ENGINE=Iceberg(path = '{storage_path}', format = Parquet, compression_method = 'auto') SETTINGS disk = '{disk_name}'"
    )
    assert instance.query(f"SELECT * FROM {table_name_5}") == instance.query(
        "SELECT number, toString(number + 1) FROM numbers(100)"
    )

    assert instance.query(
        f"SELECT * FROM iceberg(path = '{storage_path}', SETTINGS disk = '{disk_name}')"
    ) == instance.query("SELECT number, toString(number + 1) FROM numbers(100)")

    if storage_type == "s3":
        with pytest.raises(Exception):
            instance.query(
                f"SELECT * FROM icebergLocal(path = '{storage_path}', SETTINGS disk = '{disk_name}')"
            )
        instance.query(
            f"SELECT * FROM icebergS3(path = '{storage_path}',  SETTINGS disk = '{disk_name}')"
        )

    if storage_type == "local":
        assert instance.query(
            f"SELECT * FROM icebergLocal(path = '{storage_path}', SETTINGS disk = '{disk_name}')"
        ) == instance.query("SELECT number, toString(number + 1) FROM numbers(100)")
        # The disk-type mismatch check must still fire for the S3-flavoured function.
        assert "Disk type doesn't match" in instance.query_and_get_error(
            f"SELECT * FROM icebergS3(path = '{storage_path}', SETTINGS disk = '{disk_name}')"
        )


@pytest.mark.parametrize("storage_type", ["local", "s3"])
def test_iceberg_trivial_count_optimization(started_cluster, storage_type):
    # Regression guard for `StorageObjectStorage::supportsTrivialCountOptimization`. Iceberg's
    # `IcebergMetadata::supportsTotalRows` is true, so `StorageObjectStorage::totalRows` serves the
    # row count from manifest metadata and the planner short-circuits `SELECT count()` with a
    # `ReadFromPreparedSource (Optimized trivial count)` step. That override was nearly lost the
    # same way `supportsOptimizationToSubcolumns` was dropped in the storage unification (#59767,
    # fixed in #106443). Removing it makes the marker disappear and this test fail.
    instance = started_cluster.instances["node1"]
    spark = started_cluster.spark_session
    TABLE_NAME = f"test_iceberg_trivial_count_{get_uuid_str()}"

    write_iceberg_from_df(spark, generate_data(spark, 0, 100), TABLE_NAME)
    default_upload_directory(
        started_cluster,
        storage_type,
        f"/iceberg_data/default/{TABLE_NAME}/",
        f"/iceberg_data/default/{TABLE_NAME}/",
    )

    storage_path = (
        f"{TABLE_NAME}"
        if storage_type != "azure"
        else f"var/lib/datastore/user_files/iceberg_data/default/{TABLE_NAME}"
    )
    ch_table = f"{TABLE_NAME}_{storage_type}_count"
    instance.query(
        f"CREATE TABLE {ch_table} ENGINE=Iceberg('{storage_path}', 'Parquet') SETTINGS disk = 'disk_{storage_type}_common'"
    )

    explain = instance.query(
        f"EXPLAIN SELECT count() FROM {ch_table} SETTINGS optimize_trivial_count_query = 1"
    )
    assert "ReadFromPreparedSource (Optimized trivial count)" in explain, explain

    assert instance.query(f"SELECT count() FROM {ch_table}").strip() == "100"


@pytest.mark.parametrize("format_version", ["2"])
@pytest.mark.parametrize("storage_type", ["local", "s3", "azure"])
@pytest.mark.parametrize("with_cache", [False, True])
def test_many_tables(started_cluster, format_version, storage_type, with_cache):
    if with_cache and storage_type == "local":
        pytest.skip("Cache is not supported with local disk")

    instance = started_cluster.instances["node1"]
    TABLE_NAME = f"test_many_tables_{get_uuid_str()}"
    disk_name = f"disk_{storage_type}_with_cache" if with_cache else f"disk_{storage_type}_common"

    suffix = f"{storage_type}_cache" if with_cache else storage_type
    table_name = f"{TABLE_NAME}_{suffix}"
    table_name_2 = f"{TABLE_NAME}_{suffix}_2"

    storage_path = (
        f"{table_name}"
        if storage_type != "azure"
        else f"iceberg_data/default/{table_name}"
    )
    storage_path_2 = (
        f"{table_name_2}"
        if storage_type != "azure"
        else f"iceberg_data/default/{table_name_2}"
    )

    instance.query(
        f"CREATE TABLE {table_name} (col INT) ENGINE=Iceberg(path = '{storage_path}', format = Parquet, compression_method = 'auto') SETTINGS disk = '{disk_name}'",
        settings={"allow_insert_into_iceberg": 1},
    )
    instance.query(
        f"CREATE TABLE {table_name_2} (col INT) ENGINE=Iceberg(path = '{storage_path_2}', format = Parquet, compression_method = 'auto') SETTINGS disk = '{disk_name}'",
        settings={"allow_insert_into_iceberg": 1},
    )

    instance.query(
        f"INSERT INTO {table_name} VALUES (1);",
        settings={"allow_insert_into_iceberg": 1},
    )
    instance.query(
        f"INSERT INTO {table_name_2} VALUES (1);",
        settings={"allow_insert_into_iceberg": 1},
    )

    assert instance.query(f"SELECT * FROM {table_name}") == "1\n"
    assert instance.query(f"SELECT * FROM {table_name_2}") == "1\n"

    instance.query(f"DROP TABLE {table_name_2}")
    instance.query(f"DROP TABLE {table_name}")


@pytest.mark.parametrize("storage_type", ["s3", "azure"])
@pytest.mark.parametrize("with_cache", [False, True])
def test_cluster_table_function(started_cluster, storage_type, with_cache):

    instance = started_cluster.instances["node1"]
    spark = started_cluster.spark_session
    disk_name = f"disk_{storage_type}_with_cache" if with_cache else f"disk_{storage_type}_common"

    TABLE_NAME = f"test_iceberg_cluster_{get_uuid_str()}"

    def add_df(mode):
        write_iceberg_from_df(
            spark,
            generate_data(spark, 0, 100),
            TABLE_NAME,
            mode=mode,
        )

        files = default_upload_directory(
            started_cluster,
            storage_type,
            f"/iceberg_data/default/{TABLE_NAME}/",
            f"/iceberg_data/default/{TABLE_NAME}/",
        )

        logging.info(f"Adding another dataframe. result files: {files}")

        return files

    files = add_df(mode="overwrite")
    for i in range(1, len(started_cluster.instances)):
        files = add_df(mode="append")

    logging.info(f"Setup complete. files: {files}")
    assert len(files) == 5 + 4 * (len(started_cluster.instances) - 1)

    clusters = instance.query("SELECT * FROM system.clusters")
    logging.info(f"Clusters setup: {clusters}")

    # Regular Query only node1
    table_function_expr = (
        f"iceberg('{TABLE_NAME}', SETTINGS disk = '{disk_name}')"
        if storage_type == "s3"
        else f"iceberg('var/lib/datastore/user_files/iceberg_data/default/{TABLE_NAME}', SETTINGS disk = '{disk_name}')"
    )
    select_regular = (
        instance.query(f"SELECT * FROM {table_function_expr}").strip().split()
    )

    # Cluster Query with node1 as coordinator
    table_function_expr_cluster = (
        f"icebergCluster('cluster_simple', '{TABLE_NAME}', SETTINGS disk = '{disk_name}')"
        if storage_type == "s3"
        else f"icebergCluster('cluster_simple', 'var/lib/datastore/user_files/iceberg_data/default/{TABLE_NAME}', SETTINGS disk = '{disk_name}')"
    )
    select_cluster = (
        instance.query(f"SELECT * FROM {table_function_expr_cluster}").strip().split()
    )

    # Simple size check
    assert len(select_regular) == 600
    assert len(select_cluster) == 600

    # Actual check
    assert select_cluster == select_regular

    # Check query_log
    for replica in started_cluster.instances.values():
        replica.query("SYSTEM FLUSH LOGS")

    for node_name, replica in started_cluster.instances.items():
        cluster_secondary_queries = (
            replica.query(
                f"""
                SELECT query, type, is_initial_query, read_rows, read_bytes FROM system.query_log
                WHERE
                    type = 'QueryStart' AND
                    positionCaseInsensitive(query, '{storage_type}Cluster') != 0 AND
                    position(query, '{TABLE_NAME}') != 0 AND
                    position(query, 'system.query_log') = 0 AND
                    NOT is_initial_query
            """
            )
            .strip()
            .split("\n")
        )

        logging.info(
            f"[{node_name}] cluster_secondary_queries: {cluster_secondary_queries}"
        )
        assert len(cluster_secondary_queries) == 1

    # write 3 times
    assert (
        int(instance.query(f"SELECT count() FROM {table_function_expr_cluster}"))
        == 100 * 3
    )