server.py

#!/usr/bin/env python3
"""
03: GET/SET — In-memory key-value store with proper command dispatch.

Now that the parser works, wire up the smallest useful command set:
SET, GET, DEL, EXISTS, PING. Storage is a plain Python dict.

Usage:
    uv run python 03-get-set/server.py

    redis-cli -p 6380
    > SET name Param
    OK
    > GET name
    "Param"
    > DEL name
    (integer) 1
"""
from __future__ import annotations

import socket

HOST = "127.0.0.1"
PORT = 6380
CRLF = b"\r\n"

# --- Store ---
STORE: dict[str, str] = {}


# --- RESP parser + encoder (same as step 2, kept inline so each module is self-contained) ---

def parse_resp(buf: bytes):
    if not buf:
        return None
    type_byte = buf[0:1]
    end = buf.find(CRLF)
    if end < 0:
        return None
    header = buf[1:end].decode("ascii")
    after = end + 2
    if type_byte == b"+":
        return header, after
    if type_byte == b"-":
        return Exception(header), after
    if type_byte == b":":
        return int(header), after
    if type_byte == b"$":
        n = int(header)
        if n == -1:
            return None, after
        end_payload = after + n
        if len(buf) < end_payload + 2:
            return None
        return buf[after:end_payload].decode("utf-8"), end_payload + 2
    if type_byte == b"*":
        count = int(header)
        if count == -1:
            return None, after
        items = []
        cursor = after
        for _ in range(count):
            r = parse_resp(buf[cursor:])
            if r is None:
                return None
            v, c = r
            items.append(v)
            cursor += c
        return items, cursor
    raise ValueError(f"Unknown RESP type byte: {type_byte!r}")


def encode_simple(s):    return b"+" + s.encode("ascii") + CRLF
def encode_error(s):     return b"-" + s.encode("ascii") + CRLF
def encode_integer(n):   return b":" + str(n).encode("ascii") + CRLF
def encode_bulk(s):
    if s is None:
        return b"$-1\r\n"
    data = s.encode("utf-8")
    return b"$" + str(len(data)).encode("ascii") + CRLF + data + CRLF


# --- Command dispatch ---

def cmd_ping(args):
    return encode_simple("PONG") if not args else encode_bulk(args[0])


def cmd_set(args):
    if len(args) < 2:
        return encode_error("ERR wrong number of arguments for 'set'")
    key, value = args[0], args[1]
    STORE[key] = value
    return encode_simple("OK")


def cmd_get(args):
    if len(args) != 1:
        return encode_error("ERR wrong number of arguments for 'get'")
    return encode_bulk(STORE.get(args[0]))


def cmd_del(args):
    if not args:
        return encode_error("ERR wrong number of arguments for 'del'")
    deleted = 0
    for k in args:
        if k in STORE:
            del STORE[k]
            deleted += 1
    return encode_integer(deleted)


def cmd_exists(args):
    if not args:
        return encode_error("ERR wrong number of arguments for 'exists'")
    return encode_integer(sum(1 for k in args if k in STORE))


def cmd_keys(args):
    # Simple star-pattern only: KEYS *  returns everything.
    if not args:
        return encode_error("ERR wrong number of arguments for 'keys'")
    if args[0] != "*":
        return encode_error("ERR only KEYS * is supported in this workshop")
    keys = list(STORE.keys())
    parts = [b"*" + str(len(keys)).encode("ascii") + CRLF]
    for k in keys:
        parts.append(encode_bulk(k))
    return b"".join(parts)


def cmd_command(args):
    # Stub: redis-cli sends COMMAND DOCS at startup; we just return an empty array
    return b"*0\r\n"


COMMANDS = {
    "PING": cmd_ping,
    "SET": cmd_set,
    "GET": cmd_get,
    "DEL": cmd_del,
    "EXISTS": cmd_exists,
    "KEYS": cmd_keys,
    "COMMAND": cmd_command,
}


def handle_command(frame) -> bytes:
    if not isinstance(frame, list) or not frame:
        return encode_error("ERR expected array of bulk strings")
    name = frame[0]
    if not isinstance(name, str):
        return encode_error("ERR command name must be string")
    handler = COMMANDS.get(name.upper())
    if handler is None:
        return encode_error(f"ERR unknown command '{name}'")
    return handler(frame[1:])


# --- Server ---

def handle_client(client: socket.socket, addr) -> None:
    print(f"  [conn] {addr} connected")
    buf = b""
    with client:
        while True:
            chunk = client.recv(4096)
            if not chunk:
                break
            buf += chunk
            while True:
                result = parse_resp(buf)
                if result is None:
                    break
                command, consumed = result
                buf = buf[consumed:]
                print(f"  [recv] {addr}: {command!r}")
                client.sendall(handle_command(command))
    print(f"  [conn] {addr} disconnected")


def main() -> None:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind((HOST, PORT))
        server.listen(1)
        print(f"key-value server listening on {HOST}:{PORT}")
        while True:
            client, addr = server.accept()
            handle_client(client, addr)


if __name__ == "__main__":
    main()