A R R R Associates
426 lines
16 KiB
Python
426 lines
16 KiB
Python
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
|
|
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._-]+")
|
|
_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_value = _SAFE.sub("_", str(value or default).strip()).strip("._-")
|
|
return text_value[:80] or default
|
|
|
|
|
|
def role_bucket(roles: Iterable[str]) -> str | None:
|
|
role_set = set(roles)
|
|
for role, bucket in (("Partner", "Partner"), ("Manager", "Manager"), ("Branch Manager", "Manager"), ("Staff", "Staff"), ("Employee", "Staff"), ("Consultant", "Consultant")):
|
|
if role in role_set:
|
|
return bucket
|
|
return None
|
|
|
|
|
|
def can_use(roles: Iterable[str]) -> bool:
|
|
return bool(set(roles) & ALLOWED_ROLES)
|
|
|
|
|
|
def _user_root(user, roles: Iterable[str]) -> Path:
|
|
bucket = role_bucket(roles)
|
|
if not bucket:
|
|
raise PermissionError("Bank Statement Analyzer is available only to Partner, Manager, Staff and Consultant roles.")
|
|
label = _segment(getattr(user, "full_name", None) or getattr(user, "email", None) or f"user_{user.id}")
|
|
return _root() / bucket / f"{int(user.id)}_{label}" / "Bank_Statement_Analyzer"
|
|
|
|
|
|
def create_job_folder(user, roles: Iterable[str]) -> tuple[str, Path, Path]:
|
|
job_id = uuid.uuid4().hex
|
|
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
|
|
|
|
|
|
def _validate_pdf_header(data: bytes) -> None:
|
|
if not data.startswith(b"%PDF-"):
|
|
raise ValueError("Only genuine PDF files are allowed.")
|
|
|
|
|
|
async def save_uploads(files: list[UploadFile], input_dir: Path) -> list[Path]:
|
|
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] = []
|
|
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 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,
|
|
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(),
|
|
)
|
|
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 _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 _workbook_filename(metas) -> str:
|
|
"""Return BankName_ClientName.xlsx using filesystem-safe segments."""
|
|
banks = sorted({_segment(getattr(meta, "bank_name", ""), "Bank") for meta in metas if str(getattr(meta, "bank_name", "") or "").strip()})
|
|
customers = sorted({_segment(getattr(meta, "customer_name", ""), "Client") for meta in metas if str(getattr(meta, "customer_name", "") or "").strip()})
|
|
bank_label = banks[0] if len(banks) == 1 else ("Multi_Bank" if banks else "Bank")
|
|
client_label = customers[0] if len(customers) == 1 else ("Multiple_Clients" if customers else "Client")
|
|
return f"{bank_label}_{client_label}.xlsx"
|
|
|
|
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 / _workbook_filename(metas)
|
|
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"].fillna(False).astype(bool).sum()) if not all_df.empty and "exact_duplicate" in all_df.columns else 0,
|
|
"possible_duplicate_rows": int(all_df["possible_duplicate"].fillna(False).astype(bool).sum()) if not all_df.empty and "possible_duplicate" in all_df.columns 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 ""):
|
|
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 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()
|