Kalletlamadhav/sql-optimization-env

1# data/seed_database.py2# Run once: python data/seed_database.py --rows 1000003 4import sqlite3, random, string, datetime, argparse, calendar5from pathlib import Path6 7SEED = 428random.seed(SEED)9 10STATES = {11 '01':'JK','02':'HP','03':'PB','06':'HR','07':'DL','08':'RJ',12 '09':'UP','10':'BR','11':'SK','12':'AR','13':'NL','14':'MN',13 '18':'AS','19':'WB','21':'OR','22':'CG','23':'MP','24':'GJ',14 '27':'MH','29':'KA','32':'KL','33':'TN','36':'TG','37':'AP'15}16 17STATE_CODES = list(STATES.keys())18 19HSN_CODES = ['0101','0201','0301','1001','1006','1701',20 '2101','2201','3004','4901','6101','7108',21 '8471','8517','9018','9401']22 23GST_RATES = [0, 5, 12, 18, 28]24 25 26def random_gstin(state_code: str) -> str:27 chars = string.ascii_uppercase + string.digits28 pan = ''.join(random.choices(string.ascii_uppercase, k=5))29 pan += ''.join(random.choices(string.digits, k=4))30 pan += random.choice(string.ascii_uppercase)31 suffix = '1Z' + random.choice(string.ascii_uppercase + string.digits)32 return f'{state_code}{pan}{suffix}'33 34 35def random_pnr() -> str:36 return ''.join(random.choices(string.digits, k=10))37 38 39def random_date(start='2023-01-01', end='2025-12-31') -> str:40 s = datetime.date.fromisoformat(start)41 e = datetime.date.fromisoformat(end)42 delta = (e - s).days43 return str(s + datetime.timedelta(days=random.randint(0, delta)))44 45 46def random_railway_travel_date(year: int = 2025) -> str:47 """Higher booking probability in March (Holi) and Oct–Nov (Diwali season)."""48 month_pool = []49 for month, weight in (50  (1, 2), (2, 2), (3, 8), (4, 2), (5, 2), (6, 2), (7, 2), (8, 2), (9, 3),51  (10, 8), (11, 8), (12, 3),52 ):53  month_pool.extend([month] * weight)54 month = random.choice(month_pool)55 _, last = calendar.monthrange(year, month)56 day = random.randint(1, last)57 return f'{year}-{month:02d}-{day:02d}'58 59 60def seed_gst(conn, n_invoices: int):61 print(f'Seeding GST: {n_invoices} invoices...')62 invoices = []63 items = []64 65 for i in range(n_invoices):66  state = random.choice(STATE_CODES)67  supplier = random_gstin(state)68  buyer_state = random.choice(STATE_CODES)69  buyer = random_gstin(buyer_state)70 71  is_igst = (state != buyer_state)72  taxable = round(random.uniform(1000, 500000), 2)73  rate = random.choice(GST_RATES)74  tax = round(taxable * rate / 100, 2)75 76  cgst = 0 if is_igst else round(tax/2, 2)77  sgst = 0 if is_igst else round(tax/2, 2)78  igst = tax if is_igst else 079 80  # ~5% intra-state invoices: deliberate CGST/SGST split mismatch (fraud / data-quality signal)81  if (not is_igst) and rate > 0 and random.random() < 0.05:82   half = round(tax / 2, 2)83   max_delta = min(half - 0.01, tax * 0.2, 5000)84   if max_delta > 1:85    delta = round(random.uniform(1, max_delta), 2)86    if random.random() < 0.5:87     cgst = round(half - delta, 2)88     sgst = round(half, 2)89    else:90     cgst = round(half, 2)91     sgst = round(half - delta, 2)92 93  inv_id = f'INV{i:08d}'94  inv_date = random_date()95 96  invoices.append((97   inv_id, supplier, buyer, inv_date,98   'B2B', taxable, cgst, sgst, igst, 0,99   state, random.choice(HSN_CODES), 'FILED',datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')100  ))101 102  for j in range(random.randint(1, 5)):103   item_taxable = round(taxable / random.randint(1,5), 2)104 105   items.append((106    inv_id, f'Item {j+1}', random.choice(HSN_CODES),107    round(random.uniform(1,100),2),108    round(item_taxable/random.randint(1,100),2),109    item_taxable, rate110   ))111 112 conn.executemany(113  'INSERT INTO gst_invoice_records VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)',114  invoices115 )116 117 conn.executemany(118  'INSERT INTO gst_invoice_items(invoice_id,item_description,hsn_code,quantity,unit_value,taxable_value,gst_rate) VALUES(?,?,?,?,?,?,?)',119  items120 )121 122 conn.commit()123 124 print(f' GST done: {len(invoices)} invoices, {len(items)} items')125 126 127def seed_pds(conn, n_cards: int):128 print(f'Seeding PDS: {n_cards} ration cards...')129 130 CARD_TYPES = ['APL','BPL','AAY','PHH']131 COMMODITIES = ['RICE','WHEAT','SUGAR','OIL']132 133 cards, allotments = [], []134 135 for i in range(n_cards):136  state = random.choice(STATE_CODES)137  card_id = f'RC{random.randint(2020,2025)}{state}{i:07d}'138  ctype = random.choice(CARD_TYPES)139 140  cards.append((141   card_id, f'Household_{i}', state,142   f'{random.randint(1,30):03d}',143   None, None, ctype,144   random.randint(1,8), random.randint(0,1),145   None, random_date('2020-01-01','2022-12-31'), None146  ))147 148  for month_offset in range(random.randint(1,24)):149   yr = 2024 + month_offset // 12150   mo = (month_offset % 12) + 1151   month_year = f'{yr}-{mo:02d}'152 153   for commodity in random.sample(COMMODITIES, k=random.randint(1,3)):154    entitled = round(random.uniform(5,35), 1)155 156    allotments.append((157     card_id, month_year, commodity,158     entitled,159     round(entitled * random.uniform(0.7,1.0),1),160     f'FS{random.randint(10000,99999)}',161     None162    ))163 164 conn.executemany(165  'INSERT INTO ration_card_beneficiaries VALUES(?,?,?,?,?,?,?,?,?,?,?,?)',166  cards167 )168 169 conn.executemany(170  'INSERT INTO pds_allotments(card_id,month_year,commodity,entitled_qty_kg,offtake_qty_kg,fair_shop_code,offtake_date) VALUES(?,?,?,?,?,?,?)',171  allotments172 )173 174 conn.commit()175 176 print(f' PDS done: {len(cards)} cards, {len(allotments)} allotments')177 178 179def seed_railway(conn, n_bookings: int):180 print(f'Seeding Railway: {n_bookings} bookings...')181 182 TRAINS = [183  ('12001','BHOPAL EXP','NDLS','BPL',72,'2A'),184  ('12951','RAJDHANI','NDLS','BCT',48,'1A'),185  ('17031','HYDERABAD EXP','NZB','HYB',72,'SL'),186  ('12723','TELANGANA EXP','SC','NDLS',72,'3A'),187  ('16093','LUCKNOW EXP','LKO','MAS',60,'SL'),188 ]189 190 conn.executemany('INSERT INTO railway_trains VALUES(?,?,?,?,?,?)', TRAINS)191 192 STATUSES = ['CNF']*70 + ['WL/'+str(i) for i in range(1,15)] + ['RAC/'+str(i) for i in range(1,6)]193 CLASSES = ['1A','2A','3A','SL','CC']194 195 bookings = []196 197 for i in range(n_bookings):198  train = random.choice(TRAINS)199 200  bookings.append((201   random_pnr(),202   train[0],203   random_railway_travel_date(2025),204   f'Passenger_{i}',205   random.randint(5,80),206   random.choice(['M','F']),207   random.choice(STATUSES),208   random.choice(CLASSES),209   f'{random.randint(1,72)}',210   round(random.uniform(200, 3500), 2),211   None,212   random.choice(['WEB','APP','TATKAL','COUNTER'])213  ))214 215 conn.executemany(216  'INSERT INTO railway_pnr_bookings VALUES(?,?,?,?,?,?,?,?,?,?,?,?)',217  bookings218 )219 220 conn.commit()221 222 print(f' Railway done: {len(bookings)} bookings')223 224 225def seed_mgnrega(conn, n_workers: int):226 print(f'Seeding MGNREGA: {n_workers} workers...')227 228 WAGE_RATES = {'MH':273,'TG':257,'KA':309,'TN':294,'AP':257,'UP':213,'BR':194}229 230 workers, attendance, payments = [], [], []231 232 for i in range(n_workers):233  state = random.choice(list(WAGE_RATES.keys()))234  worker_id = f'WK{state}{i:07d}'235  wage = WAGE_RATES[state]236 237  workers.append((238   worker_id, f'Worker_{i}', state,239   f'{random.randint(1,30):03d}',240   f'GP{random.randint(1000,9999)}',241   f'JC{state}{i:08d}', None, None, None, wage242  ))243 244  n_days = random.randint(0, 150)245 246  for d in range(n_days):247   work_date = random_date('2024-04-01','2025-03-31')248 249   attendance.append((250    worker_id, work_date,251    f'PROJ{random.randint(10000,99999)}',252    random.choice([0.5, 1.0]),253    f'MR{random.randint(10000,99999)}', 1254   ))255 256  for m in range(12):257   yr = 2024 + m // 12258   mo = (m % 12) + 4259   if mo > 12:260    mo -= 12261    yr += 1262 263   days = round(random.uniform(0, 25), 1)264   amount_due = round(days * wage, 2)265   amount_paid = amount_due if random.random() > 0.1 else 0266 267   payments.append((268    worker_id, f'{yr}-{mo:02d}', days,269    amount_due, amount_paid,270    random_date(f'{yr}-{mo:02d}-01', f'{yr}-{mo:02d}-28') if amount_paid > 0 else None,271    f'UTR{random.randint(100000000,999999999)}' if amount_paid > 0 else None272   ))273 274 conn.executemany(275  'INSERT INTO mgnrega_workers VALUES(?,?,?,?,?,?,?,?,?,?)',276  workers277 )278 279 conn.executemany(280  'INSERT INTO mgnrega_attendance(worker_id,work_date,project_code,days_worked,muster_roll_no,verified) VALUES(?,?,?,?,?,?)',281  attendance282 )283 284 conn.executemany(285  'INSERT INTO mgnrega_payments(worker_id,payment_month,days_worked,amount_due,amount_paid,payment_date,utr_number) VALUES(?,?,?,?,?,?,?)',286  payments287 )288 289 conn.commit()290 291 print(f' MGNREGA done: {len(workers)} workers, {len(attendance)} attendance, {len(payments)} payments')292 293 294def create_all_tables(conn):295 schema_dir = Path(__file__).resolve().parent / 'schemas'296 297 for sql_file in ['gst_schema.sql','pds_schema.sql','railway_schema.sql','mgnrega_schema.sql']:298  with open(schema_dir / sql_file) as f:299   conn.executescript(f.read())300 301 302def seed_database(db, rows):303    from pathlib import Path304    import sqlite3305 306    Path(db).parent.mkdir(parents=True, exist_ok=True)307 308    conn = sqlite3.connect(db)309    conn.execute('PRAGMA foreign_keys = ON')310    conn.execute('PRAGMA journal_mode = WAL')311 312    create_all_tables(conn)313 314    n = rows315    seed_gst(conn, n)316    seed_pds(conn, n // 5)317    seed_railway(conn, n)318    seed_mgnrega(conn, n // 3)319 320    conn.close()321 322    print(f'Database seeded successfully: {db}')323 324 325def main():326 parser = argparse.ArgumentParser()327 parser.add_argument('--rows', type=int, default=1000)328 parser.add_argument('--db', default='data/fixtures/benchmark_seed42.db')329 330 args = parser.parse_args()331 332 seed_database(args.db, args.rows)333 334 335if __name__ == '__main__':336 main()