Repository: qdrant/qdrant
File: tests/consensus_tests/test_consensus_compaction.py (Python)

Code before:
    # Add cluster metadata
    put_metadata_key(peer_api_uris, 'my_metadata', 'My value!')

    # Repeatedly drop, re-create collection to accumulate Raft log entries
    for i in range(0, 5):
        print(f"recreating collection {i}")
        # Drop test_collection
        drop_collection(peer_api_uris[0], "test_collection", timeout=5)
        # Re-create test_collection
        create_collection(peer_api_uris[0], shard_number=N_SHARDS, replication_factor=N_REPLICA, timeout=3)
        res = requests.get(f"{peer_api_uris[1]}/collections")
        print(res.json())

    # Wait for collection to be ready on all peers
    for peer_uri in peer_api_uris:
        wait_for_all_replicas_active(peer_uri, "test_collection")

    # Add extra node
    # Due to aggressive consensus WAL compaction, the peer has to join by consensus snapshot
    peer_dirs.append(make_peer_folder(tmp_path, N_PEERS))
    new_url = start_peer(peer_dirs[-1], "peer_3_extra.log", bootstrap_uri, port=21000, extra_env=env)
    peer_api_uris.append(new_url)
    wait_all_peers_up([new_url])

    # Ensure cluster metadata is consistent on all peers
    # Failed before <https://github.com/qdrant/qdrant/pull/6014>
    get_metadata_key(peer_api_uris, 'my_metadata', 'My value!')


def test_consensus_compaction_shard_keys(tmp_path: pathlib.Path):
    """
    Basic consensus compaction test using user-defined shard keys.
    Does a few consensus operations on a cluster. Then adds a new peer and
    asserts it boots with correct consensus state.

    Tests:
    - <https://github.com/qdrant/qdrant/pull/6209>
    - <https://github.com/qdrant/qdrant/pull/6212>
    """

    assert_project_root()

    # Have shard keys with different types to properly test them in consensus snapshots
    SHARD_KEYS = {
        1: "some_key",
        2: "1",
        3: 2,
    }

    env = {
        # Aggressively compact consensus WAL

Review comment (suggestion):
This makes sense! If we query too early we might not have the data applied from the snapshot yet. IIRC we normally handle this through `/readyz` in cloud.

Code after:
    # Add cluster metadata
    put_metadata_key(peer_api_uris, 'my_metadata', 'My value!')

    # Repeatedly drop, re-create collection to accumulate Raft log entries
    for i in range(0, 5):
        print(f"recreating collection {i}")
        # Drop test_collection
        drop_collection(peer_api_uris[0], "test_collection", timeout=5)
        # Re-create test_collection
        create_collection(peer_api_uris[0], shard_number=N_SHARDS, replication_factor=N_REPLICA, timeout=3)
        res = requests.get(f"{peer_api_uris[1]}/collections")
        print(res.json())

    # Wait for collection to be ready on all peers
    for peer_uri in peer_api_uris:
        wait_for_all_replicas_active(peer_uri, "test_collection")

    # Add extra node
    # Due to aggressive consensus WAL compaction, the peer has to join by consensus snapshot
    peer_dirs.append(make_peer_folder(tmp_path, N_PEERS))
    new_url = start_peer(peer_dirs[-1], "peer_3_extra.log", bootstrap_uri, port=21000, extra_env=env)
    peer_api_uris.append(new_url)

    # Wait for collection to be ready on the new peer
    wait_all_peers_up([new_url])
    wait_collection_exists_and_active_on_all_peers(collection_name="test_collection", peer_api_uris=[new_url])

    # Ensure cluster metadata is consistent on all peers
    # Failed before <https://github.com/qdrant/qdrant/pull/6014>
    assert_consistent_metadata_key(peer_api_uris, 'my_metadata', 'My value!')


def test_consensus_compaction_shard_keys(tmp_path: pathlib.Path):
    """
    Basic consensus compaction test using user-defined shard keys.
    Does a few consensus operations on a cluster. Then adds a new peer and
    asserts it boots with correct consensus state.

    Tests:
    - <https://github.com/qdrant/qdrant/pull/6209>
    - <https://github.com/qdrant/qdrant/pull/6212>
    """

    assert_project_root()

    # Have shard keys with different types to properly test them in consensus snapshots
    SHARD_KEYS = {
        1: "some_key",
        2: "1",
        3: 2,
    }