rafapolo
beba7dc9ec
After init, a background thread runs COUNT(*) on the three most-queried tables so the first real user query hits a warm Parquet metadata cache.
163 lines
6.0 KiB
Python
163 lines
6.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Minimal cookie-session auth gate for DuckDB shell."""
|
|
import decimal, datetime, duckdb, hmac, hashlib, json, os, secrets, threading, time
|
|
from http.server import HTTPServer, BaseHTTPRequestHandler
|
|
from socketserver import ThreadingMixIn
|
|
from urllib.parse import parse_qs
|
|
|
|
PASSWORD = os.environ.get('BASIC_AUTH_PASSWORD', '').encode()
|
|
|
|
_con = None
|
|
_lock = threading.Lock()
|
|
|
|
def _init_db():
|
|
global _con
|
|
endpoint = os.environ.get('HETZNER_S3_ENDPOINT', '').removeprefix('https://').removeprefix('http://')
|
|
_con = duckdb.connect(':memory:')
|
|
_con.execute("INSTALL httpfs; LOAD httpfs;")
|
|
_con.execute(f"""
|
|
SET s3_endpoint='{endpoint}';
|
|
SET s3_access_key_id='{os.environ.get("AWS_ACCESS_KEY_ID", "")}';
|
|
SET s3_secret_access_key='{os.environ.get("AWS_SECRET_ACCESS_KEY", "")}';
|
|
SET s3_region='{os.environ.get("BUCKET_REGION", "")}';
|
|
SET s3_url_style='path';
|
|
SET enable_object_cache=true;
|
|
SET threads=4;
|
|
SET memory_limit='6GB';
|
|
SET preserve_insertion_order=false;
|
|
SET http_keep_alive=true;
|
|
SET http_retries=3;
|
|
""")
|
|
_con.execute("ATTACH '/app/data/basedosdados.duckdb' AS basedosdados (READ_ONLY)")
|
|
threading.Thread(target=_warm_cache, daemon=True).start()
|
|
|
|
def _warm_cache():
|
|
hot_tables = [
|
|
"br_tse_eleicoes.candidatos",
|
|
"br_tse_eleicoes.despesas_candidato",
|
|
"br_tse_eleicoes.resultados_candidato",
|
|
]
|
|
for t in hot_tables:
|
|
try:
|
|
with _lock:
|
|
_con.execute(f"SELECT COUNT(*) FROM basedosdados.{t}")
|
|
except Exception:
|
|
pass
|
|
|
|
def _json_default(obj):
|
|
if isinstance(obj, decimal.Decimal): return float(obj)
|
|
if isinstance(obj, (datetime.date, datetime.datetime)): return obj.isoformat()
|
|
return str(obj)
|
|
|
|
def _run_query(sql, json_mode=True):
|
|
with _lock:
|
|
try:
|
|
rel = _con.execute(sql)
|
|
cols = [d[0] for d in rel.description]
|
|
rows = [{cols[i]: row[i] for i in range(len(cols))} for row in rel.fetchall()]
|
|
return json.dumps(rows, default=_json_default).encode()
|
|
except Exception as e:
|
|
return json.dumps({'error': str(e)}).encode()
|
|
|
|
_SECRET = secrets.token_bytes(32)
|
|
|
|
def _make_token():
|
|
day = str(int(time.time()) // 86400)
|
|
return hmac.new(_SECRET, day.encode(), hashlib.sha256).hexdigest()
|
|
|
|
def _valid(token):
|
|
if not token:
|
|
return False
|
|
for delta in (0, 1):
|
|
day = str(int(time.time()) // 86400 - delta)
|
|
expected = hmac.new(_SECRET, day.encode(), hashlib.sha256).hexdigest()
|
|
if hmac.compare_digest(token, expected):
|
|
return True
|
|
return False
|
|
|
|
LOGIN_HTML = """<!DOCTYPE html>
|
|
<html><head><title>DB Shell</title><style>
|
|
body{display:flex;justify-content:center;align-items:center;min-height:100vh;margin:0;background:#0f1117;font-family:sans-serif}
|
|
form{background:#1a1d27;padding:2rem;border-radius:8px;display:flex;flex-direction:column;gap:1rem;min-width:280px}
|
|
h2{color:#fff;margin:0}
|
|
input{padding:.6rem;border-radius:4px;border:1px solid #333;background:#0f1117;color:#fff;font-size:1rem}
|
|
button{padding:.6rem;border-radius:4px;border:none;background:#f4c543;color:#000;font-size:1rem;cursor:pointer;font-weight:600}
|
|
</style></head>
|
|
<body><form method="POST" action="/login">
|
|
<h2>DB Shell</h2>
|
|
<input type="password" name="password" placeholder="Password" autofocus>
|
|
<button type="submit">Enter</button>
|
|
</form></body></html>""".encode()
|
|
|
|
class H(BaseHTTPRequestHandler):
|
|
def _cookie(self):
|
|
for part in self.headers.get('Cookie', '').split(';'):
|
|
part = part.strip()
|
|
if part.startswith('ddb_auth='):
|
|
return part[9:]
|
|
return ''
|
|
|
|
def do_GET(self):
|
|
if self.path == '/auth':
|
|
if _valid(self._cookie()):
|
|
self._resp(200)
|
|
else:
|
|
self.send_response(302)
|
|
self.send_header('Location', '/login')
|
|
self.end_headers()
|
|
elif self.path.startswith('/query'):
|
|
from urllib.parse import urlparse, parse_qs
|
|
pwd = self.headers.get('X-Password', '').encode()
|
|
if not hmac.compare_digest(pwd, PASSWORD):
|
|
self._resp(401, b'Unauthorized\n')
|
|
return
|
|
qs = parse_qs(urlparse(self.path).query)
|
|
sql = qs.get('q', [''])[0]
|
|
if not sql:
|
|
self._resp(400, b'missing ?q=\n')
|
|
return
|
|
self._resp(200, _run_query(sql), 'application/json')
|
|
else:
|
|
self._resp(200, LOGIN_HTML, 'text/html; charset=utf-8')
|
|
|
|
def do_POST(self):
|
|
if self.path == '/query':
|
|
pwd = self.headers.get('X-Password', '').encode()
|
|
if not hmac.compare_digest(pwd, PASSWORD):
|
|
self._resp(401, b'Unauthorized\n')
|
|
return
|
|
sql = self.rfile.read(int(self.headers.get('Content-Length', 0))).decode(errors='replace')
|
|
self._resp(200, _run_query(sql), 'application/json')
|
|
return
|
|
body = self.rfile.read(int(self.headers.get('Content-Length', 0))).decode(errors='replace')
|
|
pwd = parse_qs(body).get('password', [''])[0].encode()
|
|
if hmac.compare_digest(pwd, PASSWORD):
|
|
self.send_response(302)
|
|
self.send_header('Set-Cookie', f'ddb_auth={_make_token()}; Path=/; HttpOnly; SameSite=Strict')
|
|
self.send_header('Location', '/')
|
|
self.end_headers()
|
|
else:
|
|
self._resp(200, LOGIN_HTML, 'text/html; charset=utf-8')
|
|
|
|
def do_OPTIONS(self):
|
|
self.send_response(204)
|
|
self.end_headers()
|
|
|
|
def _resp(self, code, body=b'', ct='text/plain'):
|
|
self.send_response(code)
|
|
if body:
|
|
self.send_header('Content-Type', ct)
|
|
self.send_header('Content-Length', str(len(body)))
|
|
self.end_headers()
|
|
if body:
|
|
self.wfile.write(body)
|
|
|
|
def log_message(self, *_):
|
|
pass
|
|
|
|
class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
|
|
daemon_threads = True
|
|
|
|
_init_db()
|
|
ThreadedHTTPServer(('127.0.0.1', 8081), H).serve_forever()
|