mirror of
https://github.com/qdrant/qdrant.git
synced 2026-07-30 14:41:00 -05:00
5936926c88
* Refactor bash test for snapshots recovery * Address review: minio check * Address review
360 lines
15 KiB
Python
360 lines
15 KiB
Python
import hashlib
|
|
import io
|
|
import random
|
|
import requests
|
|
import time
|
|
import uuid
|
|
from typing import Optional, Dict, Any, List, Generator, Tuple
|
|
|
|
from qdrant_client import QdrantClient, models
|
|
|
|
VECTOR_SIZE = 4
|
|
|
|
|
|
class ClientUtils:
|
|
"""Utility class for Qdrant operations using qdrant-client."""
|
|
|
|
def __init__(self, host: str = "localhost", port: int = 6333, timeout: int = 60):
|
|
"""Initialize the ClientUtils with a QdrantClient instance."""
|
|
self.host = host
|
|
self.port = port
|
|
self.timeout = timeout
|
|
self.client = QdrantClient(host=host, port=port, timeout=timeout)
|
|
|
|
def wait_for_server(self, timeout: int = 30) -> bool:
|
|
"""Wait for Qdrant server to be ready."""
|
|
start_time = time.time()
|
|
while time.time() - start_time < timeout:
|
|
try:
|
|
info = self.client.get_collections()
|
|
if info is not None:
|
|
print("Server ready to serve traffic")
|
|
return True
|
|
except Exception:
|
|
pass
|
|
|
|
print("Waiting for server to start...")
|
|
time.sleep(1)
|
|
return False
|
|
|
|
def wait_for_collection_loaded(self, collection_name: str, timeout: int = 10) -> bool:
|
|
"""Wait for a specific collection to be loaded."""
|
|
for _ in range(timeout):
|
|
try:
|
|
collections_response = self.client.get_collections()
|
|
if collections_response and collections_response.collections:
|
|
collection_names = [col.name for col in collections_response.collections]
|
|
if collection_name in collection_names:
|
|
return True
|
|
except Exception:
|
|
pass
|
|
time.sleep(1)
|
|
return False
|
|
|
|
def list_collections_names(self) -> List[str]:
|
|
"""List all collections names."""
|
|
try:
|
|
collections_response = self.client.get_collections()
|
|
if collections_response and collections_response.collections:
|
|
return [col.name for col in collections_response.collections]
|
|
return []
|
|
except Exception as e:
|
|
raise Exception(f"Failed to list collections: {e}") from e
|
|
|
|
def get_collection_info_dict(self, collection_name: str) -> Dict[str, Any]:
|
|
"""Get detailed information about a specific collection."""
|
|
try:
|
|
collection_info = self.client.get_collection(collection_name)
|
|
return {
|
|
"result": {
|
|
"status": collection_info.status.value if hasattr(collection_info.status, 'value') else str(collection_info.status),
|
|
"vectors_count": collection_info.vectors_count,
|
|
"points_count": collection_info.points_count,
|
|
"config": collection_info.config,
|
|
"indexed_vectors_count": collection_info.indexed_vectors_count,
|
|
"payload_schema": collection_info.payload_schema
|
|
},
|
|
"status": "ok"
|
|
}
|
|
except Exception as e:
|
|
raise Exception(f"Failed to get collection info for '{collection_name}': {e}") from e
|
|
|
|
@staticmethod
|
|
def generate_points(amount: int, vector_size: int = VECTOR_SIZE) -> Generator[Dict[str, List[models.PointStruct]], None, None]:
|
|
"""Generate batches of points for insertion."""
|
|
for _ in range(amount):
|
|
points = []
|
|
for _ in range(100):
|
|
points.append(
|
|
models.PointStruct(
|
|
id=str(uuid.uuid4()),
|
|
vector=[round(random.uniform(0, 1), 2) for _ in range(vector_size)],
|
|
payload={"city": ["Berlin", "London"]}
|
|
)
|
|
)
|
|
yield {"points": points}
|
|
|
|
def create_collection(self, collection_name: str, collection_config: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Create a collection with the given configuration."""
|
|
vectors_config = collection_config.get("vectors", {})
|
|
|
|
# Handle vector configuration
|
|
if isinstance(vectors_config, dict) and "size" in vectors_config:
|
|
vector_params = models.VectorParams(
|
|
size=vectors_config["size"],
|
|
distance=models.Distance(vectors_config.get("distance", "Cosine"))
|
|
)
|
|
else:
|
|
vector_params = models.VectorParams(size=VECTOR_SIZE, distance=models.Distance.COSINE)
|
|
|
|
# Create collection
|
|
result = self.client.create_collection(
|
|
collection_name=collection_name,
|
|
vectors_config=vector_params,
|
|
shard_number=collection_config.get("shard_number"),
|
|
replication_factor=collection_config.get("replication_factor"),
|
|
write_consistency_factor=collection_config.get("write_consistency_factor"),
|
|
on_disk_payload=collection_config.get("on_disk_payload"),
|
|
hnsw_config=collection_config.get("hnsw_config"),
|
|
optimizers_config=collection_config.get("optimizers_config"),
|
|
wal_config=collection_config.get("wal_config"),
|
|
quantization_config=collection_config.get("quantization_config"),
|
|
init_from=collection_config.get("init_from"),
|
|
timeout=self.timeout
|
|
)
|
|
|
|
return {"result": result, "status": "ok"}
|
|
|
|
def update_collection(self, collection_name: str, collection_params: Dict[str, Any]) -> None:
|
|
"""Update collection parameters."""
|
|
try:
|
|
self.client.update_collection(
|
|
collection_name=collection_name,
|
|
optimizers_config=collection_params.get("optimizers_config"),
|
|
collection_params=collection_params.get("collection_params"),
|
|
vectors_config=collection_params.get("vectors_config"),
|
|
hnsw_config=collection_params.get("hnsw_config"),
|
|
quantization_config=collection_params.get("quantization_config"),
|
|
timeout=self.timeout
|
|
)
|
|
except Exception as e:
|
|
print(f"Collection patching failed with error: {e}")
|
|
raise RuntimeError(f"Collection patching failed: {e}") from e
|
|
|
|
def wait_for_status(self, collection_name: str, status: str) -> str:
|
|
"""Wait for collection to reach the specified status."""
|
|
for _ in range(30):
|
|
try:
|
|
collection_info = self.client.get_collection(collection_name)
|
|
curr_status = collection_info.status.value if hasattr(collection_info.status, 'value') else str(collection_info.status)
|
|
|
|
if curr_status.lower() == status.lower():
|
|
print(f"Status {status}: OK")
|
|
return "ok"
|
|
|
|
time.sleep(1)
|
|
print(f"Wait for status {status}")
|
|
except Exception as e:
|
|
print(f"Collection info fetching failed with error: {e}")
|
|
raise RuntimeError(f"Collection info fetching failed: {e}") from e
|
|
|
|
print(f"After 30s status is not {status}. Stop waiting.")
|
|
return "timeout"
|
|
|
|
def insert_points(self, collection_name: str, batch_data: Dict[str, Any], quit_on_ood: bool = False, wait: bool = True) -> Optional[str]:
|
|
"""Insert points into the collection."""
|
|
try:
|
|
# Convert dict format to PointStruct if needed
|
|
points = batch_data.get("points", [])
|
|
if points and isinstance(points[0], dict):
|
|
points = [
|
|
models.PointStruct(
|
|
id=p["id"],
|
|
vector=p["vector"],
|
|
payload=p.get("payload")
|
|
) for p in points
|
|
]
|
|
|
|
self.client.upsert(
|
|
collection_name=collection_name,
|
|
points=points,
|
|
wait=wait
|
|
)
|
|
|
|
except Exception as e:
|
|
expected_error_message = "No space left on device"
|
|
if expected_error_message in str(e):
|
|
if quit_on_ood:
|
|
return "ood"
|
|
else:
|
|
print(f"Points insertions failed with error: {e}")
|
|
raise RuntimeError(f"Points insertion failed: {e}") from e
|
|
|
|
def search_points(self, collection_name: str) -> Dict[str, Any]:
|
|
"""Search for points in the collection using the modern query_points API."""
|
|
try:
|
|
results = self.client.query_points(
|
|
collection_name=collection_name,
|
|
query=[round(random.uniform(0, 1), 2) for _ in range(VECTOR_SIZE)],
|
|
limit=10,
|
|
query_filter=models.Filter(
|
|
must=[
|
|
models.FieldCondition(
|
|
key="city",
|
|
match=models.MatchValue(value="Berlin")
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
# Convert results to expected format
|
|
return {
|
|
"result": [
|
|
{
|
|
"id": hit.id,
|
|
"score": hit.score,
|
|
"payload": hit.payload,
|
|
"vector": hit.vector
|
|
} for hit in results.points
|
|
],
|
|
"status": "ok"
|
|
}
|
|
|
|
except Exception as e:
|
|
raise RuntimeError("Search failed") from e
|
|
|
|
def create_snapshot(self, collection_name: str = "test_collection") -> Optional[str]:
|
|
"""Create a snapshot of the collection."""
|
|
snapshot_info = self.client.create_snapshot(collection_name=collection_name)
|
|
return snapshot_info.name if snapshot_info else None
|
|
|
|
def download_snapshot(self, collection_name: str, snapshot_name: str) -> Tuple[bytes, str]:
|
|
"""Download a snapshot and return its content and checksum."""
|
|
# Note: qdrant-client doesn't have a direct method to download snapshot content
|
|
# This would need to be implemented using the REST API directly
|
|
snapshot_url = f"http://{self.host}:{self.port}/collections/{collection_name}/snapshots/{snapshot_name}"
|
|
response = requests.get(snapshot_url)
|
|
response.raise_for_status()
|
|
|
|
content = response.content
|
|
checksum = hashlib.sha256(content).hexdigest()
|
|
return content, checksum
|
|
|
|
def recover_snapshot_from_url(self, collection_name: str, snapshot_url: str, checksum: Optional[str] = None) -> Dict[str, Any]:
|
|
"""Recover a collection from a snapshot URL."""
|
|
# Note: qdrant-client doesn't have a direct method for URL recovery
|
|
# This would need to be implemented using the REST API directly
|
|
body = {
|
|
"location": snapshot_url,
|
|
"wait": "true"
|
|
}
|
|
if checksum:
|
|
body["checksum"] = checksum
|
|
|
|
response = requests.put(
|
|
f"http://{self.host}:{self.port}/collections/{collection_name}/snapshots/recover",
|
|
json=body
|
|
)
|
|
if not response.ok:
|
|
print(f"Recovery failed with status {response.status_code}: {response.text}")
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
def upload_snapshot_file(self, collection_name: str, snapshot_content: bytes) -> Dict[str, Any]:
|
|
"""Upload a snapshot file directly."""
|
|
# Note: qdrant-client doesn't have a direct method for snapshot upload
|
|
# This would need to be implemented using the REST API directly
|
|
files = {
|
|
'snapshot': ('snapshot.tar', io.BytesIO(snapshot_content), 'application/octet-stream')
|
|
}
|
|
|
|
response = requests.post(
|
|
f"http://{self.host}:{self.port}/collections/{collection_name}/snapshots/upload",
|
|
files=files
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
def create_shard_snapshot(self, collection_name: str, shard_id: int = 0) -> str:
|
|
"""Create a snapshot of a specific shard."""
|
|
snapshot_info = self.client.create_shard_snapshot(
|
|
collection_name=collection_name,
|
|
shard_id=shard_id
|
|
)
|
|
return snapshot_info.name
|
|
|
|
def download_shard_snapshot(self, collection_name: str, shard_id: int, snapshot_name: str) -> bytes:
|
|
"""Download a shard snapshot and return its content."""
|
|
# Note: qdrant-client doesn't have a direct method to download shard snapshot content
|
|
# This would need to be implemented using the REST API directly
|
|
snapshot_url = f"http://{self.host}:{self.port}/collections/{collection_name}/shards/{shard_id}/snapshots/{snapshot_name}"
|
|
response = requests.get(snapshot_url)
|
|
response.raise_for_status()
|
|
return response.content
|
|
|
|
def recover_shard_snapshot_from_url(self, collection_name: str, shard_id: int, snapshot_url: str) -> Dict[str, Any]:
|
|
"""Recover a shard from a snapshot URL."""
|
|
result = self.client.recover_shard_snapshot(
|
|
collection_name=collection_name,
|
|
shard_id=shard_id,
|
|
location=snapshot_url,
|
|
wait=True
|
|
)
|
|
return {"result": result, "status": "ok"}
|
|
|
|
def upload_shard_snapshot_file(self, collection_name: str, shard_id: int, snapshot_content: bytes) -> Dict[str, Any]:
|
|
"""Upload a shard snapshot file directly."""
|
|
# Note: qdrant-client doesn't have a direct method for shard snapshot upload
|
|
# This would need to be implemented using the REST API directly
|
|
files = {
|
|
'snapshot': ('shard_snapshot.tar', io.BytesIO(snapshot_content), 'application/octet-stream')
|
|
}
|
|
|
|
response = requests.post(
|
|
f"http://{self.host}:{self.port}/collections/{collection_name}/shards/{shard_id}/snapshots/upload",
|
|
files=files
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
def verify_collection_exists(self, collection_name: str) -> Dict[str, Any]:
|
|
"""Verify that a collection exists and is accessible."""
|
|
collection_info = self.client.get_collection(collection_name)
|
|
return {
|
|
"result": {
|
|
"status": collection_info.status,
|
|
"points_count": collection_info.points_count,
|
|
"config": collection_info.config
|
|
},
|
|
"status": "ok"
|
|
}
|
|
|
|
def create_payload_index(self, collection_name: str, field_name: str,
|
|
field_schema: Any, wait: bool = True) -> bool:
|
|
"""
|
|
Create a payload index for a field in the collection.
|
|
|
|
Args:
|
|
collection_name: Name of the collection
|
|
field_name: Name of the field to index
|
|
field_schema: Schema for the field - can be:
|
|
- models.PayloadSchemaType.UUID for UUID fields
|
|
- models.PayloadSchemaType.KEYWORD for keyword fields
|
|
- models.TextIndexParams for text fields
|
|
- models.KeywordIndexParams for keyword fields with options
|
|
wait: Whether to wait for the operation to complete
|
|
|
|
Returns:
|
|
True if the index was created successfully
|
|
"""
|
|
try:
|
|
self.client.create_payload_index(
|
|
collection_name=collection_name,
|
|
field_name=field_name,
|
|
field_schema=field_schema,
|
|
wait=wait
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
raise Exception(f"Failed to create payload index for field '{field_name}': {e}") from e
|