Add queued bank statement analysis with global concurrency limits

This commit is contained in:
A R R R Associates
2026-07-13 15:02:58 +05:30
parent afb098aa75
commit 897c7dad38
6 changed files with 517 additions and 220 deletions
+338 -127
View File
@@ -1,58 +1,63 @@
from __future__ import annotations
import json
import math
import os
import re
import shutil
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Iterable
from fastapi import UploadFile
from sqlalchemy import func, select, text
from app.core.db.common import CommonSessionLocal
from .analyzer import analyze_files, export_excel
from .models import BankStatementAnalysisJob
ALLOWED_ROLES = {"Partner", "Manager", "Branch Manager", "Staff", "Employee", "Consultant"}
MAX_FILES = int(os.getenv("BANK_ANALYZER_MAX_FILES", "24"))
MAX_FILE_BYTES = int(os.getenv("BANK_ANALYZER_MAX_FILE_MB", "50")) * 1024 * 1024
RETENTION_HOURS = int(os.getenv("BANK_ANALYZER_FAILED_RETENTION_HOURS", "24"))
MAX_GLOBAL_PROCESSING = 3
MAX_PENDING_PER_USER = 3
COMPLETED_RETENTION_HOURS = 24
FAILED_RETENTION_HOURS = 24
POLL_SECONDS = 2
_SAFE = re.compile(r"[^A-Za-z0-9._-]+")
def _resolve_document_storage_root() -> Path:
"""Resolve the same automatic document root without importing documents.services.
Importing app.modules.documents.services here creates a circular import through
clients.__init__ -> clients.ui -> documents.services. This local resolver
intentionally mirrors the existing document module's storage-root rules while
keeping the bank analyzer independent of document models and UI modules.
"""
configured = (os.getenv("DOCUMENT_STORAGE_ROOT") or "").strip()
if configured:
configured_path = Path(configured).expanduser()
if configured_path.is_absolute():
return configured_path
return (Path.cwd() / configured_path).resolve()
if os.name == "nt":
anchor = Path.cwd().anchor or (os.environ.get("SystemDrive", "C:") + "\\")
return Path(anchor) / "AuditFirmERPDocuments" / "engagement_documents"
return (Path.cwd() / "documents" / "engagement_documents").resolve()
_QUEUE_LOCK_ID = 820260713
_executor = ThreadPoolExecutor(max_workers=MAX_GLOBAL_PROCESSING, thread_name_prefix="bank-analyzer")
_worker_lock = threading.Lock()
_worker_started = False
def _now() -> datetime:
return datetime.now(timezone.utc)
def _resolve_document_storage_root() -> Path:
configured = (os.getenv("DOCUMENT_STORAGE_ROOT") or "").strip()
if configured:
configured_path = Path(configured).expanduser()
return configured_path if configured_path.is_absolute() else (Path.cwd() / configured_path).resolve()
if os.name == "nt":
anchor = Path.cwd().anchor or (os.environ.get("SystemDrive", "C:") + "\\")
return Path(anchor) / "AuditFirmERPDocuments" / "engagement_documents"
return (Path.cwd() / "documents" / "engagement_documents").resolve()
def _root() -> Path:
return _resolve_document_storage_root().parent / "Work"
def _segment(value: object, default: str = "NA") -> str:
text = _SAFE.sub("_", str(value or default).strip()).strip("._-")
return text[:80] or default
text_value = _SAFE.sub("_", str(value or default).strip()).strip("._-")
return text_value[:80] or default
def role_bucket(roles: Iterable[str]) -> str | None:
@@ -75,26 +80,11 @@ def _user_root(user, roles: Iterable[str]) -> Path:
return _root() / bucket / f"{int(user.id)}_{label}" / "Bank_Statement_Analyzer"
def cleanup_expired(user, roles: Iterable[str]) -> None:
base = _user_root(user, roles)
if not base.exists():
return
cutoff = _now() - timedelta(hours=RETENTION_HOURS)
for child in base.iterdir():
try:
modified = datetime.fromtimestamp(child.stat().st_mtime, tz=timezone.utc)
if child.is_dir() and modified < cutoff:
shutil.rmtree(child, ignore_errors=True)
except OSError:
continue
def create_job(user, roles: Iterable[str]) -> tuple[str, Path, Path]:
cleanup_expired(user, roles)
def create_job_folder(user, roles: Iterable[str]) -> tuple[str, Path, Path]:
job_id = uuid.uuid4().hex
job = _user_root(user, roles) / job_id
input_dir = job / "Input"
output_dir = job / "Output"
job_dir = _user_root(user, roles) / job_id
input_dir = job_dir / "Input"
output_dir = job_dir / "Output"
input_dir.mkdir(parents=True, exist_ok=False)
output_dir.mkdir(parents=True, exist_ok=False)
return job_id, input_dir, output_dir
@@ -106,100 +96,321 @@ def _validate_pdf_header(data: bytes) -> None:
async def save_uploads(files: list[UploadFile], input_dir: Path) -> list[Path]:
usable = [file for file in files if file and (file.filename or "").strip()]
usable = [item for item in files if item and (item.filename or "").strip()]
if not usable:
raise ValueError("Please select at least one PDF bank statement.")
if len(usable) > MAX_FILES:
raise ValueError(f"A maximum of {MAX_FILES} PDF files can be analyzed in one job.")
saved: list[Path] = []
for index, upload in enumerate(usable, start=1):
name = Path(upload.filename or f"statement_{index}.pdf").name
if Path(name).suffix.lower() != ".pdf":
raise ValueError(f"{name}: only PDF files are allowed.")
safe_name = f"{index:02d}_{_segment(Path(name).stem, f'statement_{index}')}.pdf"
target = input_dir / safe_name
size = 0
first = b""
with target.open("wb") as handle:
while True:
chunk = await upload.read(1024 * 1024)
if not chunk:
break
if not first:
first = chunk[:8]
size += len(chunk)
if size > MAX_FILE_BYTES:
raise ValueError(f"{name}: file exceeds the {MAX_FILE_BYTES // (1024 * 1024)} MB limit.")
handle.write(chunk)
_validate_pdf_header(first)
saved.append(target)
return saved
try:
for index, upload in enumerate(usable, start=1):
name = Path(upload.filename or f"statement_{index}.pdf").name
if Path(name).suffix.lower() != ".pdf":
raise ValueError(f"{name}: only PDF files are allowed.")
target = input_dir / f"{index:02d}_{_segment(Path(name).stem, f'statement_{index}')}.pdf"
size = 0
first = b""
with target.open("wb") as handle:
while True:
chunk = await upload.read(1024 * 1024)
if not chunk:
break
if not first:
first = chunk[:8]
size += len(chunk)
if size > MAX_FILE_BYTES:
raise ValueError(f"{name}: file exceeds the {MAX_FILE_BYTES // (1024 * 1024)} MB limit.")
handle.write(chunk)
_validate_pdf_header(first)
saved.append(target)
return saved
except Exception:
shutil.rmtree(input_dir.parent, ignore_errors=True)
raise
def analyze_job(
*,
user,
roles: Iterable[str],
job_id: str,
paths: list[Path],
output_dir: Path,
customer_override: str = "",
account_override: str = "",
bank_selection: str = "auto",
financial_year: str = "",
classification_enabled: bool = True,
) -> dict:
metas, all_df, unique_df = analyze_files(
paths,
customer_override.strip(),
account_override.strip(),
bank_hint=bank_selection,
classification_enabled=classification_enabled,
)
output = output_dir / "Bank_Statement_Analysis.xlsx"
export_excel(
output,
metas,
all_df,
unique_df,
financial_year=financial_year.strip(),
def pending_count_for_user(user_id: int) -> int:
db = CommonSessionLocal()
try:
return int(db.scalar(select(func.count()).select_from(BankStatementAnalysisJob).where(BankStatementAnalysisJob.user_id == int(user_id), BankStatementAnalysisJob.status.in_(["queued", "processing"]))) or 0)
finally:
db.close()
def enqueue_job(*, user, roles: Iterable[str], job_id: str, paths: list[Path], job_dir: Path, bank_selection: str, financial_year: str, customer_override: str, account_override: str, classification_enabled: bool) -> BankStatementAnalysisJob:
if pending_count_for_user(int(user.id)) >= MAX_PENDING_PER_USER:
shutil.rmtree(job_dir, ignore_errors=True)
raise ValueError("You already have three queued or processing analyses. Please wait for one to complete before submitting another.")
bucket = role_bucket(roles)
if not bucket:
shutil.rmtree(job_dir, ignore_errors=True)
raise PermissionError("You do not have access to the Bank Statement Analyzer.")
job = BankStatementAnalysisJob(
id=job_id,
tenant_id=getattr(user, "tenant_id", None),
branch_id=getattr(user, "branch_id", None),
user_id=int(user.id),
role_bucket=bucket,
selected_bank=bank_selection,
classification_enabled=classification_enabled,
financial_year=(financial_year or "").strip() or None,
customer_override=(customer_override or "").strip() or None,
account_override=(account_override or "").strip() or None,
classification_enabled=bool(classification_enabled),
status="queued",
progress_percent=0,
file_count=len(paths),
input_files_json=json.dumps([str(path) for path in paths]),
job_directory=str(job_dir),
submitted_at_utc=_now(),
)
summary = {
"job_id": job_id,
"owner_user_id": int(user.id),
"created_at": _now().isoformat(),
"statement_count": len(metas),
"rows_extracted": int(len(all_df)),
"unique_transactions": int(len(unique_df)),
"exact_duplicate_rows": int(all_df.exact_duplicate.sum()) if not all_df.empty else 0,
"possible_duplicate_rows": int(all_df.possible_duplicate.sum()) if not all_df.empty else 0,
"review_items": int(unique_df.review_note.fillna("").ne("").sum()) if not unique_df.empty else 0,
"categories": int(unique_df.category.nunique()) if not unique_df.empty else 0,
"banks": sorted({meta.bank_name for meta in metas}),
"customer_name": next((meta.customer_name for meta in metas if meta.customer_name), ""),
"account_number": next((meta.account_number for meta in metas if meta.account_number), ""),
"financial_year": financial_year.strip(),
"classification_enabled": classification_enabled,
"output_file": output.name,
}
(output_dir.parent / "job.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
return summary
db = CommonSessionLocal()
try:
db.add(job)
db.commit()
db.refresh(job)
except Exception:
db.rollback()
shutil.rmtree(job_dir, ignore_errors=True)
raise
finally:
db.close()
ensure_worker_started()
return job
def resolve_owned_job(user, roles: Iterable[str], job_id: str) -> tuple[Path, dict]:
def _claim_jobs() -> list[str]:
db = CommonSessionLocal()
try:
try:
db.execute(text("SELECT pg_advisory_xact_lock(:lock_id)"), {"lock_id": _QUEUE_LOCK_ID})
except Exception:
db.rollback()
processing = int(db.scalar(select(func.count()).select_from(BankStatementAnalysisJob).where(BankStatementAnalysisJob.status == "processing")) or 0)
capacity = max(0, MAX_GLOBAL_PROCESSING - processing)
if capacity <= 0:
return []
stmt = select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.status == "queued").order_by(BankStatementAnalysisJob.submitted_at_utc.asc()).limit(capacity)
try:
stmt = stmt.with_for_update(skip_locked=True)
except Exception:
pass
jobs = list(db.scalars(stmt).all())
now = _now()
for job in jobs:
job.status = "processing"
job.progress_percent = 5
job.started_at_utc = now
job.error_message = None
db.commit()
return [job.id for job in jobs]
except Exception:
db.rollback()
return []
finally:
db.close()
def _process_job(job_id: str) -> None:
db = CommonSessionLocal()
try:
job = db.get(BankStatementAnalysisJob, job_id)
if not job or job.status != "processing":
return
paths = [Path(value) for value in json.loads(job.input_files_json)]
output_dir = Path(job.job_directory) / "Output"
job.progress_percent = 15
db.commit()
metas, all_df, unique_df = analyze_files(
paths,
job.customer_override or "",
job.account_override or "",
bank_hint=job.selected_bank,
classification_enabled=job.classification_enabled,
)
job.progress_percent = 75
db.commit()
output = output_dir / "Bank_Statement_Analysis.xlsx"
export_excel(
output,
metas,
all_df,
unique_df,
financial_year=job.financial_year or "",
selected_bank=job.selected_bank,
classification_enabled=job.classification_enabled,
)
summary = {
"job_id": job.id,
"statement_count": len(metas),
"rows_extracted": int(len(all_df)),
"unique_transactions": int(len(unique_df)),
"exact_duplicate_rows": int(all_df.exact_duplicate.sum()) if not all_df.empty else 0,
"possible_duplicate_rows": int(all_df.possible_duplicate.sum()) if not all_df.empty else 0,
"review_items": int(unique_df.review_note.fillna("").ne("").sum()) if not unique_df.empty and "review_note" in unique_df.columns else 0,
"categories": int(unique_df.category.nunique()) if not unique_df.empty and "category" in unique_df.columns else 0,
"banks": sorted({meta.bank_name for meta in metas}),
"customer_name": next((meta.customer_name for meta in metas if meta.customer_name), ""),
"account_number": next((meta.account_number for meta in metas if meta.account_number), ""),
"financial_year": job.financial_year or "",
"classification_enabled": job.classification_enabled,
}
for path in paths:
try:
path.unlink(missing_ok=True)
except OSError:
pass
input_dir = Path(job.job_directory) / "Input"
try:
input_dir.rmdir()
except OSError:
pass
job.output_file = str(output)
job.summary_json = json.dumps(summary)
job.status = "completed"
job.progress_percent = 100
job.completed_at_utc = _now()
job.expires_at_utc = _now() + timedelta(hours=COMPLETED_RETENTION_HOURS)
db.commit()
except Exception as exc:
db.rollback()
job = db.get(BankStatementAnalysisJob, job_id)
if job:
job.status = "failed"
job.progress_percent = 100
job.error_message = str(exc)[:4000]
job.completed_at_utc = _now()
job.expires_at_utc = _now() + timedelta(hours=FAILED_RETENTION_HOURS)
db.commit()
finally:
db.close()
def _cleanup_expired() -> None:
db = CommonSessionLocal()
try:
jobs = list(db.scalars(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.expires_at_utc.is_not(None), BankStatementAnalysisJob.expires_at_utc < _now())).all())
for job in jobs:
shutil.rmtree(job.job_directory, ignore_errors=True)
db.delete(job)
db.commit()
except Exception:
db.rollback()
finally:
db.close()
def _worker_loop() -> None:
while True:
try:
_cleanup_expired()
for job_id in _claim_jobs():
_executor.submit(_process_job, job_id)
except Exception:
pass
time.sleep(POLL_SECONDS)
def ensure_worker_started() -> None:
global _worker_started
with _worker_lock:
if _worker_started:
return
thread = threading.Thread(target=_worker_loop, name="bank-analyzer-queue", daemon=True)
thread.start()
_worker_started = True
def _summary(job: BankStatementAnalysisJob) -> dict:
return json.loads(job.summary_json) if job.summary_json else {}
def average_duration_minutes() -> float:
db = CommonSessionLocal()
try:
jobs = list(db.scalars(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.status == "completed", BankStatementAnalysisJob.started_at_utc.is_not(None), BankStatementAnalysisJob.completed_at_utc.is_not(None)).order_by(BankStatementAnalysisJob.completed_at_utc.desc()).limit(20)).all())
durations = [(job.completed_at_utc - job.started_at_utc).total_seconds() / 60 for job in jobs if job.completed_at_utc and job.started_at_utc]
return max(1.0, sum(durations) / len(durations)) if durations else 4.0
finally:
db.close()
def queue_position(job: BankStatementAnalysisJob) -> int | None:
if job.status != "queued":
return None
db = CommonSessionLocal()
try:
return int(db.scalar(select(func.count()).select_from(BankStatementAnalysisJob).where(BankStatementAnalysisJob.status == "queued", BankStatementAnalysisJob.submitted_at_utc <= job.submitted_at_utc)) or 1)
finally:
db.close()
def estimated_wait(job: BankStatementAnalysisJob) -> str:
if job.status == "processing":
return "Processing now"
if job.status != "queued":
return ""
position = queue_position(job) or 1
waves = max(1, math.ceil(position / MAX_GLOBAL_PROCESSING))
average = average_duration_minutes()
low = max(1, math.ceil((waves - 1) * average))
high = max(5, math.ceil(waves * average + 2))
return f"Approximately {low}-{high} minutes"
def get_owned_job(user_id: int, job_id: str) -> BankStatementAnalysisJob | None:
if not re.fullmatch(r"[a-f0-9]{32}", job_id or ""):
raise FileNotFoundError("Analysis job not found.")
job = _user_root(user, roles) / job_id
meta_path = job / "job.json"
if not meta_path.is_file():
raise FileNotFoundError("Analysis job not found.")
meta = json.loads(meta_path.read_text(encoding="utf-8"))
if int(meta.get("owner_user_id", 0)) != int(user.id):
raise PermissionError("You cannot access another user's analysis job.")
return job, meta
return None
db = CommonSessionLocal()
try:
job = db.scalar(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.id == job_id, BankStatementAnalysisJob.user_id == int(user_id)))
if job:
db.expunge(job)
return job
finally:
db.close()
def delete_job(job: Path) -> None:
shutil.rmtree(job, ignore_errors=True)
def list_user_jobs(user_id: int, limit: int = 25) -> list[BankStatementAnalysisJob]:
db = CommonSessionLocal()
try:
jobs = list(db.scalars(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.user_id == int(user_id)).order_by(BankStatementAnalysisJob.submitted_at_utc.desc()).limit(limit)).all())
for job in jobs:
db.expunge(job)
return jobs
finally:
db.close()
def job_view(job: BankStatementAnalysisJob) -> dict:
return {
"id": job.id,
"status": job.status,
"progress_percent": job.progress_percent,
"file_count": job.file_count,
"selected_bank": job.selected_bank,
"financial_year": job.financial_year or "",
"submitted_at": job.submitted_at_utc,
"completed_at": job.completed_at_utc,
"expires_at": job.expires_at_utc,
"error_message": job.error_message or "",
"queue_position": queue_position(job),
"estimated_wait": estimated_wait(job),
"summary": _summary(job),
"download_ready": job.status == "completed" and bool(job.output_file) and Path(job.output_file).is_file(),
}
def delete_owned_job(user_id: int, job_id: str) -> bool:
db = CommonSessionLocal()
try:
job = db.scalar(select(BankStatementAnalysisJob).where(BankStatementAnalysisJob.id == job_id, BankStatementAnalysisJob.user_id == int(user_id)))
if not job:
return False
if job.status == "processing":
raise ValueError("A processing analysis cannot be deleted. Please wait for it to finish.")
shutil.rmtree(job.job_directory, ignore_errors=True)
db.delete(job)
db.commit()
return True
finally:
db.close()