Persist job checkpoints and resource references in a synthetic lab

"""Synthetic jobs and resource-reference metadata. No upload, sending or deletion.
Company and worker are trusted test context, not authenticated request fields.
Resource bytes are hashed but NOT stored: this is not an object-storage service.
"""
import sqlite3,json,hashlib,re
from pathlib import Path
from contextlib import closing,contextmanager
def ident(v):return isinstance(v,str) and re.fullmatch('[A-Za-z0-9_-]{1,40}',v) is not None
class WorkLab:
 def __init__(self,path):
  self.path=str(path);exists=Path(path).exists() and Path(path).stat().st_size>0
  with closing(self.connect()) as db:
   if exists:
    if not db.execute("SELECT name FROM sqlite_master WHERE name='synthetic_work_marker'").fetchone() or db.execute('SELECT version FROM synthetic_work_marker').fetchall()[0][0]!='synthetic-work-v1':raise ValueError('Unrecognized synthetic store')
   else:
    db.executescript('''CREATE TABLE synthetic_work_marker(version TEXT);INSERT INTO synthetic_work_marker VALUES('synthetic-work-v1');
    CREATE TABLE jobs(company_id TEXT,id TEXT,total INTEGER,cursor INTEGER,generation INTEGER,worker TEXT,state TEXT,error_code TEXT,PRIMARY KEY(company_id,id));
    CREATE TABLE chunks(company_id TEXT,job_id TEXT,token TEXT,fingerprint TEXT,result_cursor INTEGER,PRIMARY KEY(company_id,job_id,token));
    CREATE TABLE resources(company_id TEXT,id TEXT,digest TEXT,byte_size INTEGER,reviewed INTEGER,PRIMARY KEY(company_id,id),UNIQUE(company_id,digest));
    CREATE TABLE refs(company_id TEXT,resource_id TEXT,document_id TEXT,kind TEXT,PRIMARY KEY(company_id,resource_id,document_id,kind));''')
 def connect(self):
  db=sqlite3.connect(self.path,timeout=5,isolation_level=None);db.row_factory=sqlite3.Row;return db
 @contextmanager
 def transaction(self):
  with closing(self.connect()) as db:
   db.execute('BEGIN IMMEDIATE')
   try:yield db;db.commit()
   except BaseException:db.rollback();raise
 def create_job(self,company,job,total):
  if not ident(company) or not ident(job) or type(total) is not int or not 0<=total<=1000:raise ValueError('Invalid job')
  with self.transaction() as db:db.execute('INSERT INTO jobs VALUES(?,?,?,0,0,NULL,?,NULL)',[company,job,total,'done' if total==0 else 'queued'])
 def read_job(self,company,job):
  if not ident(company) or not ident(job):return None
  with closing(self.connect()) as db:
   row=db.execute('SELECT * FROM jobs WHERE company_id=? AND id=?',[company,job]).fetchone();return dict(row) if row else None
 def claim(self,company,job,worker,expected_generation):
  if not all(ident(v) for v in [company,job,worker]) or type(expected_generation) is not int:return None
  with self.transaction() as db:
   row=db.execute('SELECT * FROM jobs WHERE company_id=? AND id=?',[company,job]).fetchone()
   if not row or row['state'] not in ['queued','failed'] or row['generation']!=expected_generation:return None
   generation=row['generation']+1;db.execute("UPDATE jobs SET generation=?,worker=?,state='running',error_code=NULL WHERE company_id=? AND id=?",[generation,worker,company,job]);return generation
 def advance(self,company,job,worker,generation,expected_cursor,token,count):
  if not all(ident(v) for v in [company,job,worker,token]) or any(type(v) is not int for v in [generation,expected_cursor,count]) or count<1:return {'ok':False,'reason':'input'}
  fingerprint=hashlib.sha256(json.dumps([expected_cursor,count]).encode()).hexdigest()
  with self.transaction() as db:
   row=db.execute('SELECT * FROM jobs WHERE company_id=? AND id=?',[company,job]).fetchone()
   if not row or row['worker']!=worker or row['generation']!=generation:return {'ok':False,'reason':'stale_worker'}
   old=db.execute('SELECT * FROM chunks WHERE company_id=? AND job_id=? AND token=?',[company,job,token]).fetchone()
   if old:return {'ok':True,'cursor':old['result_cursor'],'replayed':True} if old['fingerprint']==fingerprint else {'ok':False,'reason':'token_reused'}
   if row['state']!='running' or row['cursor']!=expected_cursor or row['cursor']+count>row['total']:return {'ok':False,'reason':'cursor_conflict'}
   next_cursor=row['cursor']+count;db.execute('INSERT INTO chunks VALUES(?,?,?,?,?)',[company,job,token,fingerprint,next_cursor]);db.execute('UPDATE jobs SET cursor=?,state=? WHERE company_id=? AND id=?',[next_cursor,'done' if next_cursor==row['total'] else 'running',company,job]);return {'ok':True,'cursor':next_cursor,'replayed':False}
 def fail(self,company,job,worker,generation):
  with self.transaction() as db:
   return db.execute("UPDATE jobs SET state='failed',error_code='synthetic_failure' WHERE company_id=? AND id=? AND worker=? AND generation=? AND state='running'",[company,job,worker,generation]).rowcount==1
 def register_resource(self,company,resource_id,data):
  if not ident(company) or not ident(resource_id) or not isinstance(data,bytes) or not 0<len(data)<=100000:raise ValueError('Invalid resource metadata input')
  digest=hashlib.sha256(data).hexdigest()
  with self.transaction() as db:
   old=db.execute('SELECT id FROM resources WHERE company_id=? AND digest=?',[company,digest]).fetchone()
   if old:return {'id':old['id'],'reused':True}
   db.execute('INSERT INTO resources VALUES(?,?,?,?,0)',[company,resource_id,digest,len(data)]);return {'id':resource_id,'reused':False}
 def review_resource(self,company,resource_id,reviewed):
  if type(reviewed) is not bool:raise ValueError('Review flag required')
  with self.transaction() as db:
   if not reviewed and db.execute("SELECT 1 FROM refs WHERE company_id=? AND resource_id=? AND kind='current'",[company,resource_id]).fetchone():return False
   return db.execute('UPDATE resources SET reviewed=? WHERE company_id=? AND id=?',[int(reviewed),company,resource_id]).rowcount==1
 def link(self,company,resource_id,document_id,kind):
  if not all(ident(v) for v in [company,resource_id,document_id]) or kind not in ['current','history','backup']:return False
  with self.transaction() as db:
   row=db.execute('SELECT reviewed FROM resources WHERE company_id=? AND id=?',[company,resource_id]).fetchone()
   if not row or not row['reviewed']:return False
   db.execute('INSERT OR IGNORE INTO refs VALUES(?,?,?,?)',[company,resource_id,document_id,kind]);return True
 def unlink(self,company,resource_id,document_id,kind):
  with self.transaction() as db:return db.execute('DELETE FROM refs WHERE company_id=? AND resource_id=? AND document_id=? AND kind=?',[company,resource_id,document_id,kind]).rowcount==1
 def resource_preview(self,company):
  with closing(self.connect()) as db:
   rows=db.execute('SELECT r.*,count(f.resource_id) AS ref_count FROM resources r LEFT JOIN refs f ON f.company_id=r.company_id AND f.resource_id=r.id WHERE r.company_id=? GROUP BY r.company_id,r.id ORDER BY r.id',[company]).fetchall()
   return [{**dict(row),'disposition':'retain' if row['ref_count'] else 'manual_review_only','deleted':False} for row in rows]