[typedb] forwarding state machine

README.md

@@ -75,7 +75,7 @@ You can install the respective extension by calling `localstack extensions insta
 | [Miniflare](https://github.com/localstack/localstack-extensions/tree/main/miniflare)               | localstack-extension-miniflare | 0.1.0   | Experimental |
 | [Stripe](https://github.com/localstack/localstack-extensions/tree/main/stripe)                     | localstack-extension-stripe | 0.2.0   | Stable |
 | [Terraform Init](https://github.com/localstack/localstack-extensions/tree/main/terraform-init)     | localstack-extension-terraform-init | 0.2.0   | Experimental |
-| [TypeDB](https://github.com/localstack/localstack-extensions/tree/main/typedb)                     | localstack-extension-typedb | 0.1.2   | Experimental |
+| [TypeDB](https://github.com/localstack/localstack-extensions/tree/main/typedb)                     | localstack-extension-typedb | 0.1.3   | Experimental |
 
 
 ## Developing Extensions

typedb/localstack_typedb/utils/h2_proxy.py

@@ -1,5 +1,6 @@
 import logging
 import socket
+from enum import Enum
 from typing import Iterable, Callable
 
 from h2.frame_buffer import FrameBuffer
@@ -63,6 +64,11 @@ def apply_http2_patches_for_grpc_support(
     )
     patched_connection = True
 
+    class ForwardingState(Enum):
+        UNDECIDED = "undecided"
+        FORWARDING = "forwarding"
+        PASSTHROUGH = "passthrough"
+
     class ForwardingBuffer:
         """
         A buffer atop the HTTP2 client connection, that will hold
@@ -72,7 +78,7 @@ class ForwardingBuffer:
 
         backend: TcpForwarder
         buffer: list
-        proxying: bool | None
+        state: ForwardingState
 
         def __init__(self, http_response_stream):
             self.http_response_stream = http_response_stream
@@ -81,7 +87,7 @@ def __init__(self, http_response_stream):
             )
             self.backend = TcpForwarder(target_port, host=target_host)
             self.buffer = []
-            self.proxying = None
+            self.state = ForwardingState.UNDECIDED
             reactor.getThreadPool().callInThread(
                 self.backend.receive_loop, self.received_from_backend
             )
@@ -91,35 +97,30 @@ def received_from_backend(self, data):
             self.http_response_stream.write(data)
 
         def received_from_http2_client(self, data, default_handler: Callable):
-            if self.proxying is False:
-                # Note: Return here only if `proxying` is `False` (a value of `None` indicates
-                # that the headers have not fully been received yet)
-                return default_handler(data)
-
-            if self.proxying:
-                assert not self.buffer
-                # Keep sending data to the backend for the lifetime of this connection
-                self.backend.send(data)
-                return
-
-            self.buffer.append(data)
-
-            if not (headers := get_headers_from_data_stream(self.buffer)):
-                # If no headers received yet, then return (method will be called again for next chunk of data)
-                return
-
-            self.proxying = should_proxy_request(headers)
-
-            buffered_data = b"".join(self.buffer)
-            self.buffer = []
-
-            if not self.proxying:
-                # if this is not a target request, then call the default handler
-                default_handler(buffered_data)
-                return
-
-            LOG.debug(f"Forwarding {len(buffered_data)} bytes to backend")
-            self.backend.send(buffered_data)
+            match self.state:
+                case ForwardingState.PASSTHROUGH:
+                    default_handler(data)
+                case ForwardingState.FORWARDING:
+                    assert not self.buffer
+                    # Keep sending data to the backend for the lifetime of this connection
+                    self.backend.send(data)
+                case ForwardingState.UNDECIDED:
+                    self.buffer.append(data)
+
+                    if headers := get_headers_from_data_stream(self.buffer):
+                        buffered_data = b"".join(self.buffer)
+                        self.buffer = []
+
+                        if should_proxy_request(headers):
+                            self.state = ForwardingState.FORWARDING
+                            LOG.debug(
+                                f"Forwarding {len(buffered_data)} bytes to backend"
+                            )
+                            self.backend.send(buffered_data)
+                        else:
+                            self.state = ForwardingState.PASSTHROUGH
+                            # if this is not a target request, then call the default handler
+                            default_handler(buffered_data)
 
         def close(self):
             self.backend.close()

typedb/pyproject.toml

@@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "localstack-extension-typedb"
-version = "0.1.2"
+version = "0.1.3"
 description = "LocalStack Extension: TypeDB on LocalStack"
 readme = {file = "README.md", content-type = "text/markdown; charset=UTF-8"}
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 authors = [
     { name = "LocalStack + TypeDB team"}
 ]