Coverage for src/lexigram/graphql/core/persisted_queries.py: 96%

117 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 04:37 +0800

1"""Automatic Persisted Queries (APQ) for GraphQL. 

2 

3This module provides APQ support to reduce bandwidth by allowing clients 

4to send query hashes instead of full query strings. 

5""" 

6 

7from __future__ import annotations 

8 

9from dataclasses import dataclass 

10from datetime import UTC, datetime, timedelta 

11from typing import Any, Protocol, runtime_checkable 

12 

13from lexigram import hashing # type: ignore[attr-defined] 

14from lexigram.logging import get_logger 

15from lexigram.primitives import clock as ambient_clock 

16 

17logger = get_logger(__name__) 

18 

19 

20@runtime_checkable 

21class PersistedQueryStore(Protocol): 

22 """Store for persisted query mappings. 

23 

24 Implement this protocol to provide custom storage for persisted queries. 

25 """ 

26 

27 async def get(self, hash: str) -> str | None: 

28 """Get a query by its hash. 

29 

30 Args: 

31 hash: SHA256 hash of the query. 

32 

33 Returns: 

34 The query string if found, None otherwise. 

35 """ 

36 ... 

37 

38 async def put(self, hash: str, query: str) -> None: 

39 """Store a query with its hash. 

40 

41 Args: 

42 hash: SHA256 hash of the query. 

43 query: The full query string. 

44 """ 

45 ... 

46 

47 

48class InMemoryPersistedQueryStore: 

49 """In-memory implementation of PersistedQueryStore. 

50 

51 Warning: Not suitable for production use with multiple instances. 

52 Use RedisPersistedQueryStore for production deployments. 

53 """ 

54 

55 def __init__(self, ttl_seconds: int | None = None): 

56 """Initialize the store. 

57 

58 Args: 

59 ttl_seconds: Optional time-to-live for entries in seconds. 

60 """ 

61 self._store: dict[str, tuple[str, datetime | None]] = {} 

62 self._ttl = timedelta(seconds=ttl_seconds) if ttl_seconds else None 

63 

64 def _now(self) -> datetime: 

65 return ambient_clock.now() 

66 

67 async def get(self, hash: str) -> str | None: 

68 """Get a query by its hash.""" 

69 if hash not in self._store: 

70 return None 

71 

72 query, expires = self._store[hash] 

73 

74 # Check expiration 

75 if expires and expires < self._now(): 

76 del self._store[hash] 

77 return None 

78 

79 return query 

80 

81 async def put(self, hash: str, query: str) -> None: 

82 """Store a query with its hash.""" 

83 expires = self._now() + self._ttl if self._ttl else None 

84 self._store[hash] = (query, expires) 

85 

86 def clear(self) -> None: 

87 """Clear all stored queries.""" 

88 self._store.clear() 

89 

90 

91class RedisPersistedQueryStore: 

92 """Redis-backed implementation of PersistedQueryStore. 

93 

94 Suitable for production use with multiple application instances. 

95 """ 

96 

97 def __init__( 

98 self, 

99 redis_client: Any, 

100 key_prefix: str = "graphql:apq:", 

101 ttl_seconds: int = 86400, # 24 hours default 

102 ): 

103 """Initialize the store. 

104 

105 Args: 

106 redis_client: Redis client instance (aioRedis or redis-py async). 

107 key_prefix: Prefix for Redis keys. 

108 ttl_seconds: Time-to-live for entries in seconds. 

109 """ 

110 self._redis = redis_client 

111 self._prefix = key_prefix 

112 self._ttl = ttl_seconds 

113 

114 async def get(self, hash: str) -> str | None: 

115 """Get a query by its hash.""" 

116 key = f"{self._prefix}{hash}" 

117 try: 

118 result = await self._redis.get(key) 

119 return result.decode("utf-8") if result else None 

120 except (OSError, RuntimeError, AttributeError) as e: 

121 logger.warning("Redis get failed for APQ: %s", e) 

122 return None 

123 

124 async def put(self, hash: str, query: str) -> None: 

125 """Store a query with its hash.""" 

126 key = f"{self._prefix}{hash}" 

127 try: 

128 await self._redis.setex(key, self._ttl, query) 

129 except (OSError, RuntimeError) as e: 

130 logger.warning("Redis put failed for APQ: %s", e) 

131 

132 

133class CacheBackendPersistedQueryStore: 

134 """CacheBackendProtocol-backed APQ store for multi-process deployments. 

135 

136 Uses the platform's :class:`~lexigram.contracts.cache.CacheBackendProtocol` 

137 abstraction rather than a raw Redis client — compatible with any 

138 configured cache provider (Redis, Memcached, etc.) without coupling 

139 this package to a specific driver. 

140 

141 Suitable for production use with multiple application instances. 

142 """ 

143 

144 def __init__( 

145 self, 

146 cache: Any, 

147 key_prefix: str = "graphql:apq:", 

148 ttl_seconds: int = 86400, 

149 ) -> None: 

150 """Initialize the store. 

151 

152 Args: 

153 cache: A :class:`~lexigram.contracts.cache.CacheBackendProtocol` instance. 

154 key_prefix: Prefix applied to all cache keys. 

155 ttl_seconds: Time-to-live for stored entries in seconds. 

156 """ 

157 self._cache = cache 

158 self._prefix = key_prefix 

159 self._ttl = ttl_seconds 

160 

161 async def get(self, hash: str) -> str | None: 

162 """Get a query by its hash.""" 

163 key = f"{self._prefix}{hash}" 

164 try: 

165 cached: str | None = await self._cache.get(key) 

166 return cached 

167 except (OSError, LookupError, RuntimeError) as e: 

168 logger.warning("cache_apq_get_failed", key=key, error=str(e)) 

169 return None 

170 

171 async def put(self, hash: str, query: str) -> None: 

172 """Store a query with its hash.""" 

173 key = f"{self._prefix}{hash}" 

174 try: 

175 await self._cache.set(key, query, ttl=self._ttl) 

176 except (OSError, LookupError, RuntimeError) as e: 

177 logger.warning("cache_apq_put_failed", key=key, error=str(e)) 

178 

179 

180@dataclass 

181class APQResult: 

182 """Result of APQ lookup. 

183 

184 Attributes: 

185 query: The full query string (if found). 

186 is_persisted: Whether the query was found in the store. 

187 hash: The hash used for lookup. 

188 """ 

189 

190 query: str | None = None 

191 is_persisted: bool = False 

192 hash: str = "" 

193 

194 

195def compute_query_hash(query: str) -> str: 

196 """Compute SHA256 hash of a GraphQL query. 

197 

198 Uses Sha256Hasher for wire compatibility with persisted queries. 

199 The algorithm must remain SHA256 to ensure client-server compatibility. 

200 

201 Args: 

202 query: The GraphQL query string. 

203 

204 Returns: 

205 Hex-encoded SHA256 hash. 

206 """ 

207 return str(hashing.hash_hex(query)) 

208 

209 

210class APQHandler: 

211 """Handler for Automatic Persisted Queries. 

212 

213 Implements the APQ protocol: 

214 1. Client sends query hash → Server looks up full query 

215 2. If not found, client sends hash + full query → Server stores 

216 3. Subsequent requests use just the hash 

217 """ 

218 

219 def __init__( 

220 self, 

221 store: PersistedQueryStore, 

222 enabled: bool = True, 

223 ): 

224 """Initialize the APQ handler. 

225 

226 Args: 

227 store: Persisted query store implementation. 

228 enabled: Whether APQ is enabled. 

229 """ 

230 self._store = store 

231 self._enabled = enabled 

232 

233 @property 

234 def enabled(self) -> bool: 

235 """Check if APQ is enabled.""" 

236 return self._enabled 

237 

238 async def resolve_query( 

239 self, 

240 query: str | None, 

241 extensions: dict | None = None, 

242 ) -> APQResult: 

243 """Resolve the full query using APQ. 

244 

245 Args: 

246 query: The query from the request (may be None for hash-only). 

247 extensions: The extensions from the GraphQL request. 

248 

249 Returns: 

250 APQResult with the resolved query. 

251 """ 

252 if not self._enabled: 

253 return APQResult(query=query, is_persisted=False) 

254 

255 # Check for APQ extension data 

256 apq_ext = None 

257 if extensions and "persistedQuery" in extensions: 

258 apq_ext = extensions["persistedQuery"] 

259 

260 # Case 1: No query, no APQ extension → error 

261 if not query and not apq_ext: 

262 return APQResult(query=query, is_persisted=False) 

263 

264 # Case 2: Has full query, no APQ → store it 

265 if query and not apq_ext: 

266 hash_value = compute_query_hash(query) 

267 await self._store.put(hash_value, query) 

268 return APQResult(query=query, hash=hash_value, is_persisted=False) 

269 

270 # Case 3: Has APQ extension 

271 if apq_ext: 

272 # Get the hash 

273 hash_value = apq_ext.get("sha256Hash") 

274 if not hash_value: 

275 return APQResult(query=query, is_persisted=False) 

276 

277 # Try to look up the query 

278 stored_query = await self._store.get(hash_value) 

279 

280 if stored_query: 

281 # Cache hit 

282 return APQResult(query=stored_query, hash=hash_value, is_persisted=True) 

283 

284 # Cache miss - client should send the full query 

285 if not query: 

286 # Client needs to send the full query 

287 return APQResult(hash=hash_value, is_persisted=False) 

288 

289 # Client sent both hash and query - store it 

290 await self._store.put(hash_value, query) 

291 return APQResult(query=query, hash=hash_value, is_persisted=False) 

292 

293 return APQResult(query=query, is_persisted=False) 

294 

295 def create_extension_response(self, hash: str) -> dict: 

296 """Create the APQ extension response. 

297 

298 Args: 

299 hash: The query hash. 

300 

301 Returns: 

302 Extension dict for the GraphQL response. 

303 """ 

304 return { 

305 "persistedQuery": { 

306 "sha256Hash": hash, 

307 "version": 1, 

308 }, 

309 } 

310 

311 

312# Default APQ handler factory 

313def create_apq_handler( 

314 store_type: str = "memory", 

315 **kwargs: Any, 

316) -> APQHandler: 

317 """Create an APQ handler with the specified store. 

318 

319 Args: 

320 store_type: Type of store ('memory' or 'redis'). 

321 **kwargs: Additional arguments for the store. 

322 

323 Returns: 

324 Configured APQHandler. 

325 """ 

326 store: PersistedQueryStore 

327 if store_type == "memory": 

328 store = InMemoryPersistedQueryStore(**kwargs) 

329 elif store_type == "redis": 

330 store = RedisPersistedQueryStore(**kwargs) 

331 else: 

332 raise ValueError(f"Unknown APQ store type: {store_type}") 

333 

334 return APQHandler(store=store) 

335 

336 

337__all__ = [ 

338 "APQHandler", 

339 "APQResult", 

340 "InMemoryPersistedQueryStore", 

341 "PersistedQueryStore", 

342 "RedisPersistedQueryStore", 

343 "compute_query_hash", 

344 "create_apq_handler", 

345]