Spaces:
Sleeping
Sleeping
| import asyncio | |
| from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession | |
| from sqlalchemy.orm import sessionmaker | |
| from sqlalchemy import text, select, and_ | |
| import os | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| DATABASE_URL = os.getenv("DATABASE_URL") | |
| engine = create_async_engine(DATABASE_URL) | |
| async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) | |
| async def migrate(): | |
| from app.models.stock_ledger import StockLedger | |
| from app.models.transaction import Transaction, TransactionItem | |
| from app.models.patti_transaction import PattiTransaction, PattiTransactionItem | |
| async with engine.connect() as conn: | |
| print("Adding poti columns to stock_ledger...") | |
| await conn.execute(text("ALTER TABLE stock_ledger ADD COLUMN IF NOT EXISTS poti_in INTEGER DEFAULT 0")) | |
| await conn.execute(text("ALTER TABLE stock_ledger ADD COLUMN IF NOT EXISTS poti_out INTEGER DEFAULT 0")) | |
| await conn.execute(text("ALTER TABLE stock_ledger ADD COLUMN IF NOT EXISTS poti_balance_after INTEGER DEFAULT 0")) | |
| await conn.commit() | |
| async with async_session() as session: | |
| print("Populating poti data and calculating bag balances...") | |
| # We need to recalculate balances per (mirchi_name, lot_number) | |
| # First, get all entries ordered by creation | |
| res = await session.execute(select(StockLedger).order_by(StockLedger.created_at.asc())) | |
| entries = res.scalars().all() | |
| balances = {} # (mirchi_name, lot_number) -> current_poti_balance | |
| for entry in entries: | |
| poti_count = 0 | |
| # Find the item to get the poti_count | |
| if entry.transaction_id: | |
| item_res = await session.execute( | |
| select(TransactionItem).where( | |
| and_( | |
| TransactionItem.transaction_id == entry.transaction_id, | |
| TransactionItem.mirchi_name == entry.mirchi_name | |
| ) | |
| ) | |
| ) | |
| item = item_res.scalar_one_or_none() | |
| if item: | |
| poti_count = item.poti_count | |
| elif entry.patti_transaction_id: | |
| item_res = await session.execute( | |
| select(PattiTransactionItem).where( | |
| and_( | |
| PattiTransactionItem.transaction_id == entry.patti_transaction_id, | |
| PattiTransactionItem.mirchi_name == entry.mirchi_name | |
| ) | |
| ) | |
| ) | |
| item = item_res.scalar_one_or_none() | |
| if item: | |
| poti_count = item.poti_count | |
| # Update entry with poti_in/poti_out | |
| if entry.movement_type in ['IN', 'RETURN_IN']: | |
| entry.poti_in = poti_count | |
| entry.poti_out = 0 | |
| else: | |
| entry.poti_in = 0 | |
| entry.poti_out = poti_count | |
| # Calculate running balance | |
| key = (entry.mirchi_name, entry.lot_number) | |
| current_bal = balances.get(key, 0) | |
| if entry.movement_type in ['IN', 'RETURN_IN']: | |
| current_bal += poti_count | |
| else: | |
| current_bal -= poti_count | |
| entry.poti_balance_after = current_bal | |
| balances[key] = current_bal | |
| await session.commit() | |
| print("Migration and poti balance recalculation completed!") | |
| if __name__ == "__main__": | |
| asyncio.run(migrate()) | |