Shared utilities for DuckDB data extraction and analysis.
| 14 | |
| 15 | |
| 16 | class DuckDBExtractor: |
| 17 | """Shared utilities for DuckDB data extraction and analysis.""" |
| 18 | |
| 19 | def __init__(self, source_dir: str = None): |
| 20 | """ |
| 21 | Initialize the DuckDB extractor. |
| 22 | |
| 23 | Args: |
| 24 | source_dir: Path to directory containing DuckDB files. If None, uses default. |
| 25 | """ |
| 26 | if source_dir is None: |
| 27 | # Default to shared/databases/duckdb relative to this script |
| 28 | script_dir = os.path.dirname(os.path.abspath(__file__)) |
| 29 | project_root = os.path.dirname(script_dir) |
| 30 | source_dir = os.path.join(project_root, "shared", "databases", "duckdb") |
| 31 | |
| 32 | self.source_dir = Path(source_dir).resolve() |
| 33 | self.excluded_files: Set[str] = set() |
| 34 | |
| 35 | def add_exclusions(self, exclusion_list: List[str]): |
| 36 | """Add files to exclusion list.""" |
| 37 | self.excluded_files.update(exclusion_list) |
| 38 | |
| 39 | def discover_duckdb_files(self, patterns: List[str] = None) -> List[Path]: |
| 40 | """Discover DuckDB files in the configured source directory.""" |
| 41 | if patterns is None: |
| 42 | patterns = ["*.duckdb", "*.db", "*.ddb"] |
| 43 | |
| 44 | if not self.source_dir.exists(): |
| 45 | print(f"Error: Source directory does not exist: {self.source_dir}", file=sys.stderr) |
| 46 | return [] |
| 47 | |
| 48 | discovered_files = [] |
| 49 | for pattern in patterns: |
| 50 | discovered_files.extend(self.source_dir.glob(pattern)) |
| 51 | |
| 52 | # Remove duplicates and sort |
| 53 | return sorted(set(discovered_files)) |
| 54 | |
| 55 | def analyze_duckdb_schema(self, db_path: Path) -> Dict: |
| 56 | """Analyze schema and tables in a DuckDB file.""" |
| 57 | try: |
| 58 | with duckdb.connect(str(db_path), read_only=True) as conn: |
| 59 | # Get all schemas |
| 60 | schemas = conn.execute(""" |
| 61 | SELECT schema_name |
| 62 | FROM information_schema.schemata |
| 63 | WHERE schema_name NOT IN ('information_schema', 'pg_catalog', 'pg_toast') |
| 64 | ORDER BY schema_name |
| 65 | """).fetchall() |
| 66 | |
| 67 | schema_info = {} |
| 68 | for (schema_name,) in schemas: |
| 69 | # Get tables in this schema |
| 70 | tables = conn.execute( |
| 71 | """ |
| 72 | SELECT table_name, table_type |
| 73 | FROM information_schema.tables |