wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
514 lines
19 KiB
Python
514 lines
19 KiB
Python
"""
|
|
CDC sync tests for partition operations.
|
|
"""
|
|
|
|
import time
|
|
import pytest
|
|
from common.common_type import CaseLabel
|
|
from .base import TestCDCSyncBase
|
|
|
|
|
|
@pytest.mark.tags(CaseLabel.CDC)
|
|
class TestCDCSyncPartition(TestCDCSyncBase):
|
|
"""Test CDC sync for partition operations."""
|
|
|
|
def setup_method(self):
|
|
"""Setup for each test method."""
|
|
self.resources_to_cleanup = []
|
|
|
|
def teardown_method(self):
|
|
"""Cleanup after each test method - only cleanup upstream, downstream will sync."""
|
|
upstream_client = getattr(self, "_upstream_client", None)
|
|
|
|
if upstream_client:
|
|
for resource_type, resource_name in self.resources_to_cleanup:
|
|
if resource_type == "collection":
|
|
self.cleanup_collection(upstream_client, resource_name)
|
|
|
|
time.sleep(1) # Allow cleanup to sync to downstream
|
|
|
|
def test_create_partition(self, upstream_client, downstream_client, sync_timeout):
|
|
"""Test CREATE_PARTITION operation sync."""
|
|
# Store upstream client for teardown
|
|
self._upstream_client = upstream_client
|
|
|
|
collection_name = self.gen_unique_name("test_col_part_create")
|
|
partition_name = self.gen_unique_name("test_part_create")
|
|
self.resources_to_cleanup.append(("collection", collection_name))
|
|
|
|
# Initial cleanup
|
|
self.cleanup_collection(upstream_client, collection_name)
|
|
|
|
# Create collection
|
|
upstream_client.create_collection(
|
|
collection_name=collection_name,
|
|
schema=self.create_default_schema(upstream_client),
|
|
)
|
|
|
|
# Wait for creation to sync
|
|
def check_create():
|
|
return downstream_client.has_collection(collection_name)
|
|
|
|
assert self.wait_for_sync(
|
|
check_create, sync_timeout, f"create collection {collection_name}"
|
|
)
|
|
|
|
# Create partition
|
|
upstream_client.create_partition(collection_name, partition_name)
|
|
|
|
# Verify partition exists in upstream
|
|
upstream_partitions = upstream_client.list_partitions(collection_name)
|
|
assert partition_name in upstream_partitions
|
|
|
|
# Wait for partition sync to downstream
|
|
def check_partition():
|
|
try:
|
|
downstream_partitions = downstream_client.list_partitions(
|
|
collection_name
|
|
)
|
|
return partition_name in downstream_partitions
|
|
except:
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_partition, sync_timeout, f"create partition {partition_name}"
|
|
)
|
|
|
|
def test_drop_partition(self, upstream_client, downstream_client, sync_timeout):
|
|
"""Test DROP_PARTITION operation sync."""
|
|
# Store upstream client for teardown
|
|
self._upstream_client = upstream_client
|
|
|
|
collection_name = self.gen_unique_name("test_col_part_drop")
|
|
partition_name = self.gen_unique_name("test_part_drop")
|
|
self.resources_to_cleanup.append(("collection", collection_name))
|
|
|
|
# Initial cleanup
|
|
self.cleanup_collection(upstream_client, collection_name)
|
|
|
|
# Create collection and partition
|
|
upstream_client.create_collection(
|
|
collection_name=collection_name,
|
|
schema=self.create_default_schema(upstream_client),
|
|
)
|
|
upstream_client.create_partition(collection_name, partition_name)
|
|
|
|
# Wait for setup to sync
|
|
def check_setup():
|
|
try:
|
|
return downstream_client.has_collection(
|
|
collection_name
|
|
) and partition_name in downstream_client.list_partitions(
|
|
collection_name
|
|
)
|
|
except:
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_setup,
|
|
sync_timeout,
|
|
f"setup collection and partition {collection_name}",
|
|
)
|
|
|
|
# Drop partition
|
|
upstream_client.drop_partition(collection_name, partition_name)
|
|
|
|
# Verify partition is dropped in upstream
|
|
upstream_partitions = upstream_client.list_partitions(collection_name)
|
|
assert partition_name not in upstream_partitions
|
|
|
|
# Wait for drop to sync to downstream
|
|
def check_drop():
|
|
try:
|
|
downstream_partitions = downstream_client.list_partitions(
|
|
collection_name
|
|
)
|
|
return partition_name not in downstream_partitions
|
|
except:
|
|
return True # If error, assume partition is dropped
|
|
|
|
assert self.wait_for_sync(
|
|
check_drop, sync_timeout, f"drop partition {partition_name}"
|
|
)
|
|
|
|
def test_load_partition(self, upstream_client, downstream_client, sync_timeout):
|
|
"""Test LOAD_PARTITION operation sync."""
|
|
# Store upstream client for teardown
|
|
self._upstream_client = upstream_client
|
|
|
|
collection_name = self.gen_unique_name("test_col_part_load")
|
|
partition_name = self.gen_unique_name("test_part_load")
|
|
self.resources_to_cleanup.append(("collection", collection_name))
|
|
|
|
# Initial cleanup
|
|
self.cleanup_collection(upstream_client, collection_name)
|
|
|
|
# Create collection, partition, and index
|
|
upstream_client.create_collection(
|
|
collection_name=collection_name,
|
|
schema=self.create_default_schema(upstream_client),
|
|
)
|
|
upstream_client.create_partition(collection_name, partition_name)
|
|
|
|
# Create index and load collection (required for querying/searching)
|
|
index_params = upstream_client.prepare_index_params()
|
|
index_params.add_index(
|
|
field_name="vector", index_type="AUTOINDEX", metric_type="L2"
|
|
)
|
|
upstream_client.create_index(collection_name, index_params)
|
|
|
|
# Wait for setup to sync
|
|
def check_setup():
|
|
try:
|
|
return downstream_client.has_collection(
|
|
collection_name
|
|
) and partition_name in downstream_client.list_partitions(
|
|
collection_name
|
|
)
|
|
except:
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_setup,
|
|
sync_timeout,
|
|
f"setup collection and partition {collection_name}",
|
|
)
|
|
|
|
# Load partition
|
|
upstream_client.load_partitions(collection_name, [partition_name])
|
|
# check partition load state in upstream
|
|
upstream_load_state = upstream_client.get_load_state(
|
|
collection_name, partition_name
|
|
)
|
|
load_state = str(upstream_load_state["state"])
|
|
print(f"DEBUG: partition load state in upstream: {load_state}")
|
|
|
|
# Wait for load to sync
|
|
def check_load():
|
|
try:
|
|
# Check partition load state
|
|
load_state = downstream_client.get_load_state(
|
|
collection_name=collection_name, partition_name=partition_name
|
|
)
|
|
print(
|
|
f"DEBUG: partition load state in check_load: {load_state['state']}"
|
|
)
|
|
return "Loaded" == str(load_state["state"])
|
|
except Exception as e:
|
|
print(f"DEBUG: get_load_state exception in check_load: {e}")
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_load, sync_timeout, f"load partition {partition_name}"
|
|
)
|
|
|
|
def test_release_partition(self, upstream_client, downstream_client, sync_timeout):
|
|
"""Test RELEASE_PARTITION operation sync."""
|
|
# Store upstream client for teardown
|
|
self._upstream_client = upstream_client
|
|
|
|
collection_name = self.gen_unique_name("test_col_part_release")
|
|
partition_name = self.gen_unique_name("test_part_release")
|
|
self.resources_to_cleanup.append(("collection", collection_name))
|
|
|
|
# Initial cleanup
|
|
self.cleanup_collection(upstream_client, collection_name)
|
|
|
|
# Create collection, partition, index, and load
|
|
upstream_client.create_collection(
|
|
collection_name=collection_name,
|
|
schema=self.create_default_schema(upstream_client),
|
|
)
|
|
upstream_client.create_partition(collection_name, partition_name)
|
|
|
|
# Create index and load collection (required for querying/searching)
|
|
index_params = upstream_client.prepare_index_params()
|
|
index_params.add_index(
|
|
field_name="vector", index_type="AUTOINDEX", metric_type="L2"
|
|
)
|
|
upstream_client.create_index(collection_name, index_params)
|
|
upstream_client.load_collection(collection_name)
|
|
upstream_client.load_partitions(collection_name, [partition_name])
|
|
for p_name in [partition_name, None]:
|
|
data = [{"vector": [0.1] * 128, "id": i} for i in range(100)]
|
|
upstream_client.insert(collection_name, data, partition_name=p_name)
|
|
|
|
# Wait for setup to sync
|
|
def check_setup():
|
|
try:
|
|
query_vector = [[0.1] * 128]
|
|
downstream_client.search(
|
|
collection_name=collection_name,
|
|
data=query_vector,
|
|
limit=1,
|
|
partition_names=[partition_name],
|
|
output_fields=[],
|
|
)
|
|
return True
|
|
except:
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_setup, sync_timeout, f"setup and load partition {partition_name}"
|
|
)
|
|
|
|
# Release partition
|
|
upstream_client.release_partitions(collection_name, [partition_name])
|
|
|
|
# check partition load state in upstream
|
|
upstream_load_state = upstream_client.get_load_state(
|
|
collection_name, partition_name
|
|
)
|
|
load_state = str(upstream_load_state["state"])
|
|
print(f"DEBUG: partition load state in upstream: {load_state}")
|
|
|
|
# check partition load state in upstream
|
|
def check_release():
|
|
try:
|
|
query_vector = [[0.1] * 128]
|
|
res = upstream_client.search(
|
|
collection_name=collection_name,
|
|
data=query_vector,
|
|
limit=1,
|
|
partition_names=[partition_name],
|
|
output_fields=[],
|
|
)
|
|
print(
|
|
f"DEBUG: released partition {partition_name} can still be searched: {res}"
|
|
)
|
|
print(
|
|
f"DEBUG: released partition {partition_name} can still be searched"
|
|
)
|
|
return False
|
|
except:
|
|
print(f"DEBUG: released partition {partition_name} cannot be searched")
|
|
return True
|
|
|
|
assert self.wait_for_sync(
|
|
check_release, sync_timeout, f"release partition {partition_name}"
|
|
)
|
|
|
|
# check partition load state in downstream
|
|
def check_release():
|
|
try:
|
|
query_vector = [[0.1] * 128]
|
|
downstream_client.search(
|
|
collection_name=collection_name,
|
|
data=query_vector,
|
|
limit=1,
|
|
partition_names=[partition_name],
|
|
output_fields=[],
|
|
)
|
|
print(
|
|
f"DEBUG: released partition {partition_name} can still be searched"
|
|
)
|
|
return False
|
|
except:
|
|
print(f"DEBUG: released partition {partition_name} cannot be searched")
|
|
return True
|
|
|
|
assert self.wait_for_sync(
|
|
check_release, sync_timeout, f"release partition {partition_name}"
|
|
)
|
|
|
|
def test_partition_insert(self, upstream_client, downstream_client, sync_timeout):
|
|
"""Test INSERT operation to partition sync."""
|
|
# Store upstream client for teardown
|
|
self._upstream_client = upstream_client
|
|
|
|
collection_name = self.gen_unique_name("test_col_part_insert")
|
|
partition_name = self.gen_unique_name("test_part_insert")
|
|
self.resources_to_cleanup.append(("collection", collection_name))
|
|
|
|
# Initial cleanup
|
|
self.cleanup_collection(upstream_client, collection_name)
|
|
|
|
# Create collection and partition
|
|
upstream_client.create_collection(
|
|
collection_name=collection_name,
|
|
schema=self.create_default_schema(upstream_client),
|
|
)
|
|
upstream_client.create_partition(collection_name, partition_name)
|
|
|
|
# Create index and load collection (required for querying/searching)
|
|
index_params = upstream_client.prepare_index_params()
|
|
index_params.add_index(
|
|
field_name="vector", index_type="AUTOINDEX", metric_type="L2"
|
|
)
|
|
upstream_client.create_index(collection_name, index_params)
|
|
upstream_client.load_collection(collection_name)
|
|
|
|
# Wait for setup to sync
|
|
def check_setup():
|
|
try:
|
|
return downstream_client.has_collection(
|
|
collection_name
|
|
) and partition_name in downstream_client.list_partitions(
|
|
collection_name
|
|
)
|
|
except:
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_setup,
|
|
sync_timeout,
|
|
f"setup collection and partition {collection_name}",
|
|
)
|
|
|
|
# Insert data to specific partition
|
|
test_data = self.generate_test_data(100)
|
|
result = upstream_client.insert(
|
|
collection_name, test_data, partition_name=partition_name
|
|
)
|
|
inserted_count = result.get("insert_count", len(test_data))
|
|
|
|
# Flush to ensure data is persisted
|
|
upstream_client.flush(collection_name)
|
|
|
|
# Wait for data sync to downstream partition by querying
|
|
def check_data():
|
|
try:
|
|
# Query data in specific partition
|
|
result = downstream_client.query(
|
|
collection_name=collection_name,
|
|
filter="",
|
|
output_fields=["count(*)"],
|
|
partition_names=[partition_name],
|
|
)
|
|
count = result[0]["count(*)"] if result else 0
|
|
return count >= inserted_count
|
|
except:
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_data, sync_timeout, f"insert data to partition {partition_name}"
|
|
)
|
|
|
|
def test_partition_delete(self, upstream_client, downstream_client, sync_timeout):
|
|
"""Test DELETE operation from partition sync."""
|
|
# Store upstream client for teardown
|
|
self._upstream_client = upstream_client
|
|
|
|
collection_name = self.gen_unique_name("test_col_part_delete")
|
|
partition_name = self.gen_unique_name("test_part_delete")
|
|
self.resources_to_cleanup.append(("collection", collection_name))
|
|
|
|
# Initial cleanup
|
|
self.cleanup_collection(upstream_client, collection_name)
|
|
|
|
# Create collection and partition
|
|
upstream_client.create_collection(
|
|
collection_name=collection_name,
|
|
schema=self.create_default_schema(upstream_client),
|
|
consistency_level="Strong",
|
|
)
|
|
upstream_client.create_partition(collection_name, partition_name)
|
|
|
|
# Create index and load collection (required for querying/searching)
|
|
index_params = upstream_client.prepare_index_params()
|
|
index_params.add_index(
|
|
field_name="vector", index_type="AUTOINDEX", metric_type="L2"
|
|
)
|
|
upstream_client.create_index(collection_name, index_params)
|
|
upstream_client.load_collection(collection_name)
|
|
|
|
# Wait for setup to sync
|
|
def check_setup():
|
|
try:
|
|
return downstream_client.has_collection(
|
|
collection_name
|
|
) and partition_name in downstream_client.list_partitions(
|
|
collection_name
|
|
)
|
|
except:
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_setup,
|
|
sync_timeout,
|
|
f"setup collection and partition {collection_name}",
|
|
)
|
|
|
|
# Insert data to partition
|
|
test_data = self.generate_test_data(100)
|
|
upstream_client.insert(
|
|
collection_name, test_data, partition_name=partition_name
|
|
)
|
|
upstream_client.flush(collection_name)
|
|
|
|
# Wait for initial data sync by querying partition
|
|
def check_data():
|
|
try:
|
|
result = downstream_client.query(
|
|
collection_name=collection_name,
|
|
filter="",
|
|
output_fields=["count(*)"],
|
|
partition_names=[partition_name],
|
|
)
|
|
count = result[0]["count(*)"] if result else 0
|
|
return count >= 100
|
|
except:
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_data, sync_timeout, f"initial data sync to partition {partition_name}"
|
|
)
|
|
|
|
# Delete some data from partition
|
|
delete_ids = list(range(20)) # Delete first 20 records
|
|
upstream_client.delete(
|
|
collection_name, filter=f"id in {delete_ids}", partition_name=partition_name
|
|
)
|
|
upstream_client.flush(collection_name)
|
|
deleted_result = upstream_client.query(
|
|
collection_name=collection_name,
|
|
filter=f"id in {delete_ids}",
|
|
output_fields=["id"],
|
|
partition_names=[partition_name],
|
|
)
|
|
total_count = upstream_client.query(
|
|
collection_name=collection_name,
|
|
filter="",
|
|
output_fields=["count(*)"],
|
|
partition_names=[partition_name],
|
|
)
|
|
total_count = total_count[0]["count(*)"] if total_count else 0
|
|
print(f"DEBUG: deleted_result in upstream: {deleted_result}")
|
|
print(f"DEBUG: total_count in upstream: {total_count}")
|
|
|
|
# Wait for delete to sync by querying partition
|
|
def check_delete():
|
|
try:
|
|
# Query for the deleted records in partition - should return empty
|
|
deleted_result = downstream_client.query(
|
|
collection_name=collection_name,
|
|
filter=f"id in {delete_ids}",
|
|
output_fields=["id"],
|
|
partition_names=[partition_name],
|
|
)
|
|
# Query total count in partition
|
|
count_result = downstream_client.query(
|
|
collection_name=collection_name,
|
|
filter="",
|
|
output_fields=["count(*)"],
|
|
partition_names=[partition_name],
|
|
)
|
|
deleted_count = len(deleted_result) if deleted_result else 0
|
|
total_count = count_result[0]["count(*)"] if count_result else 0
|
|
|
|
print(f"DEBUG: deleted_result in check_delete: {deleted_result}")
|
|
print(f"DEBUG: count_result in check_delete: {count_result}")
|
|
print(
|
|
f"DEBUG: deleted_count: {deleted_count}, total_count: {total_count}"
|
|
)
|
|
|
|
# Verify deleted records are gone and total count is correct in partition
|
|
return deleted_count == 0 and total_count == 80
|
|
except Exception as e:
|
|
print(f"DEBUG: query exception in check_delete: {e}")
|
|
return False
|
|
|
|
assert self.wait_for_sync(
|
|
check_delete, sync_timeout, f"delete data from partition {partition_name}"
|
|
)
|