Spaces:
Sleeping
Sleeping
File size: 3,006 Bytes
e8efc4f 2a944a5 e8efc4f 2a944a5 e8efc4f 2a944a5 e8efc4f 2a944a5 e8efc4f 2a944a5 e8efc4f 2a944a5 e8efc4f 2a944a5 e8efc4f 2a944a5 e8efc4f 2a944a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
# src/repository.py
from sqlmodel import Session, select
from typing import Dict, Any, Optional
import json
from datetime import date
from src.models import Invoice, LineItem
from src.database import get_session, engine, DB_CONNECTED
class InvoiceRepository:
def __init__(self, session: Optional[Session] = None):
"""
Initialize with an optional session.
If no session is provided, try to get a new one from the engine.
Only creates session if database is actually connected.
"""
if session:
self.session = session
elif engine and DB_CONNECTED:
self.session = Session(engine)
else:
self.session = None
def save_invoice(self, invoice_data: Dict[str, Any]) -> Optional[Invoice]:
"""
Saves an invoice and its line items to the database.
Returns the saved Invoice object or None if DB is disabled/failed.
"""
if not self.session:
print("⚠️ DB Session missing. Skipping save.")
return None
try:
# 1. Prepare Data
data = invoice_data.copy()
# Serialize complex types (validation_errors)
if 'validation_errors' in data and isinstance(data['validation_errors'], list):
data['validation_errors'] = json.dumps(data['validation_errors'])
# Extract items to process separately
items_data = data.pop('items', [])
# 2. Create Invoice Record
invoice = Invoice(**data)
# 3. Process Items
for item in items_data:
# Ensure item is a dict (if it's a Pydantic model, convert it)
if hasattr(item, 'model_dump'):
item_dict = item.model_dump()
elif isinstance(item, dict):
item_dict = item
else:
continue
line_item = LineItem(**item_dict)
invoice.items.append(line_item)
# 4. Commit
self.session.add(invoice)
self.session.commit()
self.session.refresh(invoice)
print(f"✅ Invoice {invoice.id} saved to DB.")
return invoice
except Exception as e:
print(f"❌ Error saving invoice to DB: {e}")
self.session.rollback()
return None
def get_by_hash(self, semantic_hash: str) -> Optional[Invoice]:
"""
Check if invoice already exists using the semantic hash.
"""
if not self.session:
return None
try:
statement = select(Invoice).where(Invoice.semantic_hash == semantic_hash)
results = self.session.exec(statement)
return results.first()
except Exception as e:
print(f"❌ Error checking hash: {e}")
return None |