import sqlite3 import json import numpy as np from typing import Dict, Any, List, Tuple, Optional import os import datetime import statistics import re from collections import defaultdict # DB-specific table-specific columns to skip # Format: {'database_name': {'table_name': ['column1', 'column2', ...]}} DB_TABLE_SPECIFIC_SKIP_COLUMNS = { 'Db-IMDB.sqlite': { 'Movie': ['index'], 'Genre': ['index'], 'Language': ['index'], 'Country': ['index'], 'Location': ['index'], 'M_Location': ['index', 'ID'], 'M_Country': ['index', 'ID'], 'M_Language': ['index', 'ID'], 'M_Genre': ['index', 'ID'], 'Person': ['index'], 'M_Producer': ['index', 'ID'], 'M_Director': ['index', 'ID'], 'M_Cast': ['index', 'ID'] } # Add more database and table-specific skip rules as needed } def quote_column_name(col_name): """Properly quote column names that need it""" if not col_name: return '""' # Special characters that require quoting special_chars = ['%', '-', '+', '*', '/', '(', ')', '[', ']', ' ', '.', '&', '|', '!', '@', '#', '$', '^', '~', '`', '=', '<', '>', '?', ',', ';', ':', "'", '"'] # SQL reserved words (common ones) reserved_words = [ 'index', 'key', 'order', 'group', 'from', 'select', 'where', 'table', 'column', 'and', 'or', 'not', 'in', 'is', 'null', 'like', 'between', 'exists', 'case', 'when', 'then', 'else', 'end', 'union', 'all', 'distinct', 'having', 'limit', 'offset', 'join', 'inner', 'outer', 'left', 'right', 'full', 'cross', 'on', 'using', 'natural', 'as', 'desc', 'asc', 'primary', 'foreign', 'references', 'constraint', 'unique', 'check', 'default', 'create', 'drop', 'alter', 'insert', 'update', 'delete', 'into', 'values', 'set', 'truncate', 'commit', 'rollback', 'transaction', 'begin', 'savepoint', 'release', 'pragma', 'vacuum', 'analyze', 'explain', 'view', 'trigger', 'procedure', 'function', 'database', 'schema' ] # Check if quoting is needed needs_quoting = ( any(char in col_name for char in special_chars) or col_name.lower() in reserved_words or col_name.isdigit() or (col_name and col_name[0].isdigit()) ) if needs_quoting: # Escape any existing double quotes by doubling them escaped_name = col_name.replace('"', '""') return f'"{escaped_name}"' return col_name def quote_table_name(table_name): """Properly quote table names that need it""" return quote_column_name(table_name) # Same logic applies class SQLiteKeyFinder: """ A tool to identify potential primary and foreign keys in SQLite databases where they are not explicitly defined. """ def __init__(self, db_path): """Initialize with the path to the SQLite database.""" if not os.path.exists(db_path): raise FileNotFoundError(f"Database file not found: {db_path}") self.db_path = db_path self.db_name = os.path.basename(db_path) self.conn = sqlite3.connect(db_path) self.cursor = self.conn.cursor() self.tables = [] self.table_columns = {} self.primary_keys = {} self.foreign_keys = defaultdict(list) def should_skip_column(self, table_name, column_name): """ Determine if a column should be skipped based on database-specific table rules. """ if self.db_name in DB_TABLE_SPECIFIC_SKIP_COLUMNS and \ table_name in DB_TABLE_SPECIFIC_SKIP_COLUMNS[self.db_name] and \ column_name.lower() in [col.lower() for col in DB_TABLE_SPECIFIC_SKIP_COLUMNS[self.db_name][table_name]]: return True return False def _get_tables(self): """Get all tables in the database.""" self.cursor.execute("SELECT name FROM sqlite_master WHERE type='table';") self.tables = [row[0] for row in self.cursor.fetchall() if not row[0].startswith('sqlite_')] return self.tables def _get_table_columns(self): """Get all columns for each table.""" for table in self.tables: quoted_table = quote_table_name(table) self.cursor.execute(f"PRAGMA table_info({quoted_table});") columns = [] for col_info in self.cursor.fetchall(): col_name = col_info[1] data_type = col_info[2] columns.append((col_name, data_type)) self.table_columns[table] = columns return self.table_columns def _check_defined_keys(self): """Check for already defined primary and foreign keys.""" defined_pk = {} defined_fk = defaultdict(list) # Check primary keys for table in self.tables: quoted_table = quote_table_name(table) self.cursor.execute(f"PRAGMA table_info({quoted_table});") pk_columns = [] for col_info in self.cursor.fetchall(): if col_info[5] > 0: # pk column pk_columns.append(col_info[1]) if pk_columns: defined_pk[table] = {'columns': pk_columns, 'origin': 'db'} # Check foreign keys for table in self.tables: quoted_table = quote_table_name(table) self.cursor.execute(f"PRAGMA foreign_key_list({quoted_table});") for fk_info in self.cursor.fetchall(): defined_fk[table].append({ 'from': fk_info[3], 'to_table': fk_info[2], 'to_column': fk_info[4], 'origin': 'db' }) return defined_pk, defined_fk def find_potential_primary_keys(self): """Identify columns that are likely to be primary keys.""" # First check for explicitly defined keys defined_pk, _ = self._check_defined_keys() if defined_pk: print("Found defined primary keys:", defined_pk) self.primary_keys.update(defined_pk) # For tables without defined primary keys, try to identify them for table in self.tables: if table in self.primary_keys: continue quoted_table = quote_table_name(table) self.cursor.execute(f"SELECT COUNT(*) FROM {quoted_table};") total_rows = self.cursor.fetchone()[0] if total_rows == 0: continue pk_candidates = {} for col_name, data_type in self.table_columns[table]: if self.should_skip_column(table, col_name): continue try: pk_score = 0 # Handle reserved keyword column names quoted_col_name = quote_column_name(col_name) # Check uniqueness self.cursor.execute(f"SELECT COUNT(DISTINCT {quoted_col_name}) FROM {quoted_table};") distinct_values = self.cursor.fetchone()[0] uniqueness_ratio = distinct_values / max(1, total_rows) if uniqueness_ratio < 0.9: continue # Check for NULL values self.cursor.execute(f"SELECT COUNT(*) FROM {quoted_table} WHERE {quoted_col_name} IS NULL;") null_count = self.cursor.fetchone()[0] null_ratio = null_count / max(1, total_rows) if null_ratio > 0.1: continue # Score uniqueness if uniqueness_ratio == 1.0: pk_score += 30 elif uniqueness_ratio > 0.98: pk_score += 20 # No nulls is good for PKs if null_count == 0: pk_score += 20 # Data type scoring if 'int' in data_type.lower(): pk_score += 15 elif data_type.lower() in ['text', 'varchar', 'char', 'string']: pk_score += 5 # Check for naming patterns name_patterns = [ (r'^id$', 15), (r'^{}_id$'.format(table), 15), (r'^{}_key$'.format(table), 15), (r'^pk_', 15), (r'^key$', 10), (r'^code$', 8), (r'^uuid$', 15), (r'^guid$', 15), (r'^serial$', 15), (r'^seq', 10), (r'id$', 5), (r'uuid$', 10), (r'code$', 5), (r'num$', 5), (r'no$', 5), (r'^record', 8), (r'^pid$', 15), (r'^mid$', 15), (r'^uid$', 15), (r'^eid$', 15), (r'^[a-z]+id$', 10), (r'^[a-z]+_id$', 10), ] for pattern, score in name_patterns: if re.search(pattern, col_name, re.IGNORECASE): pk_score += score break # Check for auto-increment indication try: self.cursor.execute(f"SELECT MIN({quoted_col_name}), MAX({quoted_col_name}) FROM {quoted_table};") min_val, max_val = self.cursor.fetchone() if (isinstance(min_val, int) and isinstance(max_val, int) and min_val in [0, 1] and max_val == total_rows): pk_score += 15 if total_rows <= 1000: self.cursor.execute(f""" WITH numbers AS ( SELECT ROW_NUMBER() OVER (ORDER BY {quoted_col_name}) + {min_val} - 1 as expected, {quoted_col_name} as actual FROM {quoted_table} ORDER BY {quoted_col_name} ) SELECT COUNT(*) FROM numbers WHERE expected != actual; """) gaps = self.cursor.fetchone()[0] if gaps == 0: pk_score += 10 except (sqlite3.OperationalError, TypeError): pass # Check for index presence self.cursor.execute(f"PRAGMA index_list({quoted_table});") indexes = self.cursor.fetchall() for idx_info in indexes: idx_name = idx_info[1] is_unique = idx_info[2] self.cursor.execute(f"PRAGMA index_info({quote_column_name(idx_name)});") index_columns = [info[2] for info in self.cursor.fetchall()] if col_name in index_columns and len(index_columns) == 1: pk_score += 10 if is_unique: pk_score += 10 # Add to candidates if score is high enough if pk_score >= 25: pk_candidates[col_name] = { 'score': pk_score, 'data_type': data_type, 'uniqueness': uniqueness_ratio, 'null_ratio': null_ratio } except sqlite3.OperationalError as e: print(f"Error analyzing column {table}.{col_name}: {e}") continue # Select the best primary key candidate(s) if pk_candidates: sorted_candidates = sorted( pk_candidates.items(), key=lambda x: x[1]['score'], reverse=True ) pk_columns = [] threshold_score = sorted_candidates[0][1]['score'] * 0.8 for col_name, info in sorted_candidates: if info['score'] >= threshold_score: pk_columns.append(col_name) self.primary_keys[table] = { 'columns': pk_columns, 'origin': 'potential' } return self.primary_keys def find_potential_foreign_keys(self): """Identify columns that are likely to be foreign keys.""" # First check for explicitly defined foreign keys _, defined_fk = self._check_defined_keys() if defined_fk: print("Found defined foreign keys:", dict(defined_fk)) self.foreign_keys.update(defined_fk) # Find potential foreign keys for src_table in self.tables: defined_fks_for_table = self.foreign_keys.get(src_table, []) defined_fk_cols = set(fk['from'] for fk in defined_fks_for_table) processed_relationships = set() for fk in defined_fks_for_table: processed_relationships.add((fk['from'], fk['to_table'], fk['to_column'])) quoted_src_table = quote_table_name(src_table) for src_col_name, src_data_type in self.table_columns[src_table]: if self.should_skip_column(src_table, src_col_name): continue if src_col_name in defined_fk_cols: continue quoted_src_col_name = quote_column_name(src_col_name) for ref_table in self.tables: if src_table == ref_table: continue if ref_table not in self.primary_keys: continue ref_col_list = self.primary_keys[ref_table].get('columns', []) if not ref_col_list: continue quoted_ref_table = quote_table_name(ref_table) for ref_col in ref_col_list: if self.should_skip_column(ref_table, ref_col): continue quoted_ref_col = quote_column_name(ref_col) if (src_col_name, ref_table, ref_col) in processed_relationships: continue # Define naming patterns for foreign keys fk_patterns = [ r'^{}_{}$'.format(ref_table, ref_col), r'^{}{}$'.format(ref_table, ref_col.capitalize()), r'^{}_id$'.format(ref_table), r'^{}Id$'.format(ref_table), r'^{}_key$'.format(ref_table), r'^fk_{}_'.format(ref_table), r'^{}$'.format(ref_col) ] name_pattern_match = False for pattern in fk_patterns: if re.match(pattern, src_col_name, re.IGNORECASE): name_pattern_match = True break ref_col_type = next((dtype for col, dtype in self.table_columns[ref_table] if col == ref_col), None) if name_pattern_match: if ref_col_type and src_data_type and ref_col_type.lower() != src_data_type.lower(): confidence = "low" else: confidence = "medium" self.foreign_keys[src_table].append({ 'from': src_col_name, 'to_table': ref_table, 'to_column': ref_col, 'origin': 'potential', 'confidence': confidence }) processed_relationships.add((src_col_name, ref_table, ref_col)) # Check data value matching else: try: if self.get_row_count(src_table) > 10000 or self.get_row_count(ref_table) > 10000: continue self.cursor.execute(f""" SELECT COUNT(*) FROM {quoted_src_table} WHERE {quoted_src_col_name} IS NOT NULL AND {quoted_src_col_name} NOT IN ( SELECT {quoted_ref_col} FROM {quoted_ref_table} ) """) invalid_refs = self.cursor.fetchone()[0] if invalid_refs == 0: self.cursor.execute(f""" SELECT COUNT(DISTINCT {quoted_src_col_name}) FROM {quoted_src_table} WHERE {quoted_src_col_name} IS NOT NULL """) distinct_values = self.cursor.fetchone()[0] self.cursor.execute(f""" SELECT COUNT(DISTINCT {quoted_ref_col}) FROM {quoted_ref_table} WHERE {quoted_ref_col} IS NOT NULL """) ref_distinct = self.cursor.fetchone()[0] coverage_ratio = distinct_values / max(1, ref_distinct) if (coverage_ratio > 0.01 and (not ref_col_type or not src_data_type or ref_col_type.lower() == src_data_type.lower())): self.foreign_keys[src_table].append({ 'from': src_col_name, 'to_table': ref_table, 'to_column': ref_col, 'origin': 'potential', 'confidence': 'high' if coverage_ratio > 0.3 else 'medium', 'evidence': 'data_match' }) processed_relationships.add((src_col_name, ref_table, ref_col)) except sqlite3.OperationalError as e: print(f"Error checking foreign key relationship {src_table}.{src_col_name} -> {ref_table}.{ref_col}: {e}") continue return dict(self.foreign_keys) def get_row_count(self, table): """Get the number of rows in a table.""" quoted_table = quote_table_name(table) self.cursor.execute(f"SELECT COUNT(*) FROM {quoted_table};") return self.cursor.fetchone()[0] def analyze(self): """Run the full analysis to find potential primary and foreign keys.""" print(f"Analyzing database: {self.db_path}") self._get_tables() print(f"Found {len(self.tables)} tables: {', '.join(self.tables)}") self._get_table_columns() print("\nFinding potential primary keys...") pk_results = self.find_potential_primary_keys() print("\nFinding potential foreign keys...") fk_results = self.find_potential_foreign_keys() return { 'tables': self.tables, 'columns': self.table_columns, 'primary_keys': pk_results, 'foreign_keys': fk_results } def close(self): """Close the database connection.""" if self.conn: self.conn.close() def extract_db_summary_for_schema(db_path: str, sample_limit=5000, include_samples=True, include_column_names=True, include_data_types=True, detect_primary_keys=True, detect_foreign_keys=True, include_key_confidence=True, include_row_count=True, # Enable: Critical for JOIN optimization include_column_count=True, include_distinct_count=True, # Enable: Cardinality for indexing decisions include_null_count=True, # Enable: NULL handling in queries include_cardinality=True, # Enable: Index selectivity include_nullability=True, # Enable: NULL safety validation include_min_max=True, # Enable: Range query optimization include_average=False, # Optional: Less critical for SQL include_median=False, # Optional: Less critical for SQL include_stddev=False, # Optional: Less critical for SQL include_avg_length=True, # Enable: String column sizing include_common_values=True, # Enable: Pattern recognition common_values_limit=5, common_values_threshold=100, include_date_range=True, # Enable: Time-based query optimization include_date_range_days=True, # Enable: Date range insights include_not_null_constraint=True, # Enable: Data quality validation include_default_values=True, # Enable: INSERT statement help include_indexes=True, # Enable: Performance optimization max_rows_for_expensive_stats=10000, max_string_display_length=100, max_binary_display_bytes=50, include_db_metadata=True, # Enable: Database-level info include_table_metadata=True, # Enable: Table-level info include_extraction_timestamp=True, # Enable: Freshness tracking skip_empty_tables=True, # Enable: Focus on relevant tables include_table_relationships=True, include_schema_summary=True): """ Extract comprehensive database summary including schema, statistics, and sample data. Args: db_path: Path to SQLite database file sample_limit: Number of sample rows per table include_samples: Include sample data include_column_names: Include column names include_data_types: Include data types detect_primary_keys: Detect primary keys detect_foreign_keys: Detect foreign keys include_key_confidence: Include key detection confidence include_row_count: Include row counts include_column_count: Include column counts include_distinct_count: Include distinct value counts include_null_count: Include NULL value counts include_cardinality: Include uniqueness ratio include_nullability: Include NULL ratio include_min_max: Include min/max values include_average: Include average values include_median: Include median (expensive) include_stddev: Include standard deviation (expensive) include_avg_length: Include average text length include_common_values: Include most frequent values common_values_limit: How many common values to show common_values_threshold: Only for columns with <= N distinct values include_date_range: Include min/max dates include_date_range_days: Include date range in days include_not_null_constraint: Include NOT NULL constraints include_default_values: Include default values include_indexes: Include index information max_rows_for_expensive_stats: Limit for median/stddev calculation max_string_display_length: String truncation length max_binary_display_bytes: Show binary data info up to N bytes include_db_metadata: Include database metadata include_table_metadata: Include table metadata include_extraction_timestamp: Include timestamp skip_empty_tables: Skip empty tables include_table_relationships: Include relationship summary include_schema_summary: Include overall schema statistics Returns: Dictionary containing database summary """ # Initialize key finder key_finder = None # Use SQLiteKeyFinder for enhanced PK/FK detection if requested if detect_primary_keys or detect_foreign_keys: key_finder = SQLiteKeyFinder(db_path) key_analysis = key_finder.analyze() else: key_analysis = {'tables': [], 'columns': {}, 'primary_keys': {}, 'foreign_keys': {}} # Still need to get basic table info conn_temp = sqlite3.connect(db_path) cursor_temp = conn_temp.cursor() cursor_temp.execute("SELECT name FROM sqlite_master WHERE type='table';") key_analysis['tables'] = [row[0] for row in cursor_temp.fetchall() if not row[0].startswith('sqlite_')] conn_temp.close() # Connect to database for additional info extraction conn = sqlite3.connect(db_path) conn.row_factory = sqlite3.Row cursor = conn.cursor() # Get database size if metadata is requested db_size = None if include_db_metadata: db_size = os.path.getsize(db_path) if os.path.exists(db_path) else None tables = key_analysis['tables'] # Create summary structure db_summary = {} # Add metadata if requested if include_db_metadata: db_summary["metadata"] = { "db_file": db_path, "db_size_bytes": db_size, "table_count": len(tables) } if include_extraction_timestamp: db_summary["metadata"]["extracted_at"] = datetime.datetime.now().isoformat() # Initialize schema summary if requested schema_summary = {} if include_schema_summary: schema_summary = { "total_tables": len(tables), "total_columns": 0, "tables_with_primary_keys": 0, "tables_with_foreign_keys": 0, "total_relationships": 0 } db_summary["tables"] = {} # Extract information for each table for table in tables: # Skip empty tables if requested quoted_table = quote_table_name(table) if skip_empty_tables: cursor.execute(f"SELECT COUNT(*) FROM {quoted_table};") if cursor.fetchone()[0] == 0: continue # Table schema cursor.execute(f"PRAGMA table_info({quoted_table});") columns_info = cursor.fetchall() # Table structure table_info = {"name": table} # Add table metadata if requested if include_table_metadata: table_info["column_count"] = 0 # We'll count only non-skipped columns # Get row count if requested if include_row_count: try: cursor.execute(f"SELECT COUNT(*) FROM {quoted_table};") table_info["row_count"] = cursor.fetchone()[0] except sqlite3.Error as e: print(f"Error getting row count for table {table}: {e}") table_info["row_count"] = -1 # Initialize columns list table_info["columns"] = [] # Initialize samples if requested if include_samples: table_info["samples"] = {} # Get column information - with enhanced PK detection if detect_primary_keys: pk_info = key_analysis['primary_keys'].get(table, {}) primary_keys = pk_info.get('columns', []) pk_origin = pk_info.get('origin', 'potential') else: primary_keys = [] pk_origin = 'unknown' # Track foreign key columns fk_columns = set() # Process columns, skipping those that should be excluded for col in columns_info: column_name = col["name"] # Skip columns that should be excluded based on DB_TABLE_SPECIFIC_SKIP_COLUMNS if key_finder and key_finder.should_skip_column(table, column_name): continue # Count this column for metadata if include_table_metadata: table_info["column_count"] += 1 # Update schema summary if include_schema_summary: schema_summary["total_columns"] += 1 # Basic column information column = {} if include_column_names: column["name"] = column_name if include_data_types: column["type"] = col["type"] # Primary key information if detect_primary_keys: is_primary_key = column_name in primary_keys column["is_primary_key"] = is_primary_key if is_primary_key and include_key_confidence: column["pk_origin"] = pk_origin # NOT NULL constraint if include_not_null_constraint: column["not_null"] = bool(col["notnull"] == 1) # Default values if include_default_values: column["default"] = col["dflt_value"] # Extract column statistics try: # Handle reserved keyword column names quoted_col_name = quote_column_name(column_name) # Get distinct value count if include_distinct_count: cursor.execute(f"SELECT COUNT(DISTINCT {quoted_col_name}) FROM {quoted_table};") column["distinct_count"] = cursor.fetchone()[0] # Get null count if include_null_count: cursor.execute(f"SELECT COUNT(*) FROM {quoted_table} WHERE {quoted_col_name} IS NULL;") column["null_count"] = cursor.fetchone()[0] # Calculate derived statistics if include_row_count and table_info.get("row_count", 0) > 0: if include_nullability and include_null_count: column["nullability"] = round(column["null_count"] / table_info["row_count"], 4) if include_cardinality and include_distinct_count: column["cardinality"] = round(column["distinct_count"] / table_info["row_count"], 4) # For numeric columns, get min, max, avg if any(num_type in col["type"].lower() for num_type in ["int", "real", "float", "double", "number", "decimal"]): try: if include_min_max or include_average: cursor.execute(f"SELECT MIN({quoted_col_name}), MAX({quoted_col_name}), AVG({quoted_col_name}) FROM {quoted_table} WHERE {quoted_col_name} IS NOT NULL;") min_val, max_val, avg_val = cursor.fetchone() if include_min_max: column["min"] = min_val column["max"] = max_val if include_average: column["avg"] = round(avg_val, 4) if avg_val is not None else None # For small to medium tables, get median and stddev if (include_median or include_stddev) and table_info.get("row_count", 0) <= max_rows_for_expensive_stats: cursor.execute(f"SELECT {quoted_col_name} FROM {quoted_table} WHERE {quoted_col_name} IS NOT NULL;") values = [row[0] for row in cursor.fetchall() if row[0] is not None] if values: try: if include_median: column["median"] = round(statistics.median(values), 4) if include_stddev and len(values) > 1: column["stddev"] = round(statistics.stdev(values), 4) except: pass # Skip if not truly numeric except: pass # Skip if not truly numeric # For text columns, get average length if any(text_type in col["type"].lower() for text_type in ["text", "char", "varchar", "string"]): try: if include_avg_length: cursor.execute(f"SELECT AVG(LENGTH({quoted_col_name})) FROM {quoted_table} WHERE {quoted_col_name} IS NOT NULL;") avg_length = cursor.fetchone()[0] column["avg_length"] = round(avg_length, 2) if avg_length is not None else None # Get most common values for columns with reasonable cardinality if (include_common_values and include_distinct_count and column.get("distinct_count", float('inf')) <= common_values_threshold): cursor.execute(f"SELECT {quoted_col_name}, COUNT(*) as cnt FROM {quoted_table} WHERE {quoted_col_name} IS NOT NULL GROUP BY {quoted_col_name} ORDER BY cnt DESC LIMIT {common_values_limit};") common_values = [] for row in cursor.fetchall(): value, count = row if isinstance(value, str) and len(value) > max_string_display_length: value = value[:max_string_display_length-3] + "..." common_values.append({"value": value, "count": count}) column["common_values"] = common_values except: pass # For date/time columns, get min and max dates if any(date_type in col["type"].lower() for date_type in ["date", "time", "timestamp", "datetime"]): try: if include_date_range: cursor.execute(f"SELECT MIN({quoted_col_name}), MAX({quoted_col_name}) FROM {quoted_table} WHERE {quoted_col_name} IS NOT NULL;") min_date, max_date = cursor.fetchone() column["min_date"] = min_date column["max_date"] = max_date # Try to calculate date range in days if include_date_range_days: try: cursor.execute(f"SELECT julianday(MAX({quoted_col_name})) - julianday(MIN({quoted_col_name})) FROM {quoted_table} WHERE {quoted_col_name} IS NOT NULL;") date_range = cursor.fetchone()[0] if date_range is not None: column["date_range_days"] = round(date_range, 2) except: pass except: pass except sqlite3.Error as e: print(f"Error getting statistics for column {table}.{column_name}: {e}") # Add column to table table_info["columns"].append(column) # Prepare for sample data if include_samples: table_info["samples"][column_name] = [] # Get foreign keys using enhanced detection with categorization if detect_foreign_keys: table_info["foreign_keys"] = [] fk_list = key_analysis['foreign_keys'].get(table, []) for fk in fk_list: # Handle both regular and composite foreign keys if 'type' in fk and fk['type'] == 'composite': # Check if any columns in this composite key should be skipped skip_columns = [col for col in fk['from_columns'] if key_finder and key_finder.should_skip_column(table, col)] skip_ref_columns = [col for col in fk['to_columns'] if key_finder and key_finder.should_skip_column(fk['to_table'], col)] # Skip this foreign key if it involves columns that should be excluded if skip_columns or skip_ref_columns: continue # Composite foreign key foreign_key = { "columns": fk['from_columns'], "references": { "table": fk['to_table'], "columns": fk['to_columns'] }, "type": "composite" } if include_key_confidence: foreign_key["fk_origin"] = fk.get('origin', 'potential') foreign_key["confidence"] = fk.get('confidence', 'low') # Update the column records to mark them as foreign keys for col_name in fk['from_columns']: for column in table_info["columns"]: if column.get("name") == col_name: column["is_foreign_key"] = True if include_key_confidence: column["fk_origin"] = fk.get('origin', 'potential') column["references_table"] = fk['to_table'] column["references_column"] = fk['to_columns'][fk['from_columns'].index(col_name)] column["composite_fk"] = True fk_columns.add(col_name) else: # Skip this foreign key if it involves columns that should be excluded if key_finder and (key_finder.should_skip_column(table, fk['from']) or key_finder.should_skip_column(fk['to_table'], fk['to_column'])): continue # Regular foreign key foreign_key = { "column": fk['from'], "references": { "table": fk['to_table'], "column": fk['to_column'] } } if include_key_confidence: foreign_key["fk_origin"] = fk.get('origin', 'potential') foreign_key["confidence"] = fk.get('confidence', 'medium') # Update the column record to mark it as a foreign key for column in table_info["columns"]: if column.get("name") == fk['from']: column["is_foreign_key"] = True if include_key_confidence: column["fk_origin"] = fk.get('origin', 'potential') column["references_table"] = fk['to_table'] column["references_column"] = fk['to_column'] fk_columns.add(fk['from']) # Add to table foreign keys table_info["foreign_keys"].append(foreign_key) # Get indexes if requested if include_indexes: table_info["indexes"] = [] cursor.execute(f"PRAGMA index_list({quoted_table});") idx_list = cursor.fetchall() for idx in idx_list: # idx fields: (seq, name, unique, origin, partial) index_name = idx["name"] index_unique = (idx["unique"] == 1) # Look up the columns used by this index cursor.execute(f"PRAGMA index_info({quote_column_name(index_name)});") idx_cols_info = cursor.fetchall() col_names = [ic["name"] for ic in idx_cols_info] # Filter out columns that should be excluded if key_finder: col_names = [col for col in col_names if not key_finder.should_skip_column(table, col)] # Skip empty indexes (after filtering) if not col_names: continue index = { "name": index_name, "unique": index_unique, "columns": col_names } table_info["indexes"].append(index) # Update column to indicate indexing for col_name in col_names: for column in table_info["columns"]: if column.get("name") == col_name: column["indexed"] = True if "indexes" not in column: column["indexes"] = [] column["indexes"].append(index_name) # Get sample data if requested if include_samples: try: # Debug info cursor.execute(f"SELECT COUNT(*) FROM {quoted_table};") row_count = cursor.fetchone()[0] print(f"Table {table} has {row_count} rows") # Get column names for SELECT, excluding those that should be skipped valid_column_names = [] for col in columns_info: col_name = col["name"] if key_finder: if not key_finder.should_skip_column(table, col_name): valid_column_names.append(col_name) else: valid_column_names.append(col_name) print(f"Valid columns for {table}: {valid_column_names}") if valid_column_names and row_count > 0: # Convert column names to quoted form for SQL quoted_col_names = [quote_column_name(col) for col in valid_column_names] select_cols = ', '.join(quoted_col_names) # Try different ordering strategies sample_rows = [] try: # First try ordering by rowid select_sql = f"SELECT {select_cols} FROM {quoted_table} ORDER BY rowid ASC LIMIT ?" cursor.execute(select_sql, (sample_limit,)) sample_rows = cursor.fetchall() print(f"Retrieved {len(sample_rows)} rows using rowid ordering") except sqlite3.OperationalError as e1: print(f"Rowid ordering failed: {e1}") try: # Fallback: order by first column first_col_quoted = quote_column_name(valid_column_names[0]) select_sql = f"SELECT {select_cols} FROM {quoted_table} ORDER BY {first_col_quoted} ASC LIMIT ?" cursor.execute(select_sql, (sample_limit,)) sample_rows = cursor.fetchall() print(f"Retrieved {len(sample_rows)} rows using first column ordering") except sqlite3.OperationalError as e2: print(f"First column ordering failed: {e2}") try: # Last resort: no ordering select_sql = f"SELECT {select_cols} FROM {quoted_table} LIMIT ?" cursor.execute(select_sql, (sample_limit,)) sample_rows = cursor.fetchall() print(f"Retrieved {len(sample_rows)} rows with no ordering") except sqlite3.OperationalError as e3: print(f"All query methods failed: {e3}") sample_rows = [] # Process each sample row if sample_rows: print(f"Processing {len(sample_rows)} sample rows for table {table}") if hasattr(sample_rows[0], 'keys'): print(f"Sample row keys: {list(sample_rows[0].keys())}") for row_idx, row in enumerate(sample_rows): # Convert sqlite3.Row to dict for easier access if hasattr(row, 'keys'): row_dict = dict(row) else: # Fallback: assume it's a tuple and map to column names row_dict = dict(zip(valid_column_names, row)) for col_name in valid_column_names: # Skip columns not in our table_info (already filtered) if col_name not in table_info["samples"]: continue # Get value - try different access methods value = None try: # Method 1: Direct access by column name value = row_dict.get(col_name) if value is None: # Method 2: Try case-insensitive lookup for key in row_dict.keys(): if key.lower() == col_name.lower(): value = row_dict[key] break except (KeyError, IndexError, AttributeError) as e: print(f"Error accessing column {col_name} in row {row_idx}: {e}") print(f"Available keys: {list(row_dict.keys()) if isinstance(row_dict, dict) else 'Not a dict'}") continue # Handle special data types for JSON serialization if isinstance(value, bytes): if len(value) <= max_binary_display_bytes: value = f"" else: value = f"" elif isinstance(value, str) and len(value) > max_string_display_length: value = value[:max_string_display_length-3] + "..." # Truncate long strings table_info["samples"][col_name].append(value) else: print(f"No sample rows retrieved for table {table}") # Debug: Print final sample counts for col_name, samples in table_info["samples"].items(): print(f"Column {col_name}: {len(samples)} samples") except sqlite3.Error as e: print(f"SQLite error getting sample data for table {table}: {e}") print(f"Query attempted: {select_sql if 'select_sql' in locals() else 'No query built'}") # Try to get actual column names from database for debugging try: cursor.execute(f"PRAGMA table_info({quoted_table});") actual_columns = [col[1] for col in cursor.fetchall()] print(f"Actual columns in {table}: {actual_columns}") print(f"Valid columns we tried to use: {valid_column_names if 'valid_column_names' in locals() else 'None'}") except Exception as debug_e: print(f"Could not get debug info: {debug_e}") except Exception as e: print(f"Unexpected error getting sample data for table {table}: {e}") import traceback traceback.print_exc() # Update schema summary if include_schema_summary: if detect_primary_keys and table in key_analysis['primary_keys']: schema_summary["tables_with_primary_keys"] += 1 if detect_foreign_keys and table in key_analysis['foreign_keys']: schema_summary["tables_with_foreign_keys"] += 1 schema_summary["total_relationships"] += len(key_analysis['foreign_keys'][table]) # Add table to db summary db_summary["tables"][table] = table_info # Add schema summary if requested if include_schema_summary: db_summary["schema_summary"] = schema_summary # Add table relationships summary if requested if include_table_relationships and detect_foreign_keys: relationships = [] for table, fk_list in key_analysis['foreign_keys'].items(): for fk in fk_list: if 'type' in fk and fk['type'] == 'composite': relationships.append({ "from_table": table, "from_columns": fk['from_columns'], "to_table": fk['to_table'], "to_columns": fk['to_columns'], "type": "composite" }) else: relationships.append({ "from_table": table, "from_column": fk['from'], "to_table": fk['to_table'], "to_column": fk['to_column'], "type": "simple" }) db_summary["relationships"] = relationships # Close database connection conn.close() if key_finder: key_finder.close() return db_summary def save_db_summary(db_summary, output_path, indent=2): """ Save database summary to JSON file """ def json_serialize(obj): if isinstance(obj, np.integer): return int(obj) if isinstance(obj, np.floating): return float(obj) if isinstance(obj, np.ndarray): return obj.tolist() if isinstance(obj, datetime.datetime): return obj.isoformat() return str(obj) # Create output directory if it doesn't exist os.makedirs(os.path.dirname(output_path), exist_ok=True) with open(output_path, 'w', encoding='utf-8') as f: json.dump(db_summary, f, default=json_serialize, indent=indent) print(f"Database summary saved to {output_path}") def main(): """ Main function to run the schema analysis """ import argparse from pathlib import Path parser = argparse.ArgumentParser(description='Analyze a SQLite database and generate schema summary') parser.add_argument('--db-path', required=True, help='Path to the SQLite database file (.sqlite)') parser.add_argument('--output', help='Directory to save the JSON output', default=str(Path(__file__).parent.parent.parent / 'data' / 'text2sql')) args = parser.parse_args() database_path = args.db_path if not os.path.exists(database_path): print(f"Error: Database file not found: {database_path}") return None print(f"Extracting schema summary from {database_path}...") try: db_summary = extract_db_summary_for_schema(database_path) except Exception as e: print(f"Error extracting database summary: {e}") return None # Save the summary db_basename = os.path.splitext(os.path.basename(database_path))[0] filename = f'{db_basename}_db_summary.json' output_path = os.path.join(args.output, filename) try: save_db_summary(db_summary, output_path) except Exception as e: print(f"Error saving database summary: {e}") return None # Print summary statistics print(f"\nSummary Statistics:") if 'metadata' in db_summary: metadata = db_summary['metadata'] db_size = metadata.get('db_size_bytes', 'unknown') if isinstance(db_size, int): db_size_mb = round(db_size / (1024 * 1024), 2) print(f" Database size: {db_size:,} bytes ({db_size_mb} MB)") else: print(f" Database size: {db_size}") print(f" Total tables: {metadata.get('table_count', 'unknown')}") if 'schema_summary' in db_summary: schema = db_summary['schema_summary'] print(f" Total columns: {schema.get('total_columns', 'unknown')}") print(f" Tables with PKs: {schema.get('tables_with_primary_keys', 'unknown')}") print(f" Tables with FKs: {schema.get('tables_with_foreign_keys', 'unknown')}") print(f" Total relationships: {schema.get('total_relationships', 'unknown')}") total_tables = len(db_summary.get('tables', {})) print(f" Processed tables: {total_tables}") # Calculate output file size try: output_size = os.path.getsize(output_path) output_size_kb = round(output_size / 1024, 2) print(f" Output file size: {output_size:,} bytes ({output_size_kb} KB)") except: pass print(f"\nDatabase summary extraction completed and saved to {output_path}") return db_summary if __name__ == "__main__": main()