# single file demo - Python 3.11+ asyncio stdlib only import asyncio from contextlib import asynccontextmanager from typing import List, Tuple, Dict, Any # ---------------------------------------------------------------------- # Mock in-memory DB # ---------------------------------------------------------------------- class MockDB: users: Dict[int, Dict[str, Any]] = {} posts: Dict[int, List[Dict[str, Any]]] = {} def seed_db(): MockDB.users.clear() MockDB.posts.clear() for uid in range(1, 26): MockDB.users[uid] = {"id": uid, "name": f"user{uid}", "label": f"label{uid % 5}"} MockDB.posts[uid] = [ {"id": pid, "user_id": uid, "title": f"post{pid} for user{uid}"} for pid in range(1, 6) ] seed_db() # ---------------------------------------------------------------------- # Mock pool / connection # ---------------------------------------------------------------------- class MockPool: def __init__(self, max_size: int = 5, acquire_timeout: float = 2.0): self.max_size = max_size self.acquire_timeout = acquire_timeout self._sem = asyncio.Semaphore(max_size) self._checked_out = 0 self._lock = asyncio.Lock() async def acquire(self) -> "MockConnection": try: await asyncio.wait_for(self._sem.acquire(), timeout=self.acquire_timeout) except asyncio.TimeoutError: raise TimeoutError(f"Pool exhausted after {self.acquire_timeout}s") async with self._lock: self._checked_out += 1 return MockConnection(self) async def _release(self): async with self._lock: self._checked_out -= 1 self._sem.release() @property def checked_out(self) -> int: return self._checked_out class MockConnection: def __init__(self, pool: MockPool): self.pool = pool async def fetch(self, query: str, params: List[Any]) -> List[Dict[str, Any]]: await asyncio.sleep(0) # simulate IO if query.startswith("SELECT COUNT"): return [{"count": len(MockDB.users)}] if "FROM users WHERE id =" in query: uid = params[0] user = MockDB.users.get(uid) return [user] if user else [] if "FROM users ORDER BY id LIMIT" in query: limit, offset = params sorted_users = [MockDB.users[i] for i in sorted(MockDB.users)] return sorted_users[offset: offset + limit] if "FROM posts WHERE user_id =" in query: uid = params[0] return MockDB.posts.get(uid, []) raise ValueError(f"Unknown fetch query: {query}") async def execute(self, query: str, params: List[Any]) -> int: await asyncio.sleep(0) if query.startswith("UPDATE users SET label"): label, uid = params if uid not in MockDB.users: raise KeyError(f"User {uid} not found") MockDB.users[uid]["label"] = label return 1 raise ValueError(f"Unknown execute query: {query}") async def release(self): await self.pool._release() @asynccontextmanager async def get_connection(pool: MockPool): conn = await pool.acquire() try: yield conn finally: await conn.release() # ---------------------------------------------------------------------- # Service API # ---------------------------------------------------------------------- class UserService: def __init__(self, pool: MockPool): self.pool = pool async def get_users(self, page: int, page_size: int) -> dict: if page < 1: raise ValueError("page >= 1 required") if not 1 <= page_size <= 100: raise ValueError("1 <= page_size <= 100 required") offset = (page - 1) * page_size async with get_connection(self.pool) as conn: total_row = await conn.fetch("SELECT COUNT(*) FROM users", []) items = await conn.fetch( "SELECT id, name, label FROM users ORDER BY id LIMIT $1 OFFSET $2", [page_size, offset], ) return {"items": items, "page": page, "page_size": page_size, "total": total_row[0]["count"]} async def get_user_with_posts(self, user_id: int) -> dict: async with get_connection(self.pool) as conn: user_rows = await conn.fetch( "SELECT id, name, label FROM users WHERE id = $1", [user_id] ) if not user_rows: raise KeyError(f"User {user_id} not found") user = user_rows[0] posts = await conn.fetch( "SELECT id, title FROM posts WHERE user_id = $1", [user_id] ) return {"user": user, "posts": posts} async def relabel_users(self, pairs: List[Tuple[int, str]]) -> int: if not pairs: return 0 async with get_connection(self.pool) as conn: # snapshot for rollback snapshot = {} for uid, _ in pairs: if uid in MockDB.users: snapshot[uid] = MockDB.users[uid]["label"] try: for uid, label in pairs: await conn.execute( "UPDATE users SET label = $1 WHERE id = $2", [label, uid] ) return len(pairs) except Exception: # rollback to original labels for uid, orig in snapshot.items(): MockDB.users[uid]["label"] = orig return 0 # ---------------------------------------------------------------------- # Tests # ---------------------------------------------------------------------- async def run_tests(): pool = MockPool(max_size=5, acquire_timeout=2.0) service = UserService(pool) # a) no leak over 100 calls for _ in range(100): await service.get_users(page=1, page_size=10) assert pool.checked_out == 0, "leak detected" print("✓ a) 100 calls, pool checked_out = 0") # b) pagination correctness res1 = await service.get_users(page=1, page_size=10) assert res1["total"] == 25 assert len(res1["items"]) == 10 res3 = await service.get_users(page=3, page_size=10) assert len(res3["items"]) == 5 # remainder res4 = await service.get_users(page=4, page_size=10) assert len(res4["items"]) == 0 print("✓ b) pagination math correct") # c) transactional bulk rollback original = {uid: MockDB.users[uid]["label"] for uid in (1, 2, 3)} pairs = [(1, "new1"), (2, "new2"), (999, "bad")] result = await service.relabel_users(pairs) assert result == 0 for uid in (1, 2): assert MockDB.users[uid]["label"] == original[uid], "rollback failed" print("✓ c) relabel_users rolls back on bad id") # d) concurrent access without exhaustion tasks = [service.get_user_with_posts(uid) for uid in range(1, 11)] results = await asyncio.gather(*tasks) assert len(results) == 10 assert pool.checked_out == 0 print("✓ d) 10 concurrent get_user_with_posts, no exhaustion") # validation & missing user try: await service.get_users(page=0, page_size=10) assert False except ValueError: pass try: await service.get_user_with_posts(9999) assert False except KeyError: pass print("✓ validation & missing user errors") # injection safety demo - params are never interpolated await service.relabel_users([(1, "label'; DROP TABLE users;--")]) assert MockDB.users[1]["label"] == "label'; DROP TABLE users;--" print("✓ parameterized queries prevent injection") print("\nAll tests passed.") if __name__ == "__main__": asyncio.run(run_tests())