boc/atm-anomaly-detection/scripts/setup_db.py

"""
Database Setup Script

Creates database schema for ATM anomaly detection.
Supports PostgreSQL and SQLite.
"""

import os
import sys
import argparse
from pathlib import Path

def setup_sqlite(db_path: str = 'data/atm_anomaly.db'):
    """Setup SQLite database."""
    import sqlite3
    
    os.makedirs(os.path.dirname(db_path), exist_ok=True)
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # Read schema
    schema_path = Path(__file__).parent.parent / 'config' / 'database.sql'
    with open(schema_path, 'r') as f:
        schema = f.read()
    
    # Execute schema (SQLite compatible)
    # Replace PostgreSQL-specific syntax
    schema = schema.replace('SERIAL PRIMARY KEY', 'INTEGER PRIMARY KEY AUTOINCREMENT')
    schema = schema.replace('JSONB', 'JSON')
    schema = schema.replace('DECIMAL(10, 8)', 'REAL')
    schema = schema.replace('DECIMAL(11, 8)', 'REAL')
    schema = schema.replace('DECIMAL(10, 2)', 'REAL')
    schema = schema.replace('DECIMAL(10, 6)', 'REAL')
    schema = schema.replace('DECIMAL(5, 4)', 'REAL')
    schema = schema.replace('TIMESTAMP', 'DATETIME')
    schema = schema.replace('CHECK (severity_level BETWEEN 1 AND 5)', '')
    schema = schema.replace('CHECK (priority BETWEEN 1 AND 5)', '')
    
    # Split and execute statements
    statements = schema.split(';')
    for stmt in statements:
        stmt = stmt.strip()
        if stmt:
            try:
                cursor.execute(stmt)
            except sqlite3.Error as e:
                print(f"Warning: {e}")
                print(f"Statement: {stmt[:100]}...")
    
    conn.commit()
    conn.close()
    
    print(f"SQLite database created: {db_path}")


def setup_postgres(connection_string: str):
    """Setup PostgreSQL database."""
    import psycopg2
    
    conn = psycopg2.connect(connection_string)
    cursor = conn.cursor()
    
    schema_path = Path(__file__).parent.parent / 'config' / 'database.sql'
    with open(schema_path, 'r') as f:
        schema = f.read()
    
    cursor.execute(schema)
    conn.commit()
    conn.close()
    
    print("PostgreSQL database initialized")


def seed_demo_data(db_path: str = 'data/atm_anomaly.db'):
    """Insert demo data for testing."""
    import sqlite3
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # Insert demo ATMs
    atms = [
        ('ATM-001', 'Swedbank', 'Stockholm Central', 'Sergels Torg 1, Stockholm', 59.3326, 18.0649, 'Stockholm', 'Sweden'),
        ('ATM-002', 'SEB', 'Göteborg Central', 'Drottningtorget 2, Göteborg', 57.7089, 11.9746, 'Göteborg', 'Sweden'),
        ('ATM-003', 'Nordea', 'Malmö Central', 'Centralplan 1, Malmö', 55.6090, 13.0007, 'Malmö', 'Sweden'),
    ]
    
    cursor.executemany('''
        INSERT OR IGNORE INTO atm_locations 
        (atm_id, bank_name, branch_name, address, latitude, longitude, city, country)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
    ''', atms)
    
    # Insert demo captures
    captures = [
        ('ATM-001', '2026-07-11 06:00:00', 'front', 'data/raw/atm_001_20260711060000_front.jpg', 'day'),
        ('ATM-001', '2026-07-11 06:05:00', 'side', 'data/raw/atm_001_20260711060500_side.jpg', 'day'),
        ('ATM-002', '2026-07-11 06:00:00', 'front', 'data/raw/atm_002_20260711060000_front.jpg', 'day'),
    ]
    
    cursor.executemany('''
        INSERT INTO atm_captures 
        (atm_id, capture_timestamp, camera_angle, image_path, lighting_condition)
        VALUES (?, ?, ?, ?, ?)
    ''', captures)
    
    conn.commit()
    conn.close()
    
    print("Demo data inserted")


def main():
    parser = argparse.ArgumentParser(description='Setup ATM Anomaly Database')
    parser.add_argument('--db-type', choices=['sqlite', 'postgres'], default='sqlite')
    parser.add_argument('--connection', help='PostgreSQL connection string')
    parser.add_argument('--db-path', default='data/atm_anomaly.db', help='SQLite database path')
    parser.add_argument('--seed', action='store_true', help='Insert demo data')
    
    args = parser.parse_args()
    
    if args.db_type == 'sqlite':
        setup_sqlite(args.db_path)
        if args.seed:
            seed_demo_data(args.db_path)
    elif args.db_type == 'postgres':
        if not args.connection:
            print("Error: --connection required for PostgreSQL")
            sys.exit(1)
        setup_postgres(args.connection)
    
    print("Database setup complete!")


if __name__ == "__main__":
    main()