client.py

"""CLI client for the Raft KV store.

Usage:
    python client.py --server http://127.0.0.1:8001 set mykey myvalue
    python client.py --server http://127.0.0.1:8001 get mykey
    python client.py --server http://127.0.0.1:8001 delete mykey
    python client.py --server http://127.0.0.1:8001 status

Features:
    - Follows 307 redirects to find the leader
    - Generates unique client_id and request_id for duplicate detection
    - Retries on 503 with backoff
"""

import urllib.request
import urllib.error
import json
import uuid
import argparse
import sys
import time
from typing import Optional


class RaftClient:
    MAX_REDIRECTS = 5
    MAX_RETRIES = 3
    RETRY_BACKOFF = 0.5

    def __init__(self, server_url: str):
        self.server_url = server_url.rstrip("/")
        self.client_id = str(uuid.uuid4())
        self._request_counter = 0

    def set(self, key: str, value: str) -> dict:
        return self._request("PUT", f"/key/{key}", body=value)

    def get(self, key: str) -> dict:
        return self._request("GET", f"/key/{key}")

    def delete(self, key: str) -> dict:
        return self._request("DELETE", f"/key/{key}")

    def status(self) -> dict:
        return self._request("GET", "/status")

    def _next_request_id(self) -> str:
        self._request_counter += 1
        return f"{self.client_id}-{self._request_counter}"

    def _request(
        self, method: str, path: str, body: Optional[str] = None
    ) -> dict:
        url = self.server_url + path
        request_id = self._next_request_id()

        for retry in range(self.MAX_RETRIES):
            current_url = url
            for _ in range(self.MAX_REDIRECTS):
                headers = {
                    "X-Client-ID": self.client_id,
                    "X-Request-ID": request_id,
                }
                data = body.encode("utf-8") if body else None

                req = urllib.request.Request(
                    current_url, data=data, headers=headers, method=method
                )

                try:
                    with urllib.request.urlopen(req) as resp:
                        resp_body = resp.read().decode("utf-8")
                        return json.loads(resp_body) if resp_body else {}
                except urllib.error.HTTPError as e:
                    if e.code == 307:
                        location = e.headers.get("Location")
                        if location:
                            current_url = location
                            continue
                    elif e.code == 503:
                        break  # retry after backoff
                    elif e.code == 404:
                        return {"error": "key not found"}
                    else:
                        resp_body = e.read().decode("utf-8")
                        try:
                            return json.loads(resp_body)
                        except json.JSONDecodeError:
                            return {"error": f"HTTP {e.code}: {resp_body}"}
                except (urllib.error.URLError, ConnectionError) as e:
                    if retry < self.MAX_RETRIES - 1:
                        break  # retry after backoff
                    return {"error": str(e)}

            if retry < self.MAX_RETRIES - 1:
                time.sleep(self.RETRY_BACKOFF * (retry + 1))

        return {"error": "max retries exceeded"}


def main():
    parser = argparse.ArgumentParser(description="Raft KV store client")
    parser.add_argument("--server", default="http://127.0.0.1:8001", help="Server URL")
    subparsers = parser.add_subparsers(dest="command", required=True)

    set_parser = subparsers.add_parser("set", help="Set a key")
    set_parser.add_argument("key")
    set_parser.add_argument("value")

    get_parser = subparsers.add_parser("get", help="Get a key")
    get_parser.add_argument("key")

    del_parser = subparsers.add_parser("delete", help="Delete a key")
    del_parser.add_argument("key")

    subparsers.add_parser("status", help="Node status")

    args = parser.parse_args()
    client = RaftClient(args.server)

    if args.command == "set":
        result = client.set(args.key, args.value)
    elif args.command == "get":
        result = client.get(args.key)
    elif args.command == "delete":
        result = client.delete(args.key)
    elif args.command == "status":
        result = client.status()
    else:
        parser.print_help()
        sys.exit(1)

    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()