kedis-python 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

kedis_python/commands.py

@@ -136,57 +136,82 @@ class CommandHandler:
         if not tokens:
             return "-ERR empty command or syntax error", []
 
-        cmd = tokens[0].upper()
-        args = tokens[
-            1:
-        ]  # (Though your current code passes tokens, let's keep your routing exact)
-
         # 🚀 TRACK A: Start the high-precision telemetry stopwatch (in nanoseconds)
         start_time = time.perf_counter_ns()
+        try:
+            with self._engine_lock:
+                return self._run_locked(tokens, client_socket)
+        finally:
+            self._record_duration(tokens, start_time)
+
+    def execute_batch(self, batch: list, client_socket=None):
+        """Runs several commands as one unit: the engine lock is taken once, so
+        no other client's command can run between them, and the AOF records
+        them back to back. Returns (results, forwards): one result per command,
+        and the commands a master must send to replicas, in order (see
+        replication_payloads). Commands are not skipped or reordered; one that
+        fails gives an error result and the rest still run, like the one-at-a-
+        time path."""
+        results: list = []
+        forwards: list = []
+        write_commands = self.WRITE_COMMANDS
+        starts: list = []
+        with self._engine_lock:
+            for tokens in batch:
+                if not tokens:
+                    results.append("-ERR empty command or syntax error")
+                    starts.append(None)
+                    continue
+                starts.append(time.perf_counter_ns())
+                result, records = self._run_locked(tokens, client_socket)
+                results.append(result)
+                if tokens[0].upper() in write_commands:
+                    forwards.extend(self.replication_payloads(tokens, result, records))
+        for tokens, started in zip(batch, starts):
+            if started is not None:
+                self._record_duration(tokens, started)
+        return results, forwards
+
+    def _run_locked(self, tokens: list[str], client_socket=None):
+        """Runs one command. The caller holds the engine lock."""
+        cmd = tokens[0].upper()
         result = None
         records = []
 
         try:
-            with self._engine_lock:
-                try:
-                    if cmd == "XGAPPLY" and client_socket is not None:
-                        # Internal replication command: a network client
-                        # must not be able to write raw group records.
-                        result = f"-ERR unknown command '{cmd}'"
-                    elif cmd in self._commands:
-                        if cmd in ["SUBSCRIBE", "PUBLISH"]:
-                            handler_function = self._commands[cmd]
-                            result = handler_function(tokens, client_socket)
-                        else:
-                            handler_function = self._commands[cmd]
-                            # Note: Depending on how your handlers are registered,
-                            # they might expect 'tokens' or 'args'. Keeping your existing call:
-                            result = handler_function(tokens)
+            try:
+                if cmd == "XGAPPLY" and client_socket is not None:
+                    # Internal replication command: a network client
+                    # must not be able to write raw group records.
+                    result = f"-ERR unknown command '{cmd}'"
+                elif cmd in self._commands:
+                    handler_function = self._commands[cmd]
+                    if cmd in ["SUBSCRIBE", "PUBLISH"]:
+                        result = handler_function(tokens, client_socket)
                     else:
-                        result = f"-ERR unknown command '{cmd}'"
-                finally:
-                    records = self.store.drain_group_records()
-            return result, records
-
-        except TypeError as e:
-            result = f"-ERR wrong number of arguments for '{cmd}' command"
+                        # Handlers take the whole token list, command name included.
+                        result = handler_function(tokens)
+                else:
+                    result = f"-ERR unknown command '{cmd}'"
+            finally:
+                records = self.store.drain_group_records()
             return result, records
+        except TypeError:
+            return f"-ERR wrong number of arguments for '{cmd}' command", records
         except Exception as e:
-            result = f"-ERR {str(e)}"
-            return result, records
+            return f"-ERR {str(e)}", records
 
-        finally:
-            # 🚀 TRACK A: Stop the stopwatch, convert nanoseconds to microseconds
-            end_time = time.perf_counter_ns()
-            duration_us = (end_time - start_time) // 1000
-
-            # Feed the metrics to the store's slowlog ring buffer if available
-            if (
-                hasattr(self, "store")
-                and self.store
-                and hasattr(self.store, "_log_slow_command")
-            ):
-                self.store._log_slow_command(tokens, duration_us)
+    def _record_duration(self, tokens: list[str], start_time: int) -> None:
+        # 🚀 TRACK A: Stop the stopwatch, convert nanoseconds to microseconds
+        duration_us = (time.perf_counter_ns() - start_time) // 1000
+
+        # Feed the metrics to the store's slowlog ring buffer if available
+        if (
+            hasattr(self, "store")
+            and self.store
+            and hasattr(self.store, "_log_slow_command")
+        ):
+            self.store._log_slow_command(tokens, duration_us)
 
     # ---------------------------------------------------------
     # COMMAND HANDLERS (The isolated engine components)

kedis_python/main.py

@@ -447,7 +447,7 @@ class KedisClient:
             return True
 
         if cmd == "INFO":
-            version = "0.2.0"
+            version = "0.3.0"
             codename = "Echo"
 
             if self.local_mode:

kedis_python/parser.py

@@ -36,7 +36,7 @@ class CommandParser:
     _COUNT = re.compile(rb"(?:0|[1-9][0-9]*)")
 
     @staticmethod
-    def parse(raw_data, start: int = 0):
+    def parse(raw_data, start: int = 0, normalise: bool = True):
         """
         Parses incoming network data. Returns (tokens, bytes_consumed).
 
@@ -51,12 +51,16 @@ class CommandParser:
         raw_data may be bytes or a bytearray. Parsing starts at offset start and
         bytes_consumed counts from there, so a caller can work through a buffer
         holding many pipelined commands without copying or trimming it each time.
+
+        The first token is taken to be a command name and is uppercased. Pass
+        normalise=False for a frame that is data from its first token, such as the
+        baseline snapshot a leader sends: its JSON must reach the caller as sent.
         """
-        tokens, consumed, _ = CommandParser.parse_buffer(raw_data, start)
+        tokens, consumed, _ = CommandParser.parse_buffer(raw_data, start, normalise)
         return tokens, consumed
 
     @staticmethod
-    def parse_buffer(raw_data, start: int = 0):
+    def parse_buffer(raw_data, start: int = 0, normalise: bool = True):
         """parse(), plus a hint for incomplete frames: returns (tokens, consumed,
         need). When the next frame is incomplete, need is the length the bytes
         from start + consumed must reach before parsing can make any progress
@@ -75,7 +79,8 @@ class CommandParser:
                 # Only the command name is normalised. Every other token is data
                 # (keys, values, field names) until a handler says otherwise, so
                 # handlers match their own options case-insensitively.
-                tokens[0] = tokens[0].upper()
+                if normalise:
+                    tokens[0] = tokens[0].upper()
                 return tokens, end - start, 0
             pos = end  # blank frame: skip it and keep scanning
         return [], pos - start, 0
@@ -267,6 +272,11 @@ class KESPDecoder:
     MAX_DEPTH = 32  # nesting guard, so a hostile reply can't blow the stack
     MAX_LINE = 64 * 1024  # a status/error line longer than this is garbage
     MAX_HEADER_LINE = 32  # I<n> / S<n> / A<n> headers are never longer
+    # Whole-reply budgets. MAX_DEPTH and the per-array/per-string caps bound one
+    # node; these bound the *sum* over a reply, so many small nodes cannot add up
+    # to an unbounded amount of memory or work. Looked up at call time.
+    MAX_REPLY_ELEMENTS = 4_194_304  # every value in the reply, containers included
+    MAX_REPLY_PAYLOAD = 1024 * 1024 * 1024  # string and text-line bytes, summed
 
     _COUNT = re.compile(rb"(?:0|[1-9][0-9]*)")
     _INTEGER = re.compile(rb"(?:0|-?[1-9][0-9]*)")
@@ -275,7 +285,12 @@ class KESPDecoder:
     _NUMERIC_SIGILS = (b"I", b"S", b"A")
 
     @staticmethod
-    def decode(buf, pos: int = 0, _depth: int = 0, lenient: bool = False):
+    def decode(
+        buf, pos: int = 0, _depth: int = 0, lenient: bool = False, _budget=None
+    ):
+        if _budget is None:
+            # [elements left, payload bytes left], shared by the whole reply.
+            _budget = [KESPDecoder.MAX_REPLY_ELEMENTS, KESPDecoder.MAX_REPLY_PAYLOAD]
         if _depth > KESPDecoder.MAX_DEPTH:
             raise ReplyProtocolError("reply nested too deeply")
         if pos >= len(buf):
@@ -309,6 +324,16 @@ class KESPDecoder:
         after = nl + 1
         rest = line[1:]
 
+        # Charge this value to the whole-reply budget. A reply that is still
+        # incomplete never gets here, so a retry is not charged twice.
+        if _budget[0] <= 0:
+            raise ReplyProtocolError("reply has too many elements")
+        _budget[0] -= 1
+        if sigil not in (b"N", b"I", b"S", b"A"):  # line-carried text
+            if len(line) > _budget[1]:
+                raise ReplyProtocolError("reply payload too large")
+            _budget[1] -= len(line)
+
         if sigil == b"N":
             if rest and not lenient:
                 raise ReplyProtocolError("invalid nil reply")
@@ -326,6 +351,10 @@ class KESPDecoder:
             length = KESPDecoder._count(rest)
             if length > CommandParser.MAX_BULK_LENGTH:
                 raise ReplyProtocolError("invalid bulk length")
+            if length > _budget[1]:
+                # Refuse on the header, before waiting for the bytes to arrive.
+                raise ReplyProtocolError("reply payload too large")
+            _budget[1] -= length
             end = after + length
             if len(buf) < end + 1:
                 raise IncompleteReply(end + 1)
@@ -337,9 +366,14 @@ class KESPDecoder:
             count = KESPDecoder._count(rest)
             if count > CommandParser.MAX_ARRAY_LENGTH:
                 raise ReplyProtocolError("invalid array length")
+            if count > _budget[0]:
+                # Every element costs at least one; refuse on the header.
+                raise ReplyProtocolError("reply has too many elements")
             items = []
             for _ in range(count):
-                item, after = KESPDecoder.decode(buf, after, _depth + 1, lenient)
+                item, after = KESPDecoder.decode(
+                    buf, after, _depth + 1, lenient, _budget
+                )
                 items.append(item)
             return items, after
 
@@ -388,10 +422,81 @@ class KESPEncoder:
         """
         Translates raw Python objects from the router into strict KESP network bytes.
         """
+        # The same whole-reply budgets the decoder enforces, so a reply that
+        # leaves the server is one the client will accept.
+        budget = [KESPDecoder.MAX_REPLY_ELEMENTS, KESPDecoder.MAX_REPLY_PAYLOAD]
         try:
-            return KESPEncoder._encode(data, 0, True)
+            return KESPEncoder._encode(data, 0, True, budget)
         except ReplyEncodeError as exc:
-            return KESPEncoder._error_line(f"Internal server error: {exc}")
+            return KESPEncoder.failure_reply(exc)
+
+    @staticmethod
+    def failure_reply(exc) -> bytes:
+        """The one error line that stands in for a reply that could not be encoded."""
+        return KESPEncoder._error_line(f"Internal server error: {exc}")
+
+    @staticmethod
+    def encode_chunks(data, chunk: int = 512):
+        """encode(), in pieces: b"".join(encode_chunks(x)) == encode(x) for every
+        x that encode() can encode.
+
+        A reply with more than `chunk` elements comes out as its header and then
+        one piece per `chunk` elements, so a caller can hand control back to the
+        event loop between pieces instead of holding it for the whole reply.
+        Smaller replies, and replies that are not a flat top-level array or
+        dict, come out as a single piece.
+
+        Unlike encode() this raises ReplyEncodeError when the reply cannot be
+        encoded, possibly after some pieces were produced: the caller must hold
+        every piece back until the last one is made, then send either all of
+        them or failure_reply(exc), so that a failure never leaves half a frame
+        on the wire.
+        """
+        budget = [KESPDecoder.MAX_REPLY_ELEMENTS, KESPDecoder.MAX_REPLY_PAYLOAD]
+
+        if isinstance(data, (list, set, tuple)) and len(data) > chunk:
+            # The same checks _encode makes for a top-level array.
+            if budget[0] <= 0:
+                raise ReplyEncodeError("reply has too many elements")
+            budget[0] -= 1
+            if len(data) > CommandParser.MAX_ARRAY_LENGTH:
+                raise ReplyEncodeError("array too large")
+            if len(data) > budget[0]:
+                raise ReplyEncodeError("reply has too many elements")
+            yield f"A{len(data)}\n".encode("utf-8")
+            items = data if isinstance(data, (list, tuple)) else list(data)
+            for i in range(0, len(items), chunk):
+                yield b"".join(
+                    [
+                        KESPEncoder._encode(item, 1, False, budget)
+                        for item in items[i : i + chunk]
+                    ]
+                )
+
+        elif isinstance(data, dict) and len(data) > chunk:
+            # The same checks _encode makes for a top-level dict (a flat array of
+            # alternating keys and values).
+            if budget[0] <= 0:
+                raise ReplyEncodeError("reply has too many elements")
+            budget[0] -= 1
+            if len(data) * 2 > CommandParser.MAX_ARRAY_LENGTH:
+                raise ReplyEncodeError("array too large")
+            if len(data) * 2 > budget[0]:
+                raise ReplyEncodeError("reply has too many elements")
+            yield f"A{len(data) * 2}\n".encode("utf-8")
+            entries = list(data.items())
+            for i in range(0, len(entries), chunk):
+                parts = []
+                for k, v in entries[i : i + chunk]:
+                    if budget[0] <= 0:
+                        raise ReplyEncodeError("reply has too many elements")
+                    budget[0] -= 1  # the key
+                    parts.append(KESPEncoder._bulk(str(k), budget))
+                    parts.append(KESPEncoder._encode(v, 1, False, budget))
+                yield b"".join(parts)
+
+        else:
+            yield KESPEncoder._encode(data, 0, True, budget)
 
     @staticmethod
     def encode_command(tokens) -> bytes:
@@ -411,22 +516,39 @@ class KESPEncoder:
             raw = to_bytes(text)
         except UnicodeEncodeError:
             raw = text.encode("utf-8", "replace")
+        # Errors often echo client input (an unknown command's name, say). A line
+        # longer than the decoder's MAX_LINE would be rejected by our own client,
+        # so cut it, at a character boundary, and say so.
+        limit = KESPDecoder.MAX_LINE - 1  # the 'E' counts towards the line
+        if len(raw) > limit:
+            cut = limit - 3
+            while cut > 0 and raw[cut] & 0xC0 == 0x80:
+                cut -= 1  # raw[cut] would be mid-character: step back to its start
+            raw = raw[:cut] + b"..."
         return b"E" + raw + b"\n"
 
     @staticmethod
-    def _bulk(text: str) -> bytes:
+    def _bulk(text: str, budget=None) -> bytes:
         try:
             raw = to_bytes(text)
         except UnicodeEncodeError:
             raise ReplyEncodeError("string is not valid UTF-8") from None
         if len(raw) > CommandParser.MAX_BULK_LENGTH:
             raise ReplyEncodeError("string too large")
+        if budget is not None:
+            if len(raw) > budget[1]:
+                raise ReplyEncodeError("reply payload too large")
+            budget[1] -= len(raw)
         return b"S%d\n%s\n" % (len(raw), raw)
 
     @staticmethod
-    def _encode(data, depth: int, top: bool) -> bytes:
+    def _encode(data, depth: int, top: bool, budget) -> bytes:
         if depth > KESPDecoder.MAX_DEPTH:
             raise ReplyEncodeError("reply nested too deeply")
+        # One element per value, containers included (as the decoder counts).
+        if budget[0] <= 0:
+            raise ReplyEncodeError("reply has too many elements")
+        budget[0] -= 1
 
         # 1. Null / Missing Data
         if data is None:
@@ -450,7 +572,7 @@ class KESPEncoder:
         # 4. Strings & Status Messages
         if isinstance(data, str):
             if isinstance(data, BulkString) or not top:
-                return KESPEncoder._bulk(data)
+                return KESPEncoder._bulk(data, budget)
 
             if data in KESPEncoder._STATUSES:
                 val = data if data.startswith("+") else f"+{data}"
@@ -462,16 +584,18 @@ class KESPEncoder:
                     return KESPEncoder._error_line(data[len(prefix) + 1 :])
 
             # Otherwise, it's a standard Binary-Safe Bulk String
-            return KESPEncoder._bulk(data)
+            return KESPEncoder._bulk(data, budget)
 
         # 5. Arrays (Lists/Sets)
         if isinstance(data, (list, set, tuple)):
             if len(data) > CommandParser.MAX_ARRAY_LENGTH:
                 raise ReplyEncodeError("array too large")
+            if len(data) > budget[0]:
+                raise ReplyEncodeError("reply has too many elements")
             header = f"A{len(data)}\n".encode("utf-8")
             # Recursively encode every element inside the array
             elements = b"".join(
-                KESPEncoder._encode(item, depth + 1, False) for item in data
+                KESPEncoder._encode(item, depth + 1, False, budget) for item in data
             )
             return header + elements
 
@@ -480,11 +604,17 @@ class KESPEncoder:
             # A dictionary of 3 items becomes a KESP array of 6 items [k1, v1, k2, v2...]
             if len(data) * 2 > CommandParser.MAX_ARRAY_LENGTH:
                 raise ReplyEncodeError("array too large")
+            if len(data) * 2 > budget[0]:
+                raise ReplyEncodeError("reply has too many elements")
             header = f"A{len(data) * 2}\n".encode("utf-8")
-            elements = b"".join(
-                KESPEncoder._bulk(str(k)) + KESPEncoder._encode(v, depth + 1, False)
-                for k, v in data.items()
-            )
+            parts = []
+            for k, v in data.items():
+                if budget[0] <= 0:
+                    raise ReplyEncodeError("reply has too many elements")
+                budget[0] -= 1  # the key
+                parts.append(KESPEncoder._bulk(str(k), budget))
+                parts.append(KESPEncoder._encode(v, depth + 1, False, budget))
+            elements = b"".join(parts)
             return header + elements
 
         # Fallback