refactor: extract KnowledgeBaseHandler to separate file

Author: mdesmet

src/datapilot/core/knowledge/cli.py

@@ -1,13 +1,5 @@
-import json
-import re
-from http.server import BaseHTTPRequestHandler
 from http.server import HTTPServer
-from urllib.error import HTTPError
-from urllib.error import URLError
-from urllib.parse import urlparse
-from urllib.request import Request
-from urllib.request import urlopen
-
+from .server import KnowledgeBaseHandler
 import click
 
 
@@ -32,72 +24,10 @@ def serve(ctx, port):
         )
         ctx.exit(1)
 
-    class KnowledgeBaseHandler(BaseHTTPRequestHandler):
-        def do_GET(self):
-            """Handle GET requests."""
-            path = urlparse(self.path).path
-
-            # Match /knowledge_bases/{uuid} pattern
-            match = re.match(r"^/knowledge_bases/([a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12})$", path)
-
-            if match:
-                public_id = match.group(1)
-                self.handle_knowledge_base(public_id)
-            elif path == "/health":
-                self.handle_health()
-            else:
-                self.send_error(404, "Not Found")
-
-        def handle_knowledge_base(self, public_id):
-            """Fetch and return knowledge base data."""
-            url = f"{backend_url}/knowledge_bases/public/{public_id}"
-
-            # Validate URL scheme for security
-            parsed_url = urlparse(url)
-            if parsed_url.scheme not in ("http", "https"):
-                self.send_response(400)
-                self.send_header("Content-Type", "application/json")
-                self.end_headers()
-                error_msg = json.dumps({"error": "Invalid URL scheme. Only HTTP and HTTPS are allowed."})
-                self.wfile.write(error_msg.encode("utf-8"))
-                return
-
-            headers = {"Authorization": f"Bearer {token}", "X-Tenant": instance_name, "Content-Type": "application/json"}
-
-            req = Request(url, headers=headers)  # noqa: S310
-
-            try:
-                # URL scheme validated above - only HTTP/HTTPS allowed
-                with urlopen(req, timeout=30) as response:  # noqa: S310
-                    data = response.read()
-                    self.send_response(200)
-                    self.send_header("Content-Type", "application/json")
-                    self.end_headers()
-                    self.wfile.write(data)
-            except HTTPError as e:
-                error_body = e.read()
-                error_data = error_body.decode("utf-8") if error_body else '{"error": "HTTP Error"}'
-                self.send_response(e.code)
-                self.send_header("Content-Type", "application/json")
-                self.end_headers()
-                self.wfile.write(error_data.encode("utf-8"))
-            except URLError as e:
-                self.send_response(500)
-                self.send_header("Content-Type", "application/json")
-                self.end_headers()
-                error_msg = json.dumps({"error": str(e)})
-                self.wfile.write(error_msg.encode("utf-8"))
-
-        def handle_health(self):
-            """Handle health check endpoint."""
-            self.send_response(200)
-            self.send_header("Content-Type", "application/json")
-            self.end_headers()
-            self.wfile.write(json.dumps({"status": "ok"}).encode("utf-8"))
-
-        def log_message(self, format, *args):
-            """Override to use click.echo for logging."""
-            click.echo(f"{self.address_string()} - {format % args}")
+    # Set context data for the handler
+    KnowledgeBaseHandler.token = token
+    KnowledgeBaseHandler.instance_name = instance_name
+    KnowledgeBaseHandler.backend_url = backend_url
 
     server_address = ("", port)
     httpd = HTTPServer(server_address, KnowledgeBaseHandler)

src/datapilot/core/knowledge/server.py

@@ -0,0 +1,83 @@
+import json
+import re
+from http.server import BaseHTTPRequestHandler
+from urllib.error import HTTPError
+from urllib.error import URLError
+from urllib.parse import urlparse
+from urllib.request import Request
+from urllib.request import urlopen
+
+import click
+
+class KnowledgeBaseHandler(BaseHTTPRequestHandler):
+    """HTTP request handler for serving knowledge bases and health checks."""
+    
+    token: str = ""
+    instance_name: str = ""
+    backend_url: str = ""
+
+    def do_GET(self):
+        """Handle GET requests."""
+        path = urlparse(self.path).path
+
+        # Match /knowledge_bases/{uuid} pattern
+        match = re.match(r"^/knowledge_bases/([a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12})$", path)
+
+        if match:
+            public_id = match.group(1)
+            self.handle_knowledge_base(public_id)
+        elif path == "/health":
+            self.handle_health()
+        else:
+            self.send_error(404, "Not Found")
+
+    def handle_knowledge_base(self, public_id):
+        """Fetch and return knowledge base data."""
+        url = f"{self.backend_url}/knowledge_bases/public/{public_id}"
+
+        # Validate URL scheme for security
+        parsed_url = urlparse(url)
+        if parsed_url.scheme not in ("http", "https"):
+            self.send_response(400)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+            error_msg = json.dumps({"error": "Invalid URL scheme. Only HTTP and HTTPS are allowed."})
+            self.wfile.write(error_msg.encode("utf-8"))
+            return
+
+        headers = {"Authorization": f"Bearer {self.token}", "X-Tenant": self.instance_name, "Content-Type": "application/json"}
+
+        req = Request(url, headers=headers)  # noqa: S310
+
+        try:
+            # URL scheme validated above - only HTTP/HTTPS allowed
+            with urlopen(req, timeout=30) as response:  # noqa: S310
+                data = response.read()
+                self.send_response(200)
+                self.send_header("Content-Type", "application/json")
+                self.end_headers()
+                self.wfile.write(data)
+        except HTTPError as e:
+            error_body = e.read()
+            error_data = error_body.decode("utf-8") if error_body else '{"error": "HTTP Error"}'
+            self.send_response(e.code)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+            self.wfile.write(error_data.encode("utf-8"))
+        except URLError as e:
+            self.send_response(500)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+            error_msg = json.dumps({"error": str(e)})
+            self.wfile.write(error_msg.encode("utf-8"))
+
+    def handle_health(self):
+        """Handle health check endpoint."""
+        self.send_response(200)
+        self.send_header("Content-Type", "application/json")
+        self.end_headers()
+        self.wfile.write(json.dumps({"status": "ok"}).encode("utf-8"))
+
+    def log_message(self, format, *args):
+        """Override to use click.echo for logging."""
+        click.echo(f"{self.address_string()} - {format % args}")