From 5eb8fa586b2e30f603b03137d5ff88f1ea9706f9 Mon Sep 17 00:00:00 2001 From: Christian Krakau-Louis Date: Tue, 25 Mar 2025 13:55:15 +0100 Subject: [PATCH 1/3] added duplicate check and database --- app/database.py | 9 ++++- app/main.py | 5 +++ app/models.py | 38 ++++++++++++++++++- app/tasks/upload_to_s3.py | 77 +++++++++++++++++++++++++++++++-------- app/utils.py | 16 ++++++++ 5 files changed, 125 insertions(+), 20 deletions(-) create mode 100644 app/utils.py diff --git a/app/database.py b/app/database.py index a513c72d..fe71b5e7 100644 --- a/app/database.py +++ b/app/database.py @@ -1,14 +1,19 @@ +# app/database.py #!/usr/bin/env python3 -from sqlalchemy import create_engine, Column, String, Integer +from sqlalchemy import create_engine from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker -from .config import settings +from app.config import settings Base = declarative_base() engine = create_engine(settings.database_url, connect_args={"check_same_thread": False}) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) +def init_db(): + """Call this once (e.g. on startup) to create tables if they don't exist.""" + Base.metadata.create_all(bind=engine) + def get_db(): db = SessionLocal() try: diff --git a/app/main.py b/app/main.py index 9d03538f..aeaa70cc 100644 --- a/app/main.py +++ b/app/main.py @@ -2,6 +2,7 @@ import os from fastapi import FastAPI, HTTPException, UploadFile, File +from app.database import init_db from app.config import settings from app.tasks.upload_to_s3 import upload_to_s3 from app.tasks.upload_to_dropbox import upload_to_dropbox @@ -12,6 +13,10 @@ from app.frontend import router as frontend_router app = FastAPI(title="Document Processing API") +@app.on_event("startup") +def on_startup(): + init_db() # Create tables if they don't exist + @app.get("/") def root(): return {"message": "Document Processing API"} diff --git a/app/models.py b/app/models.py index 0ad894e4..6b129b7d 100644 --- a/app/models.py +++ b/app/models.py @@ -1,7 +1,9 @@ +# app/models.py #!/usr/bin/env python3 -from .database import Base -from sqlalchemy import Column, String, Integer +from sqlalchemy import Column, String, Integer, DateTime, func +from sqlalchemy.ext.declarative import declarative_base +from app.database import Base class DocumentMetadata(Base): __tablename__ = "documents" @@ -12,3 +14,35 @@ class DocumentMetadata(Base): recipient = Column(String) tags = Column(String) summary = Column(String) + +class FileRecord(Base): + __tablename__ = "files" + + id = Column(Integer, primary_key=True, index=True) + + # Hash of the file content (e.g. SHA-256) + filehash = Column(String, unique=True, index=True, nullable=False) + + # The name of the file as it was originally uploaded (if known) + original_filename = Column(String) + + # The name/path we store on disk (e.g. /workdir/tmp/.pdf) + local_filename = Column(String, nullable=False) + + # Size of the file in bytes + file_size = Column(Integer, nullable=False) + + # MIME type or extension (optional) + mime_type = Column(String) + + # Timestamp when we inserted this record + created_at = Column(DateTime(timezone=True), server_default=func.now()) + +class ProcessingLog(Base): + __tablename__ = "processing_logs" + id = Column(Integer, primary_key=True, index=True) + file_id = Column(Integer, ForeignKey("files.id")) + step_name = Column(String) # e.g. "OCR", "convert_to_pdf", "upload_s3" + status = Column(String) # "success" / "failure" + message = Column(String) # error text or success note + timestamp = Column(DateTime(timezone=True), server_default=func.now()) diff --git a/app/tasks/upload_to_s3.py b/app/tasks/upload_to_s3.py index 10919bed..45e09ecb 100644 --- a/app/tasks/upload_to_s3.py +++ b/app/tasks/upload_to_s3.py @@ -4,15 +4,20 @@ import os import uuid import boto3 import shutil +import mimetypes import fitz # PyMuPDF for checking embedded text + from app.config import settings from app.tasks.retry_config import BaseTaskWithRetry from app.tasks.process_with_textract import process_with_textract from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt - -# Import the shared Celery instance from app.celery_app import celery +# NEW imports for the DB +from app.database import SessionLocal +from app.models import FileRecord +from app.utils import hash_file + # Initialize S3 client s3_client = boto3.client( "s3", @@ -21,13 +26,20 @@ s3_client = boto3.client( region_name=settings.aws_region, ) + @celery.task(base=BaseTaskWithRetry) def upload_to_s3(original_local_file: str): """ Uploads a file to S3 with a UUID-based filename and triggers processing. - - If the PDF already contains embedded text, skip Textract and extract text locally. - - Otherwise, upload to S3 and process with Textract. + + Steps: + 1. Check if we have a FileRecord entry (via SHA-256 hash). If found, skip re-processing. + 2. If not found, insert a new DB row and continue with the pipeline: + - Copy file to /workdir/tmp + - Check for embedded text. If present, skip S3 and run local GPT extraction + - Otherwise, upload to S3 and queue Textract-based OCR """ + bucket_name = settings.s3_bucket_name if not bucket_name: print("[ERROR] S3 bucket name not set.") @@ -37,22 +49,54 @@ def upload_to_s3(original_local_file: str): print(f"[ERROR] File {original_local_file} not found.") return {"error": "File not found"} - # Generate UUID and create a new filename - file_ext = os.path.splitext(original_local_file)[1] # Preserve original file extension - file_uuid = str(uuid.uuid4()) - new_filename = f"{file_uuid}{file_ext}" + # 0. Compute the file hash and check for duplicates + filehash = hash_file(original_local_file) + original_filename = os.path.basename(original_local_file) + file_size = os.path.getsize(original_local_file) + mime_type, _ = mimetypes.guess_type(original_local_file) + if not mime_type: + mime_type = "application/octet-stream" - # Construct the new local path using settings.workdir and a 'tmp' subdirectory - tmp_dir = os.path.join(settings.workdir, "tmp") - new_local_path = os.path.join(tmp_dir, new_filename) + # Acquire DB session in the task + with SessionLocal() as db: + existing = db.query(FileRecord).filter_by(filehash=filehash).one_or_none() + if existing: + print(f"[INFO] Duplicate file detected (hash={filehash[:10]}...) Skipping processing.") + return { + "status": "duplicate_file", + "file_id": existing.id, + "detail": "File already processed." + } - # Ensure the target tmp directory exists - os.makedirs(tmp_dir, exist_ok=True) + # Not a duplicate -> insert a new record + new_record = FileRecord( + filehash=filehash, + original_filename=original_filename, + local_filename="", # Will fill in after we move it + file_size=file_size, + mime_type=mime_type, + ) + db.add(new_record) + db.commit() + db.refresh(new_record) - # Copy the file instead of moving it - shutil.copy(original_local_file, new_local_path) + # 1. Generate a UUID-based filename and place it in /workdir/tmp + file_ext = os.path.splitext(original_local_file)[1] + file_uuid = str(uuid.uuid4()) + new_filename = f"{file_uuid}{file_ext}" - # Check for embedded text + tmp_dir = os.path.join(settings.workdir, "tmp") + os.makedirs(tmp_dir, exist_ok=True) + new_local_path = os.path.join(tmp_dir, new_filename) + + # Copy the file instead of moving it + shutil.copy(original_local_file, new_local_path) + + # Update the DB with final local filename + new_record.local_filename = new_local_path + db.commit() + + # 2. Check for embedded text (outside the DB session to avoid long open transactions) pdf_doc = fitz.open(new_local_path) has_text = any(page.get_text() for page in pdf_doc) pdf_doc.close() @@ -72,6 +116,7 @@ def upload_to_s3(original_local_file: str): return {"file": new_local_path, "status": "Text extracted locally"} + # 3. If no embedded text, upload to S3 and queue Textract processing try: print(f"[INFO] Uploading {new_local_path} to s3://{bucket_name}/{new_filename}...") s3_client.upload_file(new_local_path, bucket_name, new_filename) diff --git a/app/utils.py b/app/utils.py new file mode 100644 index 00000000..05d0026e --- /dev/null +++ b/app/utils.py @@ -0,0 +1,16 @@ +# app/utils.py +import hashlib + +def hash_file(filepath, chunk_size=65536): + """ + Returns the SHA-256 hash of the file at 'filepath'. + Reads the file in chunks to handle large files efficiently. + """ + sha256 = hashlib.sha256() + with open(filepath, "rb") as f: + while True: + data = f.read(chunk_size) + if not data: + break + sha256.update(data) + return sha256.hexdigest() From 260358a78b0a2224bee7a102ed22b871703a1dba Mon Sep 17 00:00:00 2001 From: Christian Krakau-Louis Date: Tue, 25 Mar 2025 13:58:54 +0100 Subject: [PATCH 2/3] added ForeignKey object --- app/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/models.py b/app/models.py index 6b129b7d..ea07a074 100644 --- a/app/models.py +++ b/app/models.py @@ -1,7 +1,7 @@ # app/models.py #!/usr/bin/env python3 -from sqlalchemy import Column, String, Integer, DateTime, func +from sqlalchemy import Column, String, Integer, DateTime, func, ForeignKey from sqlalchemy.ext.declarative import declarative_base from app.database import Base From ae674e8fc4ae042d94e038a0f78b445d8dd0481a Mon Sep 17 00:00:00 2001 From: Christian Krakau-Louis Date: Tue, 25 Mar 2025 14:33:46 +0100 Subject: [PATCH 3/3] added robust code to create the sqlite file --- app/database.py | 53 ++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/app/database.py b/app/database.py index fe71b5e7..8fff3c3d 100644 --- a/app/database.py +++ b/app/database.py @@ -1,20 +1,63 @@ # app/database.py -#!/usr/bin/env python3 -from sqlalchemy import create_engine +import os +import logging + +from sqlalchemy import create_engine, exc from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker +from sqlalchemy.engine.url import make_url + from app.config import settings +logger = logging.getLogger(__name__) + Base = declarative_base() -engine = create_engine(settings.database_url, connect_args={"check_same_thread": False}) + +# Parse the DATABASE_URL +DB_URL = settings.database_url +engine = create_engine(DB_URL, connect_args={"check_same_thread": False}) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) + def init_db(): - """Call this once (e.g. on startup) to create tables if they don't exist.""" - Base.metadata.create_all(bind=engine) + """ + Ensures the SQLite database file and its parent directory exist (if using sqlite). + Then runs Base.metadata.create_all(bind=engine) to initialize tables. + Logs a message if a new SQLite DB file is created. + """ + # 1. Parse the DB URL to see if it's sqlite + url = make_url(DB_URL) + if url.get_backend_name() == "sqlite": + # 2. Extract the database path from the URL + database_path = url.database # e.g. "/workdir/db/database.db" or ":memory:" + + if database_path != ":memory:": + # 3. Ensure directory exists + db_dir = os.path.dirname(database_path) + if db_dir and not os.path.exists(db_dir): + logger.info(f"Creating directory for SQLite DB: {db_dir}") + os.makedirs(db_dir, exist_ok=True) + + # 4. If the file does not exist, create an empty one + if not os.path.exists(database_path): + logger.info(f"Creating new SQLite database file at {database_path}") + open(database_path, "a").close() + + # 5. Now create tables if they don't exist yet + try: + Base.metadata.create_all(bind=engine) + logger.info("Database initialization complete (tables created if not exist).") + except exc.SQLAlchemyError as e: + logger.error(f"Error initializing database: {e}") + raise + def get_db(): + """ + Dependency for FastAPI routes or general DB usage. + Yields a SQLAlchemy session, and closes it upon exit. + """ db = SessionLocal() try: yield db