gateway.py

#!/usr/bin/env python3
"""
gateway.py — Lightweight HTTP gateway for the Observatory

Sits on port 7860 (the only port HF Spaces exposes) and routes:

  GET  /api/refresh   →  git pull all repos + recreate all views (on-demand)
  GET  /api/sources   →  return current sources.yaml as JSON
  *    /*             →  proxy everything else to ClickHouse on 8123

This lets you push to GitHub, then immediately:
  curl https://your-space.hf.space/api/refresh
  curl https://your-space.hf.space/?query=SELECT * FROM ohlc LIMIT 5
"""

import json
import os
import subprocess
import sys
import time
import urllib.request
import urllib.error
import yaml
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler

CLICKHOUSE_URL = "http://127.0.0.1:8123"
CONFIG_PATH = "/app/sources.yaml"
REFRESH_SCRIPT = "/app/refresh_sources.py"
GATEWAY_PORT = 7860


class GatewayHandler(BaseHTTPRequestHandler):

    # ── /api/refresh ─────────────────────────────────────────────────────
    def _handle_refresh(self):
        """Run refresh_sources.sh and return JSON result."""
        try:
            start = time.time()
            result = subprocess.run(
                ["python3", REFRESH_SCRIPT, "full"],
                capture_output=True, text=True, timeout=1200,
            )
            elapsed = round(time.time() - start, 2)

            body = {
                "action": "refresh",
                "elapsed_seconds": elapsed,
                "exit_code": result.returncode,
            }

            # Try to parse script's JSON output
            try:
                body["result"] = json.loads(result.stdout)
            except (json.JSONDecodeError, ValueError):
                body["stdout"] = result.stdout

            if result.stderr:
                body["stderr"] = result.stderr

            self._json_response(200, body)

        except subprocess.TimeoutExpired:
            self._json_response(504, {"error": "refresh timed out after 1200s"})
        except Exception as e:
            self._json_response(500, {"error": str(e)})

    # ── /api/sources ─────────────────────────────────────────────────────
    def _handle_sources(self):
        """Return current sources.yaml as JSON."""
        try:
            with open(CONFIG_PATH, "r") as f:
                config = yaml.safe_load(f)
            self._json_response(200, config)
        except Exception as e:
            self._json_response(500, {"error": str(e)})

    # ── Proxy to ClickHouse ──────────────────────────────────────────────
    def _proxy_to_clickhouse(self):
        """Forward the request to ClickHouse HTTP interface."""
        # Read request body if present
        content_length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(content_length) if content_length else None

        url = f"{CLICKHOUSE_URL}{self.path}"

        req = urllib.request.Request(url, data=body, method=self.command)

        # Forward relevant headers
        for header in self.headers:
            lower = header.lower()
            if lower not in ("host", "content-length", "transfer-encoding"):
                req.add_header(header, self.headers[header])

        try:
            resp = urllib.request.urlopen(req, timeout=600)
            self.send_response(resp.status)
            # Forward response headers
            for key, val in resp.headers.items():
                if key.lower() != "transfer-encoding":
                    self.send_header(key, val)
            # Add CORS headers for browser-based SQL playgrounds
            self.send_header("Access-Control-Allow-Origin", "*")
            self.end_headers()
            # Stream response body
            while True:
                chunk = resp.read(65536)
                if not chunk:
                    break
                self.wfile.write(chunk)

        except urllib.error.HTTPError as e:
            self.send_response(e.code)
            for key, val in e.headers.items():
                if key.lower() != "transfer-encoding":
                    self.send_header(key, val)
            self.send_header("Access-Control-Allow-Origin", "*")
            self.end_headers()
            self.wfile.write(e.read())

        except Exception as e:
            self._json_response(502, {"error": f"ClickHouse unreachable: {e}"})

    # ── HTTP method handlers ─────────────────────────────────────────────
    def do_GET(self):
        if self.path == "/api/refresh":
            self._handle_refresh()
        elif self.path == "/api/sources":
            self._handle_sources()
        else:
            self._proxy_to_clickhouse()

    def do_POST(self):
        self._proxy_to_clickhouse()

    def do_OPTIONS(self):
        """CORS preflight."""
        self.send_response(204)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "*")
        self.end_headers()

    # ── Helpers ──────────────────────────────────────────────────────────
    def _json_response(self, code, data):
        body = json.dumps(data, indent=2).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, format, *args):
        """Cleaner log format. Ignores HF container ping logs for cleanliness."""
        msg = format % args
        if "?logs=container" in msg:
            return
        sys.stderr.write(f"[gateway] {self.address_string()} {msg}\n")


def wait_for_clickhouse(timeout=60):
    """Block until ClickHouse is responding on 8123."""
    print(f"[gateway] Waiting for ClickHouse on {CLICKHOUSE_URL}...")
    for i in range(timeout):
        try:
            urllib.request.urlopen(f"{CLICKHOUSE_URL}/ping", timeout=2)
            print(f"[gateway] ClickHouse is ready (took {i+1}s)")
            return True
        except Exception:
            time.sleep(1)
    print(f"[gateway] WARNING: ClickHouse not ready after {timeout}s, starting anyway")
    return False


if __name__ == "__main__":
    # Wait for ClickHouse to be ready before accepting traffic
    wait_for_clickhouse()

    # Run initial full refresh (clone repos + create views)
    print("[gateway] Running initial data sync + view creation...")
    subprocess.run(
        ["python3", REFRESH_SCRIPT, "full"],
        timeout=1200,
    )

    # Start gateway
    server = ThreadingHTTPServer(("0.0.0.0", GATEWAY_PORT), GatewayHandler)
    print(f"[gateway] Listening on port {GATEWAY_PORT}")
    print(f"[gateway]   /api/refresh  → on-demand git pull + view recreation")
    print(f"[gateway]   /api/sources  → current config as JSON")
    print(f"[gateway]   /*            → proxy to ClickHouse")
    server.serve_forever()