qdrant cluster distribution

qdrant-cluster-shards.py

import argparse
import requests

def get_data_from_api(base_url, endpoint):
    response = requests.get(base_url + endpoint)
    response.raise_for_status()
    return response.json()

def parse_cluster_peers(cluster_data):
    peers = cluster_data.get("result", {}).get("peers", {})
    ip_peer_map = {}
    for peer_id, peer_info in peers.items():
        uri = peer_info.get("uri", "")
        ip_address = uri.split("//")[-1].split(":")[0]
        ip_peer_map[ip_address] = int(peer_id)
    return ip_peer_map

def parse_shards(collections_data):
    local_shards = collections_data.get("result", {}).get("local_shards", [])
    remote_shards = collections_data.get("result", {}).get("remote_shards", [])

    peer_shard_map = {}

    for shard in local_shards:
        peer_id = collections_data.get("result", {}).get("peer_id")
        shard_id = shard.get("shard_id")
        peer_shard_map.setdefault(peer_id, []).append(shard_id)

    for shard in remote_shards:
        peer_id = shard.get("peer_id")
        shard_id = shard.get("shard_id")
        peer_shard_map.setdefault(peer_id, []).append(shard_id)

    return peer_shard_map

def main():
    parser = argparse.ArgumentParser(description="Fetch Qdrant cluster and shard information.")
    parser.add_argument("ip", help="Cluster node IP (e.g., 10.0.0.6)")
    parser.add_argument("collection", help="Collection name")
    args = parser.parse_args()

    base_url = f"http://{args.ip}:6333"
    cluster_endpoint = "/cluster"
    collections_endpoint = f"/collections/{args.collection}/cluster"

    try:
        cluster_data = get_data_from_api(base_url, cluster_endpoint)
        collections_data = get_data_from_api(base_url, collections_endpoint)
    except requests.exceptions.RequestException as e:
        print(f"Error fetching data from API: {e}")
        return

    ip_peer_map = parse_cluster_peers(cluster_data)
    peer_shard_map = parse_shards(collections_data)

    ip_shard_map = {}
    for ip, peer_id in ip_peer_map.items():
        ip_shard_map[ip] = peer_shard_map.get(peer_id, [])

    for ip, shard_ids in ip_shard_map.items():
        peer_id = ip_peer_map[ip]
        print(f"IP: {ip}, Peer ID: {peer_id}, Shard IDs: {shard_ids}")

if __name__ == "__main__":
    main()