add durable jobs and live job monitor
This commit is contained in:
+111
-2
@@ -1,5 +1,5 @@
|
||||
from __future__ import annotations
|
||||
import argparse, hashlib, json, os, re, secrets, sqlite3, sys
|
||||
import argparse, hashlib, json, os, re, secrets, sqlite3, sys, threading, time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from http.cookies import SimpleCookie
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
@@ -15,8 +15,23 @@ SCHEMA = Path(__file__).resolve().parents[1] / "schema.sql"
|
||||
SESSION_DAYS = 7
|
||||
PBKDF2_ITERATIONS = 300_000
|
||||
MUTATING_ROLES = {"owner", "admin", "researcher"}
|
||||
JOB_TYPES = {"noop", "prospect_recalculate"}
|
||||
JOB_PAGE_SIZE = 100
|
||||
SECRET_KEYS = {"password", "passwd", "secret", "token", "api_key", "apikey", "authorization", "credential", "private_key"}
|
||||
CHILD_TABLES = {"contacts": ("name", "email", "phone", "title", "do_not_contact"), "domains": ("domain", "kind"), "websites": ("url", "website_class"), "evidence": ("kind", "url", "claim"), "notes": ("body",)}
|
||||
|
||||
def redact(value):
|
||||
if isinstance(value, dict): return {k: ("[REDACTED]" if str(k).lower() in SECRET_KEYS or any(s in str(k).lower() for s in ("password", "token", "secret", "api_key")) else redact(v)) for k,v in value.items()}
|
||||
if isinstance(value, list): return [redact(v) for v in value[:100]]
|
||||
if isinstance(value, str): return value[:2000]
|
||||
return value
|
||||
|
||||
def job_json(row):
|
||||
result = row_json(row)
|
||||
try: result["payload"] = json.loads(result.get("payload") or "{}")
|
||||
except (ValueError, TypeError): result["payload"] = {}
|
||||
return result
|
||||
|
||||
def hash_password(password: str, salt: bytes | None = None) -> tuple[str, str]:
|
||||
salt = salt or secrets.token_bytes(16); return hashlib.pbkdf2_hmac("sha256", password.encode(), salt, PBKDF2_ITERATIONS).hex(), salt.hex()
|
||||
def verify_password(password, encoded_hash, encoded_salt):
|
||||
@@ -95,6 +110,8 @@ class ApiHandler(BaseHTTPRequestHandler):
|
||||
if path=="/api/v1/dashboard/summary":
|
||||
row=db.execute("SELECT COUNT(*) businesses,COALESCE(AVG(score),0) average_score FROM businesses WHERE organization_id=?",(org,)).fetchone(); return self.send_json(200,{"organization_id":org,"businesses":row["businesses"],"average_score":round(row["average_score"],2),"suppressed":db.execute("SELECT COUNT(*) FROM suppressions WHERE organization_id=?",(org,)).fetchone()[0]})
|
||||
if path=="/api/v1/businesses": return self.list_businesses(db,org,parse_qs(parsed.query))
|
||||
if path=="/api/v1/jobs": return self.list_jobs(db,org,parse_qs(parsed.query))
|
||||
if path.startswith("/api/v1/jobs/"): return self.get_job_route(db,org,path,parse_qs(parsed.query))
|
||||
if path.startswith("/api/v1/businesses/"):
|
||||
bits=path.split("/"); ident=bits[4] if len(bits)>4 else ""
|
||||
if not ident.isdigit(): return self.send_json(404,{"error":"not_found"})
|
||||
@@ -103,6 +120,59 @@ class ApiHandler(BaseHTTPRequestHandler):
|
||||
payload=row_json(row); payload.update(self.nested(db,int(ident),org)); return self.send_json(200,payload)
|
||||
return self.send_json(404,{"error":"not_found"})
|
||||
finally: db.close()
|
||||
def list_jobs(self, db, org, query):
|
||||
try: limit=max(1,min(int(query.get("page_size",[50])[0]),JOB_PAGE_SIZE)); offset=max(0,int(query.get("offset",[0])[0]))
|
||||
except (ValueError, TypeError): return self.send_json(400,{"error":"invalid_pagination"})
|
||||
rows=db.execute("SELECT * FROM jobs WHERE organization_id=? ORDER BY id DESC LIMIT ? OFFSET ?",(org,limit+1,offset)).fetchall(); more=len(rows)>limit
|
||||
return self.send_json(200,{"organization_id":org,"items":[job_json(r) for r in rows[:limit]],"limit":limit,"offset":offset,"has_more":more})
|
||||
|
||||
def get_job_route(self, db, org, path, query):
|
||||
bits=path.split("/")
|
||||
if len(bits)<5 or not bits[4].isdigit(): return self.send_json(404,{"error":"not_found"})
|
||||
job=db.execute("SELECT * FROM jobs WHERE id=? AND organization_id=?",(int(bits[4]),org)).fetchone()
|
||||
if not job:return self.send_json(404,{"error":"not_found"})
|
||||
if len(bits)==5:return self.send_json(200,job_json(job))
|
||||
if len(bits)==6 and bits[5]=="events":
|
||||
try: after=max(0,int(query.get("after",[0])[0]))
|
||||
except (ValueError,TypeError): return self.send_json(400,{"error":"invalid_sequence"})
|
||||
events=[row_json(r) for r in db.execute("SELECT * FROM job_events WHERE job_id=? AND organization_id=? AND sequence>? ORDER BY sequence",(job["id"],org,after))]
|
||||
return self.send_json(200,{"items":events,"after":after})
|
||||
if len(bits)==7 and bits[5]=="events" and bits[6]=="stream":
|
||||
try: after=max(0,int(query.get("after",[0])[0]))
|
||||
except (ValueError,TypeError): return self.send_json(400,{"error":"invalid_sequence"})
|
||||
events=[row_json(r) for r in db.execute("SELECT * FROM job_events WHERE job_id=? AND organization_id=? AND sequence>? ORDER BY sequence",(job["id"],org,after))]
|
||||
body=b"".join((b"event: "+str(e["event_type"]).encode()+b"\\ndata: "+json.dumps(e,sort_keys=True).encode()+b"\\n\\n") for e in events)
|
||||
self.send_response(200);self.send_header("Content-Type","text/event-stream");self.send_header("Cache-Control","no-cache");self.send_header("Content-Length",str(len(body)));self.end_headers();self.wfile.write(body);return
|
||||
return self.send_json(404,{"error":"not_found"})
|
||||
|
||||
def create_job(self, payload, db, user):
|
||||
kind=str(payload.get("type","")).strip(); key=str(payload.get("idempotency_key","")).strip(); data=payload.get("payload",{})
|
||||
if kind not in JOB_TYPES:return self.send_json(400,{"error":"invalid_job_type"})
|
||||
if not key or len(key)>200 or not isinstance(data,dict):return self.send_json(400,{"error":"invalid_job_request"})
|
||||
safe=json.dumps(redact(data),sort_keys=True,separators=(",",":")); max_attempts=max(1,min(int(payload.get("max_attempts",3)),5)) if str(payload.get("max_attempts",3)).isdigit() else 3
|
||||
try:
|
||||
cur=db.execute("INSERT INTO jobs(organization_id,idempotency_key,type,payload,max_attempts) VALUES(?,?,?,?,?)",(user["organization_id"],key,kind,safe,max_attempts)); jid=cur.lastrowid
|
||||
self.add_job_event(db,jid,user["organization_id"],"queued","Job queued",0); self.audit(db,user,"job.created",str(jid)); db.commit(); getattr(self.server,"job_wakeup",threading.Event()).set()
|
||||
return self.send_json(201,job_json(db.execute("SELECT * FROM jobs WHERE id=?",(jid,)).fetchone()))
|
||||
except sqlite3.IntegrityError:
|
||||
row=db.execute("SELECT * FROM jobs WHERE organization_id=? AND idempotency_key=?",(user["organization_id"],key)).fetchone(); return self.send_json(200,job_json(row))
|
||||
|
||||
def add_job_event(self, db, jid, org, event_type, message, progress=0, error_code=None):
|
||||
seq=db.execute("SELECT COALESCE(MAX(sequence),0)+1 FROM job_events WHERE job_id=?",(jid,)).fetchone()[0]
|
||||
db.execute("INSERT INTO job_events(job_id,organization_id,sequence,event_type,message,progress,error_code) VALUES(?,?,?,?,?,?,?)",(jid,org,seq,event_type,str(message)[:500],progress,error_code)); return seq
|
||||
|
||||
def job_action(self, db, user, path):
|
||||
bits=path.split("/"); jid=int(bits[4]) if len(bits)>4 and bits[4].isdigit() else -1; action=bits[5] if len(bits)>5 else ""
|
||||
job=db.execute("SELECT * FROM jobs WHERE id=? AND organization_id=?",(jid,user["organization_id"])).fetchone()
|
||||
if not job:return self.send_json(404,{"error":"not_found"})
|
||||
if action=="cancel":
|
||||
if job["status"] in ("queued","running"): db.execute("UPDATE jobs SET status='cancelled',completed_at=CURRENT_TIMESTAMP,updated_at=CURRENT_TIMESTAMP WHERE id=?",(jid,));self.add_job_event(db,jid,user["organization_id"],"cancelled","Job cancelled",job["progress"])
|
||||
self.audit(db,user,"job.cancelled",str(jid));db.commit();return self.send_json(200,job_json(db.execute("SELECT * FROM jobs WHERE id=?",(jid,)).fetchone()))
|
||||
if action=="retry":
|
||||
if job["status"]!="failed":return self.send_json(409,{"error":"job_not_failed"})
|
||||
db.execute("UPDATE jobs SET status='queued',error_code=NULL,completed_at=NULL,updated_at=CURRENT_TIMESTAMP WHERE id=?",(jid,));self.add_job_event(db,jid,user["organization_id"],"retry","Job retry queued",job["progress"]);self.audit(db,user,"job.retried",str(jid));db.commit();getattr(self.server,"job_wakeup",threading.Event()).set();return self.send_json(200,job_json(db.execute("SELECT * FROM jobs WHERE id=?",(jid,)).fetchone()))
|
||||
return self.send_json(404,{"error":"not_found"})
|
||||
|
||||
def list_businesses(self,db,org,query):
|
||||
def number(name, default=None):
|
||||
raw=query.get(name,[None])[0]
|
||||
@@ -132,8 +202,13 @@ class ApiHandler(BaseHTTPRequestHandler):
|
||||
c=SimpleCookie();c.load(self.headers.get("Cookie",""));t=c.get("session");
|
||||
if t:db.execute("DELETE FROM sessions WHERE token_hash=?",(hashlib.sha256(t.value.encode()).hexdigest(),))
|
||||
self.audit(db,user,"logout");db.commit();return self.send_json(200,{"ok":True},{"Set-Cookie":self.auth_cookie("",0)})
|
||||
if path=="/api/v1/jobs":
|
||||
if user["role"] not in MUTATING_ROLES:return self.send_json(403,{"error":"forbidden"})
|
||||
return self.create_job(self.read_json(),db,user)
|
||||
if user["role"] not in MUTATING_ROLES:return self.send_json(403,{"error":"forbidden"})
|
||||
payload=self.read_json(); org=user["organization_id"]
|
||||
if path.startswith("/api/v1/jobs/"):
|
||||
return self.job_action(db,user,path)
|
||||
if path=="/api/v1/businesses":return self.create_business(payload,db,user)
|
||||
if path=="/api/v1/suppressions":return self.create_suppression(payload,db,user)
|
||||
if path=="/api/v1/imports/preview":return self.preview_import(payload,db,org)
|
||||
@@ -213,8 +288,42 @@ class ApiHandler(BaseHTTPRequestHandler):
|
||||
return self.send_json(200,{"accepted":len(accepted),"duplicates":len(rows)-len(normalized)+len(normalized)-len(accepted)-suppressed,"suppressed":suppressed,"rows":accepted})
|
||||
def log_message(self,*_):pass
|
||||
|
||||
def _job_worker(server):
|
||||
while not server.job_stop.is_set():
|
||||
db=connect(server.db_path)
|
||||
try:
|
||||
job=db.execute("SELECT * FROM jobs WHERE status='queued' ORDER BY id LIMIT 1").fetchone()
|
||||
if not job:
|
||||
db.close(); server.job_wakeup.wait(.1); server.job_wakeup.clear(); continue
|
||||
changed=db.execute("UPDATE jobs SET status='running',attempts=attempts+1,started_at=COALESCE(started_at,CURRENT_TIMESTAMP),updated_at=CURRENT_TIMESTAMP WHERE id=? AND status='queued'",(job["id"],)).rowcount
|
||||
if not changed: db.close(); continue
|
||||
db.commit(); org=job["organization_id"]; jid=job["id"]; server_handler=object.__new__(ApiHandler)
|
||||
server_handler.add_job_event(db,jid,org,"started","Job started",0); db.commit()
|
||||
try: payload=json.loads(job["payload"] or "{}")
|
||||
except ValueError: payload={}
|
||||
try: steps=1 if job["type"]=="noop" else max(1,min(int(payload.get("steps",5)),20))
|
||||
except (ValueError,TypeError): steps=5
|
||||
cancelled=False
|
||||
for i in range(steps):
|
||||
time.sleep(.01)
|
||||
fresh=db.execute("SELECT status FROM jobs WHERE id=?",(jid,)).fetchone()
|
||||
if not fresh or fresh["status"]=="cancelled": cancelled=True; break
|
||||
progress=int((i+1)*100/steps); db.execute("UPDATE jobs SET progress=?,updated_at=CURRENT_TIMESTAMP WHERE id=? AND status='running'",(progress,jid)); server_handler.add_job_event(db,jid,org,"progress",f"Job progress {progress}%",progress); db.commit()
|
||||
if cancelled: continue
|
||||
if payload.get("force_fail") or (payload.get("fail_once") and job["attempts"] == 0):
|
||||
db.execute("UPDATE jobs SET status='failed',error_code='DEMO_FAILURE',completed_at=CURRENT_TIMESTAMP,updated_at=CURRENT_TIMESTAMP WHERE id=?",(jid,));server_handler.add_job_event(db,jid,org,"failed","Job failed",job["progress"],"DEMO_FAILURE")
|
||||
else:
|
||||
db.execute("UPDATE jobs SET status='succeeded',progress=100,completed_at=CURRENT_TIMESTAMP,updated_at=CURRENT_TIMESTAMP WHERE id=?",(jid,));server_handler.add_job_event(db,jid,org,"succeeded","Job completed",100)
|
||||
db.commit()
|
||||
finally: db.close()
|
||||
|
||||
def create_server(host="127.0.0.1",port=8000,db_path="prospects.db"):
|
||||
server=ThreadingHTTPServer((host,port),ApiHandler);server.db_path=db_path;connect(db_path).close();return server
|
||||
server=ThreadingHTTPServer((host,port),ApiHandler);server.db_path=db_path;connect(db_path).close();server.job_stop=threading.Event();server.job_wakeup=threading.Event();server.job_thread=threading.Thread(target=_job_worker,args=(server,),daemon=True);server.job_thread.start()
|
||||
original_close=server.server_close
|
||||
def close():
|
||||
server.job_stop.set();server.job_wakeup.set();server.job_thread.join(timeout=2);original_close()
|
||||
server.server_close=close
|
||||
return server
|
||||
if __name__=="__main__":
|
||||
parser=argparse.ArgumentParser();parser.add_argument("--host",default="127.0.0.1");parser.add_argument("--port",type=int,default=int(os.environ.get("PROSPECT_API_PORT","8000")));parser.add_argument("--db",default=os.environ.get("PROSPECT_API_DB","prospects.db"));args=parser.parse_args();server=create_server(args.host,args.port,args.db);print(f"Prospect API listening on http://{args.host}:{args.port}",flush=True)
|
||||
try:server.serve_forever()
|
||||
|
||||
Reference in New Issue
Block a user