Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/redis/_himport_exec.py: 9%

1"""Shared HIMPORT wire-execution helpers for the synchronous clients. 

2 

3The PREPARE / SET / DISCARD packed-write drain loops, per-connection version 

4bookkeeping, and ``NoSuchFieldsetError`` re-prepare-and-retry are identical for 

5the standalone (:class:`redis.Redis`) and cluster (:class:`redis.RedisCluster`) 

6sync clients, differing only in (a) which object provides ``parse_response`` and 

7(b) the cluster-only ASK-redirect handling. These functions take that object as 

8``node`` and an ``asking`` flag (``False`` -- and a no-op -- for standalone), so 

9both clients share one implementation instead of copy-pasting the logic. The 

10per-class ``_himport_*`` methods are thin delegators to these functions. 

11 

12The async mirror lives in :mod:`redis.asyncio._himport_exec`; the two are kept 

13separate on purpose (the project maintains parallel sync/async stacks by hand). 

14""" 

15 

16from redis.exceptions import NoSuchFieldsetError, ResponseError 

17from redis.himport import ( 

18 HIMPORT_DISCARD, 

19 HIMPORT_PREPARE, 

20 HIMPORT_SET, 

21 HImportRegistry, 

22 himport_discard_command, 

23 himport_prepare_command, 

24 himport_set_command, 

25 parse_himport_set_args, 

26) 

27 

28 

29def reconcile_discards(node, conn): 

30 """DISCARD, on ``conn``, any prepared fieldset removed from the registry. 

31 

32 Runs at most once per registry mutation: the connection records the registry 

33 ``revision`` it last reconciled against, so unchanged registries are a no-op. 

34 ``node`` supplies ``parse_response`` (the standalone client itself, or the 

35 owning node's client in cluster mode). 

36 """ 

37 registry = conn.himport_registry 

38 if registry is None or conn._himport_reconciled_revision == registry.revision: 

39 return 

40 # Snapshot the revision *before* computing ``stale`` so the value stamped at 

41 # the end is never newer than the registry state ``stale`` reflects. A 

42 # concurrent ``himport_discard`` (thread-shared sync client) that lands after 

43 # this point only leaves the connection marked behind the live revision, so 

44 # the next reconcile re-runs and catches it. Re-reading ``registry.revision`` 

45 # at the end instead would stamp a discard this connection never sent. 

46 reconciled_to = registry.revision 

47 stale = registry.names_to_discard(list(conn._himport_prepared)) 

48 if stale: 

49 conn.send_packed_command( 

50 conn.pack_commands([himport_discard_command(n) for n in stale]) 

51 ) 

52 # One reply per packed DISCARD must be read regardless of a per-command 

53 # ResponseError, otherwise the unread replies desync the pooled socket. 

54 # Drain every reply, then surface the first error (ConnectionError is not 

55 # caught: it tears the socket down, so no desync is possible). 

56 first_error = None 

57 for n in stale: 

58 try: 

59 node.parse_response(conn, HIMPORT_DISCARD) 

60 except ResponseError as e: 

61 first_error = first_error or e 

62 conn._himport_prepared.pop(n, None) 

63 if first_error is not None: 

64 raise first_error 

65 conn._himport_reconciled_revision = reconciled_to 

66 

67 

68def prepare_and_set(node, conn, key, fieldset_name, values, fieldset, asking=False): 

69 """PREPARE ``fieldset`` bundled with the SET on ``conn`` (one packed write). 

70 

71 When ``asking`` is set (an ASK-redirected cluster SET) the batch becomes 

72 ``[PREPARE, ASKING, SET]`` so the per-command ASKING allowance falls 

73 immediately before the SET -- the only slot-scoped command. PREPARE is a 

74 connection-session command the ASKING flag does not gate, so placing it 

75 before ASKING is safe. Every reply is drained even on a per-command error so 

76 the packed replies never desync the pooled socket. 

77 """ 

78 commands = [himport_prepare_command(fieldset_name, fieldset.fields)] 

79 if asking: 

80 commands.append(("ASKING",)) 

81 commands.append(himport_set_command(key, fieldset_name, values)) 

82 conn.send_packed_command(conn.pack_commands(commands)) 

83 prep_error = ask_error = set_error = None 

84 set_resp = None 

85 try: 

86 node.parse_response(conn, HIMPORT_PREPARE) 

87 except ResponseError as e: 

88 prep_error = e 

89 if asking: 

90 try: 

91 node.parse_response(conn, "ASKING") 

92 except ResponseError as e: 

93 ask_error = e 

94 try: 

95 set_resp = node.parse_response(conn, HIMPORT_SET) 

96 except ResponseError as e: 

97 set_error = e 

98 

99 if prep_error: 

100 raise prep_error # PREPARE failure is the root cause 

101 else: 

102 conn._himport_prepared[fieldset_name] = fieldset.version 

103 

104 if ask_error: 

105 raise ask_error 

106 if set_error: 

107 raise set_error 

108 return set_resp 

109 

110 

111def execute_set(node, conn, key, fieldset_name, values, asking=False): 

112 """Execute an ``HIMPORT SET`` on ``conn`` with the required session setup. 

113 

114 Reconciles deferred discards, lazily bundles PREPARE with the SET on first 

115 use of a fieldset, and recovers once from a mid-connection fieldset loss 

116 (``NoSuchFieldsetError``) by re-PREPARE-and-retry. When ``asking`` is set the 

117 ASKING allowance is folded into the SET's own packed write so it immediately 

118 precedes the (slot-scoped) SET; the session setup runs first, since those are 

119 connection-session commands the flag does not gate. 

120 """ 

121 reconcile_discards(node, conn) 

122 

123 registry = conn.himport_registry 

124 fieldset = registry.get(fieldset_name) if registry is not None else None 

125 # Lazy PREPARE bundled with SET on first use of this fieldset. 

126 if ( 

127 fieldset is not None 

128 and conn._himport_prepared.get(fieldset_name) != fieldset.version 

129 ): 

130 return prepare_and_set( 

131 node, conn, key, fieldset_name, values, fieldset, asking=asking 

132 ) 

133 

134 # Believed already prepared (or an unregistered fieldset): bare SET, with 

135 # ASKING packed immediately before it when this is an ASK redirect. 

136 if asking: 

137 conn.send_packed_command( 

138 conn.pack_commands( 

139 [("ASKING",), himport_set_command(key, fieldset_name, values)] 

140 ) 

141 ) 

142 try: 

143 node.parse_response(conn, "ASKING") 

144 except ResponseError as ask_error: 

145 # ASKING and SET were one packed write, so the SET reply is still 

146 # queued. Drain it before surfacing the ASKING error, otherwise the 

147 # connection returns to the pool with an unread reply and desyncs the 

148 # next borrower. 

149 try: 

150 node.parse_response(conn, HIMPORT_SET) 

151 except ResponseError: 

152 pass 

153 raise ask_error 

154 else: 

155 conn.send_command(*himport_set_command(key, fieldset_name, values)) 

156 try: 

157 return node.parse_response(conn, HIMPORT_SET) 

158 except NoSuchFieldsetError: 

159 # Server dropped the fieldset mid-connection without dropping the socket 

160 # (e.g. RESET / maxmemory-clients eviction): re-PREPARE on this healthy 

161 # connection and retry the SET once rather than reconnecting. Only for 

162 # registry-backed fieldsets; manual/unregistered usage propagates. 

163 if fieldset is None: 

164 raise 

165 conn._himport_prepared.pop(fieldset_name, None) 

166 return prepare_and_set( 

167 node, conn, key, fieldset_name, values, fieldset, asking=asking 

168 ) 

169 

170 

171def prepare_pipeline(node, conn, command_arg_lists): 

172 """Pre-flight ``conn`` for a pipeline batch containing ``HIMPORT SET``s. 

173 

174 The packed pipeline write bypasses the per-command lazy-PREPARE path, so the 

175 fieldsets referenced by the buffered SETs must be PREPAREd on ``conn`` first. 

176 Reconciles deferred discards, then PREPAREs every distinct registered fieldset 

177 the batch references that this connection has not already prepared, in one 

178 packed write. ``command_arg_lists`` is the batch's per-command positional-arg 

179 sequences (the caller extracts them from its own command representation). 

180 No-op when the batch has no registry-backed ``HIMPORT SET``. 

181 """ 

182 # Selection (registry check, deferred-discard reconcile, scan/dedup/version) 

183 # is shared with pipeline_prepares. This path differs only in that it sends the 

184 # PREPAREs as their own packed exchange -- rather than folding them into a 

185 # queued write -- then drains their replies and raises the first error. 

186 to_prepare = pipeline_prepares(node, conn, command_arg_lists) 

187 if not to_prepare: 

188 return 

189 conn.send_packed_command(conn.pack_commands(prepare_wire_commands(to_prepare))) 

190 # Every reply must be drained even on a per-command error, or the unread 

191 # replies desync the socket before the buffered batch is sent; then raise. 

192 first_error = drain_pipeline_prepares(node, conn, to_prepare) 

193 if first_error is not None: 

194 raise first_error 

195 

196 

197def pipeline_prepares(node, conn, command_arg_lists): 

198 """Return the fieldsets that must be PREPAREd on ``conn`` for this batch. 

199 

200 Like :func:`prepare_pipeline`, but does **not** send the PREPAREs: the caller 

201 folds them into the same packed write as the queued commands (see the pipeline 

202 executors), so the first pipeline use of a fieldset on a fresh or reconnected 

203 connection stays a single round trip instead of a separate PREPARE exchange 

204 followed by the batch. Deferred-discard reconciliation is still performed here, 

205 but it only touches the socket when discards are actually pending (rare); the 

206 common warm-up cost -- the first-use PREPARE -- is what gets folded. Returns an 

207 empty list when the batch references no not-yet-prepared registered fieldset, 

208 or when ``conn`` carries no real HIMPORT registry. 

209 """ 

210 registry = getattr(conn, "himport_registry", None) 

211 if not isinstance(registry, HImportRegistry): 

212 return [] 

213 reconcile_discards(node, conn) 

214 to_prepare = [] 

215 seen = set() 

216 for args in command_arg_lists: 

217 parsed = parse_himport_set_args(args) 

218 if parsed is None: 

219 continue 

220 fieldset_name = parsed[1] 

221 if fieldset_name in seen: 

222 continue 

223 seen.add(fieldset_name) 

224 fieldset = registry.get(fieldset_name) 

225 if ( 

226 fieldset is not None 

227 and conn._himport_prepared.get(fieldset_name) != fieldset.version 

228 ): 

229 to_prepare.append(fieldset) 

230 return to_prepare 

231 

232 

233def prepare_wire_commands(fieldsets): 

234 """The leading ``HIMPORT PREPARE`` wire commands the caller folds into a batch.""" 

235 return [himport_prepare_command(fs.name, fs.fields) for fs in fieldsets] 

236 

237 

238def drain_pipeline_prepares(node, conn, fieldsets): 

239 """Drain the ``len(fieldsets)`` leading PREPARE replies of a folded pipeline 

240 write, marking each fieldset prepared on success. 

241 

242 Returns the first ``ResponseError`` (or ``None``). The caller must still drain 

243 the queued command replies and only then surface this error: every reply on 

244 the wire has to be read before raising, or the pooled socket desyncs. 

245 """ 

246 first_error = None 

247 for fs in fieldsets: 

248 try: 

249 node.parse_response(conn, HIMPORT_PREPARE) 

250 except ResponseError as e: 

251 first_error = first_error or e 

252 continue 

253 conn._himport_prepared[fs.name] = fs.version 

254 return first_error