deltacat 2.0.0b11__py3-none-any.whl → 2.0.0b12__py3-none-any.whl

deltacat/tests/compute/test_compact_partition_rebase.py

@@ -1,43 +1,40 @@
-import ray
+import tempfile
 import os
-from moto import mock_s3
+from typing import Any, Callable, Dict, List, Optional, Set
 import pytest
-import boto3
-from boto3.resources.base import ServiceResource
 import pyarrow as pa
+import ray
+
 from deltacat.io.file_object_store import FileObjectStore
 from pytest_benchmark.fixture import BenchmarkFixture
-import tempfile
 
 from deltacat.tests.compute.test_util_constant import (
-    TEST_S3_RCF_BUCKET_NAME,
     DEFAULT_NUM_WORKERS,
     DEFAULT_WORKER_INSTANCE_CPUS,
 )
 from deltacat.tests.compute.test_util_common import (
-    get_rcf,
+    get_rci_from_partition,
+    read_audit_file,
+    PartitionKey,
+    get_compacted_delta_locator_from_partition,
 )
-from deltacat.tests.test_utils.utils import read_s3_contents
-from deltacat.compute.compactor.model.compactor_version import CompactorVersion
 from deltacat.tests.compute.test_util_common import (
-    get_compacted_delta_locator_from_rcf,
+    create_src_w_deltas_destination_rebase_w_deltas_strategy_main,
 )
+
+from deltacat.compute.compactor.model.compactor_version import CompactorVersion
 from deltacat.compute.compactor.model.compaction_session_audit_info import (
     CompactionSessionAuditInfo,
 )
-from deltacat.tests.compute.test_util_create_table_deltas_repo import (
-    create_src_w_deltas_destination_rebase_w_deltas_strategy,
-)
 from deltacat.tests.compute.compact_partition_rebase_test_cases import (
     REBASE_TEST_CASES,
 )
-from typing import Any, Callable, Dict, List, Optional, Set
-from deltacat.types.media import StorageType
+from deltacat.types.media import StorageType, ContentType
 from deltacat.storage import (
     DeltaLocator,
     Partition,
+    metastore,
 )
-from deltacat.types.media import ContentType
 from deltacat.compute.compactor.model.compact_partition_params import (
     CompactPartitionParams,
 )
@@ -48,6 +45,7 @@ from deltacat.utils.placement import (
     PlacementGroupManager,
 )
 
+
 """
 MODULE scoped fixtures
 """
@@ -60,29 +58,24 @@ def setup_ray_cluster():
     ray.shutdown()
 
 
-@pytest.fixture(autouse=True, scope="module")
-def mock_aws_credential():
-    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
-    os.environ["AWS_SECRET_ACCESS_ID"] = "testing"
-    os.environ["AWS_SECURITY_TOKEN"] = "testing"
-    os.environ["AWS_SESSION_TOKEN"] = "testing"
-    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
-    yield
-
+"""
+FUNCTION scoped fixtures
+"""
 
-@pytest.fixture(scope="module")
-def s3_resource(mock_aws_credential):
-    with mock_s3():
-        yield boto3.resource("s3")
 
+@pytest.fixture(autouse=True, scope="function")
+def enable_bucketing_spec_validation(monkeypatch):
+    """
+    Enable the bucketing spec validation for all tests.
+    This will help catch hash bucket drift in testing.
+    """
+    import deltacat.compute.compactor_v2.steps.merge
 
-@pytest.fixture(autouse=True, scope="module")
-def setup_compaction_artifacts_s3_bucket(s3_resource: ServiceResource):
-    s3_resource.create_bucket(
-        ACL="authenticated-read",
-        Bucket=TEST_S3_RCF_BUCKET_NAME,
+    monkeypatch.setattr(
+        deltacat.compute.compactor_v2.steps.merge,
+        "BUCKETING_SPEC_COMPLIANCE_PROFILE",
+        "ASSERT",
     )
-    yield
 
 
 @pytest.mark.parametrize(
@@ -155,14 +148,13 @@ def setup_compaction_artifacts_s3_bucket(s3_resource: ServiceResource):
     ],
     ids=[test_name for test_name in REBASE_TEST_CASES],
 )
-def test_compact_partition_rebase_same_source_and_destination(
+def test_compact_partition_rebase_same_source_and_destination_main(
     mocker,
-    s3_resource: ServiceResource,
-    local_deltacat_storage_kwargs: Dict[str, Any],
+    main_deltacat_storage_kwargs: Dict[str, Any],
     test_name: str,
     primary_keys: Set[str],
     sort_keys: List[Optional[Any]],
-    partition_keys_param: Optional[List[Any]],
+    partition_keys_param: Optional[List[PartitionKey]],
     partition_values_param: List[Optional[str]],
     input_deltas_param: List[pa.Array],
     input_deltas_delta_type: str,
@@ -181,20 +173,20 @@ def test_compact_partition_rebase_same_source_and_destination(
     compact_partition_func: Callable,
     benchmark: BenchmarkFixture,
 ):
-    import deltacat.tests.local_deltacat_storage as ds
-
-    ds_mock_kwargs = local_deltacat_storage_kwargs
+    ds_mock_kwargs = main_deltacat_storage_kwargs
     """
     This test tests the scenario where source partition locator == destination partition locator,
     but rebase source partition locator is different.
     This scenario could occur when hash bucket count changes.
+
+    This version uses the main metastore implementation instead of local storage.
     """
     partition_keys = partition_keys_param
     (
         source_table_stream,
         _,
         rebased_table_stream,
-    ) = create_src_w_deltas_destination_rebase_w_deltas_strategy(
+    ) = create_src_w_deltas_destination_rebase_w_deltas_strategy_main(
         sort_keys,
         partition_keys,
         input_deltas_param,
@@ -202,14 +194,31 @@ def test_compact_partition_rebase_same_source_and_destination(
         partition_values_param,
         ds_mock_kwargs,
     )
-    source_partition: Partition = ds.get_partition(
+
+    # Convert partition values for partition lookup (same as in the helper function)
+    converted_partition_values_for_lookup = partition_values_param
+    if partition_values_param and partition_keys:
+        converted_partition_values_for_lookup = []
+        for i, (value, pk) in enumerate(zip(partition_values_param, partition_keys)):
+            if pk.key_type.value == "int":  # Use .value to get string representation
+                converted_partition_values_for_lookup.append(int(value))
+            else:
+                converted_partition_values_for_lookup.append(value)
+
+    source_partition: Partition = metastore.get_partition(
         source_table_stream.locator,
-        partition_values_param,
+        converted_partition_values_for_lookup,
         **ds_mock_kwargs,
     )
-    rebased_partition: Partition = ds.get_partition(
+    rebased_partition: Partition = metastore.get_partition(
         rebased_table_stream.locator,
-        partition_values_param,
+        converted_partition_values_for_lookup,
+        **ds_mock_kwargs,
+    )
+    all_column_names = metastore.get_table_version_column_names(
+        rebased_table_stream.locator.table_locator.namespace,
+        rebased_table_stream.locator.table_locator.table_name,
+        rebased_table_stream.locator.table_version_locator.table_version,
         **ds_mock_kwargs,
     )
     num_workers, worker_instance_cpu = DEFAULT_NUM_WORKERS, DEFAULT_WORKER_INSTANCE_CPUS
@@ -224,10 +233,10 @@ def test_compact_partition_rebase_same_source_and_destination(
     with tempfile.TemporaryDirectory() as test_dir:
         compact_partition_params = CompactPartitionParams.of(
             {
-                "compaction_artifact_s3_bucket": TEST_S3_RCF_BUCKET_NAME,
+                "catalog": ds_mock_kwargs.get("inner"),
                 "compacted_file_content_type": ContentType.PARQUET,
                 "dd_max_parallelism_ratio": 1.0,
-                "deltacat_storage": ds,
+                "deltacat_storage": metastore,
                 "deltacat_storage_kwargs": ds_mock_kwargs,
                 "destination_partition_locator": rebased_partition.locator,
                 "hash_bucket_count": hash_bucket_count_param,
@@ -239,11 +248,11 @@ def test_compact_partition_rebase_same_source_and_destination(
                 "object_store": FileObjectStore(test_dir),
                 "pg_config": pgm,
                 "primary_keys": primary_keys,
+                "all_column_names": all_column_names,
                 "read_kwargs_provider": read_kwargs_provider_param,
                 "rebase_source_partition_locator": source_partition.locator,
                 "rebase_source_partition_high_watermark": rebased_partition.stream_position,
                 "records_per_compacted_file": records_per_compacted_file_param,
-                "s3_client_kwargs": {},
                 "source_partition_locator": rebased_partition.locator,
                 "sort_keys": sort_keys if sort_keys else None,
                 "drop_duplicates": drop_duplicates_param,
@@ -260,16 +269,14 @@ def test_compact_partition_rebase_same_source_and_destination(
         object_store_put_many_spy = mocker.spy(FileObjectStore, "put_many")
 
         # execute
-        rcf_file_s3_uri = benchmark(compact_partition_func, compact_partition_params)
+        benchmark(compact_partition_func, compact_partition_params)
 
-        round_completion_info: RoundCompletionInfo = get_rcf(
-            s3_resource, rcf_file_s3_uri
-        )
-        audit_bucket, audit_key = RoundCompletionInfo.get_audit_bucket_name_and_key(
-            round_completion_info.compaction_audit_url
+        # Get RoundCompletionInfo from the compacted partition
+        round_completion_info: RoundCompletionInfo = get_rci_from_partition(
+            rebased_partition.locator, metastore, catalog=ds_mock_kwargs.get("inner")
         )
 
-        # assert if RCF covers all files
+        # assert if RCI covers all files
         if compactor_version != CompactorVersion.V1.value:
             previous_end = None
             for start, end in round_completion_info.hb_index_to_entry_range.values():
@@ -280,8 +287,12 @@ def test_compact_partition_rebase_same_source_and_destination(
                 == round_completion_info.compacted_pyarrow_write_result.files
             )
 
-        compaction_audit_obj: Dict[str, Any] = read_s3_contents(
-            s3_resource, audit_bucket, audit_key
+        # Get catalog root for audit file resolution
+        catalog = ds_mock_kwargs.get("inner")
+        catalog_root = catalog.root
+
+        compaction_audit_obj: Dict[str, Any] = read_audit_file(
+            round_completion_info.compaction_audit_url, catalog_root
         )
         compaction_audit: CompactionSessionAuditInfo = CompactionSessionAuditInfo(
             **compaction_audit_obj
@@ -291,13 +302,17 @@ def test_compact_partition_rebase_same_source_and_destination(
         assert (
             execute_compaction_result_spy.call_args.args[-1] is False
         ), "Table version erroneously marked as in-place compacted!"
-        compacted_delta_locator: DeltaLocator = get_compacted_delta_locator_from_rcf(
-            s3_resource, rcf_file_s3_uri
+        compacted_delta_locator: DeltaLocator = (
+            get_compacted_delta_locator_from_partition(
+                rebased_partition.locator,
+                metastore,
+                catalog=ds_mock_kwargs.get("inner"),
+            )
         )
         assert (
             compacted_delta_locator.stream_position == last_stream_position_to_compact
         ), "Compacted delta locator must be equal to last stream position"
-        tables = ds.download_delta(
+        tables = metastore.download_delta(
             compacted_delta_locator, storage_type=StorageType.LOCAL, **ds_mock_kwargs
         )
         actual_rebase_compacted_table = pa.concat_tables(tables)