Pythonsharding.py
Copy
1
import sqlite3
2
import hashlib
3
import threading
4
import time
5
import random
6
import string
7
from dataclasses import dataclass
8
 
9
# ── Shard setup ───────────────────────────────────────────────────────
10
NUM_SHARDS = 4
11
 
12
shards: list[sqlite3.Connection] = []
13
shard_locks: list[threading.Lock] = []
14
 
15
for i in range(NUM_SHARDS):
16
    conn = sqlite3.connect(':memory:', check_same_thread=False)
17
    conn.execute("""
18
        CREATE TABLE users (
19
            id INTEGER PRIMARY KEY,
20
            username TEXT NOT NULL,
21
            email TEXT NOT NULL,
22
            created_at REAL
23
        )
24
    """)
25
    conn.commit()
26
    shards.append(conn)
27
    shard_locks.append(threading.Lock())
28
 
29
# ── Hash-based shard router ───────────────────────────────────────────
30
def get_shard_index(user_id: int) -> int:
31
    """Route user to shard based on user_id."""
32
    return user_id % NUM_SHARDS
33
 
34
def get_shard(user_id: int) -> tuple[sqlite3.Connection, threading.Lock]:
35
    idx = get_shard_index(user_id)
36
    return shards[idx], shard_locks[idx]
37
 
38
# ── CRUD operations ───────────────────────────────────────────────────
39
def create_user(user_id: int, username: str, email: str) -> dict:
40
    conn, lock = get_shard(user_id)
41
    with lock:
42
        conn.execute(
43
            "INSERT INTO users (id, username, email, created_at) VALUES (?, ?, ?, ?)",
44
            (user_id, username, email, time.time())
45
        )
46
        conn.commit()
47
    return {"id": user_id, "shard": get_shard_index(user_id)}
48
 
49
def get_user(user_id: int) -> dict | None:
50
    conn, lock = get_shard(user_id)
51
    with lock:
52
        row = conn.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()
53
    if row:
54
        return {"id": row[0], "username": row[1], "email": row[2]}
55
    return None
56
 
57
# ── Cross-shard query ──────────────────────────────────────────────────
58
def count_all_users() -> int:
59
    """Scatter-gather query across all shards."""
60
    total = 0
61
    for i, (conn, lock) in enumerate(zip(shards, shard_locks)):
62
        with lock:
63
            count = conn.execute("SELECT COUNT(*) FROM users").fetchone()[0]
64
        total += count
65
    return total
66
 
67
def search_users_by_username(pattern: str) -> list[dict]:
68
    """Fan-out search across all shards."""
69
    results = []
70
    for conn, lock in zip(shards, shard_locks):
71
        with lock:
72
            rows = conn.execute(
73
                "SELECT id, username, email FROM users WHERE username LIKE ?",
74
                (f"%{pattern}%",)
75
            ).fetchall()
76
        results.extend({"id": r[0], "username": r[1], "email": r[2]} for r in rows)
77
    return results
78
 
79
def get_shard_stats() -> list[dict]:
80
    """Get distribution stats per shard."""
81
    stats = []
82
    for i, (conn, lock) in enumerate(zip(shards, shard_locks)):
83
        with lock:
84
            count = conn.execute("SELECT COUNT(*) FROM users").fetchone()[0]
85
        stats.append({"shard": i, "count": count})
86
    return stats
87
 
88
# ── Write performance benchmark ───────────────────────────────────────
89
def bench_writes(n: int = 1000):
90
    """Benchmark concurrent writes across shards."""
91
    errors = []
92
    
93
    def write_batch(start: int, end: int):
94
        for i in range(start, end):
95
            try:
96
                create_user(i, f"user_{i}", f"user{i}@example.com")
97
            except Exception as e:
98
                errors.append(str(e))
99
    
100
    threads = []
101
    batch_size = n // 4
102
    for t in range(4):
103
        thread = threading.Thread(
104
            target=write_batch,
105
            args=(t * batch_size, (t + 1) * batch_size)
106
        )
107
        threads.append(thread)
108
    
109
    start = time.time()
110
    for t in threads:
111
        t.start()
112
    for t in threads:
113
        t.join()
114
    elapsed = time.time() - start
115
    
116
    return elapsed, errors
117
 
118
# ── Run demo ──────────────────────────────────────────────────────────
119
print("=== Sharding Demo ===
120
")
121
 
122
# Write 1000 users
123
elapsed, errors = bench_writes(1000)
124
print(f"Wrote 1000 users in {elapsed:.3f}s ({1000/elapsed:.0f} writes/sec)")
125
print(f"Errors: {len(errors)}
126
")
127
 
128
# Distribution
129
print("=== Shard Distribution ===")
130
for stat in get_shard_stats():
131
    bar = "█" * (stat["count"] // 10)
132
    print(f"  Shard {stat['shard']}: {stat['count']:4d} users  {bar}")
133
 
134
# Cross-shard query
135
total = count_all_users()
136
print(f"
137
Cross-shard count: {total} users
138
")
139
 
140
# Read by key
141
user = get_user(42)
142
print(f"get_user(42): {user} (shard {get_shard_index(42)})")
143
 
144
user = get_user(999)
145
print(f"get_user(999): {user} (shard {get_shard_index(999)})")
146
 
147
# Search
148
results = search_users_by_username("user_50")
149
print(f"
150
Search 'user_50': {len(results)} results across shards")
151