yaml_shredder.data_loader

  1"""Load generated tables into databases."""
  2
  3import sqlite3
  4from pathlib import Path
  5
  6import pandas as pd
  7
  8
  9class SQLiteLoader:
 10    """Load DataFrames into SQLite database."""
 11
 12    def __init__(self, db_path: str | Path):
 13        """
 14        Initialize SQLite loader.
 15
 16        Args:
 17            db_path: Path to SQLite database file (will be created if doesn't exist)
 18        """
 19        self.db_path = Path(db_path)
 20        self.connection = None
 21        self.loaded_tables = []
 22
 23    def connect(self) -> None:
 24        """Establish connection to SQLite database."""
 25        self.connection = sqlite3.connect(str(self.db_path))
 26        print(f"Connected to SQLite database: {self.db_path}")
 27
 28    def disconnect(self) -> None:
 29        """Close database connection."""
 30        if self.connection:
 31            self.connection.close()
 32            self.connection = None
 33            print("Database connection closed")
 34
 35    def load_tables(
 36        self,
 37        tables: dict[str, pd.DataFrame],
 38        if_exists: str = "replace",
 39        create_indexes: bool = True,
 40    ) -> None:
 41        """
 42        Load multiple tables into SQLite.
 43
 44        Args:
 45            tables: Dictionary of table_name -> DataFrame
 46            if_exists: How to behave if table exists ('fail', 'replace', 'append')
 47            create_indexes: Whether to create indexes on foreign key columns
 48        """
 49        if not self.connection:
 50            self.connect()
 51
 52        self.loaded_tables = []
 53
 54        for table_name, df in tables.items():
 55            self._load_table(table_name, df, if_exists)
 56            self.loaded_tables.append(table_name)
 57
 58            if create_indexes:
 59                self._create_indexes(table_name, df)
 60
 61        print(f"\n✓ Loaded {len(tables)} tables into {self.db_path}")
 62
 63    def _load_table(self, table_name: str, df: pd.DataFrame, if_exists: str = "replace") -> None:
 64        """
 65        Load a single table into SQLite.
 66
 67        Args:
 68            table_name: Name of the table
 69            df: DataFrame to load
 70            if_exists: How to behave if table exists
 71        """
 72        # Convert boolean columns to integers for SQLite
 73        df_copy = df.copy()
 74        for col in df_copy.columns:
 75            if df_copy[col].dtype == "bool":
 76                df_copy[col] = df_copy[col].astype(int)
 77
 78        df_copy.to_sql(table_name, self.connection, if_exists=if_exists, index=False)
 79        print(f"  Loaded {len(df_copy)} rows into table: {table_name}")
 80
 81    def _create_indexes(self, table_name: str, df: pd.DataFrame) -> None:
 82        """
 83        Create indexes on key columns.
 84
 85        Args:
 86            table_name: Name of the table
 87            df: DataFrame to analyze for index creation
 88        """
 89        cursor = self.connection.cursor()
 90
 91        # Create index on id column if exists
 92        if "id" in [c.lower() for c in df.columns]:
 93            id_col = next(c for c in df.columns if c.lower() == "id")
 94            index_name = f"idx_{table_name}_id"
 95            try:
 96                cursor.execute(f'CREATE INDEX IF NOT EXISTS "{index_name}" ON "{table_name}" ("{id_col}")')
 97            except sqlite3.OperationalError:
 98                pass  # Index might already exist
 99
100        # Create indexes on foreign key columns (parent_*)
101        fk_columns = [c for c in df.columns if c.startswith("parent_")]
102        for fk_col in fk_columns:
103            index_name = f"idx_{table_name}_{fk_col}"
104            try:
105                cursor.execute(f'CREATE INDEX IF NOT EXISTS "{index_name}" ON "{table_name}" ("{fk_col}")')
106            except sqlite3.OperationalError:
107                pass
108
109        self.connection.commit()
110
111    def execute_ddl(self, ddl_statements: dict[str, str]) -> None:
112        """
113        Execute DDL statements to create tables.
114
115        Args:
116            ddl_statements: Dictionary of table_name -> CREATE TABLE statement
117        """
118        if not self.connection:
119            self.connect()
120
121        cursor = self.connection.cursor()
122
123        for table_name, ddl in ddl_statements.items():
124            # SQLite doesn't support ALTER TABLE ADD CONSTRAINT for foreign keys
125            # Skip ALTER statements
126            statements = ddl.split(";")
127            for stmt in statements:
128                stmt = stmt.strip()
129                if stmt and not stmt.startswith("ALTER TABLE"):
130                    try:
131                        cursor.execute(stmt)
132                    except sqlite3.OperationalError as e:
133                        print(f"Warning: Could not execute DDL for {table_name}: {e}")
134
135        self.connection.commit()
136        print(f"✓ Executed DDL for {len(ddl_statements)} tables")
137
138    def query(self, sql: str) -> pd.DataFrame:
139        """
140        Execute a SQL query and return results as DataFrame.
141
142        Args:
143            sql: SQL query to execute
144
145        Returns:
146            Query results as DataFrame
147        """
148        if not self.connection:
149            self.connect()
150
151        return pd.read_sql_query(sql, self.connection)
152
153    def get_table_info(self, table_name: str) -> pd.DataFrame:
154        """
155        Get schema information for a table.
156
157        Args:
158            table_name: Name of the table
159
160        Returns:
161            DataFrame with table schema information
162        """
163        return self.query(f'PRAGMA table_info("{table_name}")')
164
165    def list_tables(self) -> list[str]:
166        """
167        List all tables in the database.
168
169        Returns:
170            List of table names
171        """
172        cursor = self.connection.cursor()
173        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
174        return [row[0] for row in cursor.fetchall()]
175
176    def print_summary(self) -> None:
177        """Print summary of loaded database."""
178        if not self.connection:
179            print("Not connected to database")
180            return
181
182        print(f"\n{'=' * 60}")
183        print("DATABASE SUMMARY")
184        print(f"{'=' * 60}")
185        print(f"Database: {self.db_path}")
186        print(f"Size: {self.db_path.stat().st_size / 1024:.1f} KB")
187
188        tables = self.list_tables()
189        print(f"\nTables: {len(tables)}")
190
191        for table in tables:
192            cursor = self.connection.cursor()
193            cursor.execute(f'SELECT COUNT(*) FROM "{table}"')
194            count = cursor.fetchone()[0]
195            print(f"  - {table}: {count} rows")
196
197    def __enter__(self):
198        """Context manager entry."""
199        self.connect()
200        return self
201
202    def __exit__(self, exc_type, exc_val, exc_tb):
203        """Context manager exit."""
204        self.disconnect()
205
206
207def load_to_sqlite(
208    tables: dict[str, pd.DataFrame],
209    db_path: str | Path,
210    if_exists: str = "replace",
211    create_indexes: bool = True,
212) -> SQLiteLoader:
213    """
214    Quick function to load tables into SQLite.
215
216    Args:
217        tables: Dictionary of table_name -> DataFrame
218        db_path: Path to SQLite database
219        if_exists: How to behave if table exists
220        create_indexes: Whether to create indexes
221
222    Returns:
223        SQLiteLoader instance (still connected)
224    """
225    loader = SQLiteLoader(db_path)
226    loader.connect()
227    loader.load_tables(tables, if_exists, create_indexes)
228    return loader