import pathlib from .fixtures import create_collection, upsert_random_points, random_dense_vector, search, random_sparse_vector from .utils import * N_PEERS = 3 N_SHARDS = 3 COLLECTION_NAME = "test_collection" def create_snapshot(peer_api_uri): r = requests.post(f"{peer_api_uri}/collections/{COLLECTION_NAME}/snapshots") assert_http_ok(r) return r.json()["result"]["name"] def create_payload_index(peer_api_uri, collection_name, field_name, field_schema): r = requests.put(f"{peer_api_uri}/collections/{COLLECTION_NAME}/index", json={ "field_name": field_name, "field_schema": field_schema, }) assert_http_ok(r) return r.json()["result"] def get_peer_id(peer_api_uri): r = requests.get(f"{peer_api_uri}/cluster") assert_http_ok(r) return r.json()["result"]["peer_id"] def get_local_shards(peer_api_uri): r = requests.get(f"{peer_api_uri}/collections/{COLLECTION_NAME}/cluster") assert_http_ok(r) return r.json()["result"]['local_shards'] def get_remote_shards(peer_api_uri): r = requests.get(f"{peer_api_uri}/collections/{COLLECTION_NAME}/cluster") assert_http_ok(r) return r.json()["result"]['remote_shards'] def recover_snapshot(peer_api_uri, snapshot_url): r = requests.put(f"{peer_api_uri}/collections/{COLLECTION_NAME}/snapshots/recover", json={"location": snapshot_url}) assert_http_ok(r) return r.json()["result"] def test_recover_from_snapshot_1(tmp_path: pathlib.Path): recover_from_snapshot(tmp_path, 1) def test_recover_from_snapshot_2(tmp_path: pathlib.Path): recover_from_snapshot(tmp_path, 2) def recover_from_snapshot(tmp_path: pathlib.Path, n_replicas): assert_project_root() peer_api_uris, peer_dirs, bootstrap_uri = start_cluster(tmp_path, N_PEERS) create_collection(peer_api_uris[0], shard_number=N_SHARDS, replication_factor=n_replicas) wait_collection_exists_and_active_on_all_peers(collection_name=COLLECTION_NAME, peer_api_uris=peer_api_uris) create_payload_index(peer_api_uris[0], COLLECTION_NAME, "city", "keyword") wait_for_same_commit(peer_api_uris=peer_api_uris) upsert_random_points(peer_api_uris[0], 100) query_city = "London" dense_query_vector = random_dense_vector() dense_search_result = search(peer_api_uris[0], dense_query_vector, query_city) assert len(dense_search_result) > 0 sparse_query_vector = {"name": "sparse-text", "vector": random_sparse_vector()} sparse_search_result = search(peer_api_uris[0], sparse_query_vector, query_city) assert len(sparse_search_result) > 0 collection_info = get_collection_info(peer_api_uris[0], COLLECTION_NAME) points_with_indexed_payload = collection_info["payload_schema"]["city"]["points"] assert points_with_indexed_payload snapshot_name = create_snapshot(peer_api_uris[-1]) assert snapshot_name is not None # move file snapshot_path = Path(peer_dirs[-1]) / "snapshots" / COLLECTION_NAME / snapshot_name assert snapshot_path.exists() snapshot_path.rename(Path(peer_dirs[0]) / "snapshots" / COLLECTION_NAME / snapshot_name) process_peer_id = get_peer_id(peer_api_uris[-1]) local_shards = get_local_shards(peer_api_uris[-1]) # peer_0_remote_shards = get_remote_shards(peer_api_uris[0]) # Kill last peer p = processes.pop() p.kill() # Remove last peer from cluster res = requests.delete(f"{peer_api_uris[0]}/cluster/peer/{process_peer_id}?force=true") assert_http_ok(res) new_peer_dir = make_peer_folder(tmp_path, N_PEERS + 1) new_url = start_peer(new_peer_dir, f"peer_snapshot_{N_PEERS + 1}.log", bootstrap_uri) # Wait node is up and synced while True: try: res = requests.get(f"{new_url}/collections") except requests.exceptions.ConnectionError: time.sleep(1) continue if not res.ok: time.sleep(1) # Wait to node is up continue collections = set(collection['name'] for collection in res.json()["result"]['collections']) if COLLECTION_NAME not in collections: time.sleep(1) # Wait to sync with consensus continue break # Recover snapshot # All nodes share the same snapshot directory, so it is fine to use any snapshot_url = f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/snapshots/{snapshot_name}" print(f"Recovering snapshot {snapshot_url} on {new_url}") recover_snapshot(new_url, snapshot_url) wait_collection_exists_and_active_on_all_peers(collection_name=COLLECTION_NAME, peer_api_uris=[new_url]) new_local_shards = get_local_shards(new_url) new_local_shards = sorted(new_local_shards, key=lambda x: x['shard_id']) local_shards = sorted(local_shards, key=lambda x: x['shard_id']) assert len(new_local_shards) == len(local_shards) for i in range(len(new_local_shards)): assert new_local_shards[i] == local_shards[i] # check that the dense vectors are still the same new_dense_search_result = search(new_url, dense_query_vector, query_city) assert len(new_dense_search_result) == len(dense_search_result) for i in range(len(new_dense_search_result)): assert new_dense_search_result[i] == dense_search_result[i] # check that the sparse vectors are still the same new_sparse_search_result = search(new_url, sparse_query_vector, query_city) assert len(new_sparse_search_result) == len(sparse_search_result) for i in range(len(new_sparse_search_result)): # skip score check because it is not deterministic new_sparse_search_result[i]["score"] = sparse_search_result[i]["score"] assert new_sparse_search_result[i] == sparse_search_result[i] new_collection_info = get_collection_info(new_url, COLLECTION_NAME) new_points_with_indexed_payload = new_collection_info["payload_schema"]["city"]["points"] assert new_points_with_indexed_payload == points_with_indexed_payload peer_0_remote_shards_new = get_remote_shards(peer_api_uris[0]) for shard in peer_0_remote_shards_new: print("remote shard", shard) assert shard['state'] == 'Active' assert len(peer_0_remote_shards_new) == 2 * n_replicas