sibi-flux 2025.12.0__py3-none-any.whl

sibi_flux/datacube/generator.py

@@ -0,0 +1,677 @@
+"""
+DataCube Generator module.
+
+Handles introspection of database schemas and generation of Python DataCube classes.
+"""
+
+import sqlalchemy as sa
+from pathlib import Path
+from typing import Optional, Set, Any, Callable, Tuple
+from collections.abc import Mapping, MutableMapping, Iterable, Sequence
+from sqlalchemy import inspect
+import re
+import importlib
+import sys
+import os
+
+
+def get_sa_type_classes(
+    type_names: Iterable[str], logger: Optional[Any] = None
+) -> Tuple[Any, ...]:
+    """Converts a list of strings to actual SQLAlchemy type classes."""
+    classes = []
+    for name in type_names:
+        type_cls = getattr(sa, name, None)
+        if type_cls:
+            classes.append(type_cls)
+        elif logger:
+            logger.warning(f"SQLAlchemy type '{name}' not found.")
+    return tuple(classes)
+
+
+def is_secure_path(target_file: str, allowed_dirs: Iterable[str]) -> bool:
+    """Verifies path is within sandbox using resolution to prevent traversal."""
+    try:
+        target_path = Path(target_file).resolve()
+        for allowed in allowed_dirs:
+            if target_path.is_relative_to(Path(allowed).resolve()):
+                return True
+    except (ValueError, RuntimeError, OSError):
+        return False
+    return False
+
+
+def generate_class_code(
+    inspector: Optional[sa.engine.Inspector],
+    table_name: str,
+    config_obj: str,
+    base_class: str,
+    mappings: Mapping[str, tuple],
+    class_name: Optional[str] = None,
+    class_suffix: str = "Dc",
+    backend: str = "sqlalchemy",
+    field_map_var: Optional[str] = None,
+    legacy_filters: bool = False,
+    field_map_dict: Optional[Mapping[str, str]] = None,
+    sticky_filters: Optional[Mapping[str, Any]] = None,
+) -> str:
+    """Introspects DB table and generates Python class string. If inspector is None, generates shell."""
+
+    # Initialize metadata dict based on the keys in our YAML mapping
+    meta: dict[str, list[str]] = {label: [] for label in mappings.keys()}
+
+    if inspector:
+        try:
+            columns = inspector.get_columns(table_name)
+            for col in columns:
+                col_name = col["name"]
+                # Apply alias if map is provided
+                if field_map_dict:
+                    col_name = field_map_dict.get(col_name, col_name)
+
+                for label, sa_classes in mappings.items():
+                    if isinstance(col["type"], sa_classes):
+                        meta[label].append(col_name)
+        except (sa.exc.SQLAlchemyError, KeyError):
+            # If introspection fails for specific table (e.g. missing), just skip columns
+            pass
+
+    if not class_name:
+        class_name = (
+            "".join(w.capitalize() for w in table_name.split("_")) + class_suffix
+        )
+
+    # Generate attribute lines for every label defined in the mapping
+    metadata_lines = []
+    for label, cols in meta.items():
+        metadata_lines.append(f"    {label}: ClassVar[List[str]] = {repr(cols)}")
+
+    # Configuration block extensions
+    extra_config = []
+    if field_map_var:
+        extra_config.append(f"    field_map: ClassVar[Dict] = {field_map_var}")
+    if legacy_filters:
+        extra_config.append(f"    legacy_filters: ClassVar[bool] = {legacy_filters}")
+    if sticky_filters:
+        extra_config.append(
+            f"    sticky_filters: ClassVar[Dict] = {repr(sticky_filters)}"
+        )
+
+    attributes_str = "\n".join(metadata_lines)
+    extra_config_str = "\n".join(extra_config)
+    if extra_config_str:
+        extra_config_str = "\n" + extra_config_str
+
+    return f"""
+class {class_name}({base_class}):
+    df: Optional[dd.DataFrame] = None
+    
+    # Config
+    backend: ClassVar[str] = "{backend}"
+    connection_url: ClassVar[str] = {config_obj}.get("db_url")
+    table: ClassVar[str] = "{table_name}"{extra_config_str}
+    
+    # Transformations & Metadata
+{attributes_str}
+"""
+
+
+def check_drift(
+    dc_class: Any,
+    db_column_names: Iterable[str],
+    attribute_names: Iterable[str],
+    field_map: Optional[Mapping[str, str]] = None,
+) -> list[str]:
+    """
+    Checks for drift between the datacube class metadata and database columns.
+    Returns a list of issue strings.
+    """
+    code_cols: Set[str] = set()
+    for attr in attribute_names:
+        # Get the list from the class (default to empty list)
+        cols = getattr(dc_class, attr, [])
+        code_cols.update(cols)
+
+    # Apply alias mapping to DB columns if present
+    # Map: DB_NAME -> CODE_NAME
+    if field_map:
+        mapped_db_cols = set()
+        for c in db_column_names:
+            mapped_db_cols.add(field_map.get(c, c))  # Use alias if exists, else raw
+        final_db_cols = mapped_db_cols
+    else:
+        final_db_cols = set(db_column_names)
+
+    missing_in_code = final_db_cols - code_cols
+    extra_in_code = code_cols - final_db_cols
+
+    issues = []
+    if missing_in_code:
+        issues.append(f"DB has new cols: {missing_in_code}")
+    if extra_in_code:
+        issues.append(f"Code has stale cols: {extra_in_code}")
+
+    return issues
+
+
+def resolve_db_url(config_obj_name: str, imports: Iterable[str]) -> Optional[str]:
+    """
+    Attempts to resolve a database URL by dynamically importing the config object
+    described in the import statements.
+    """
+    # Ensure current working directory is in path to allow local imports
+    if os.getcwd() not in sys.path:
+        sys.path.insert(0, os.getcwd())
+
+    for imp in imports:
+        # Simple parsing for "from module import var" or "from module import var1, var2"
+        if imp.startswith("from ") and " import " in imp:
+            parts = imp.split(" import ")
+            module_name = parts[0][5:].strip()
+            imported_vars = [v.strip() for v in parts[1].split(",")]
+
+            if config_obj_name in imported_vars:
+                try:
+                    mod = importlib.import_module(module_name)
+                    conf = getattr(mod, config_obj_name)
+                    # Support dict-like (.get) or object attribute access
+                    if hasattr(conf, "get"):
+                        return conf.get("db_url")
+                    elif hasattr(conf, "db_url"):
+                        return conf.db_url
+                except (ImportError, AttributeError):
+                    pass
+    return None
+
+
+def load_environment(
+    env_file: Optional[Path] = None, logger: Optional[Any] = None
+) -> None:
+    """
+    Loads environment variables from the specified file or .env.linux by default.
+    """
+    target = env_file if env_file else Path(os.getcwd()) / ".env"
+    if target.exists():
+        if logger:
+            logger(f"Loading environment from {target.name}...")
+        with open(target, "r") as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith("#"):
+                    continue
+                if "=" in line:
+                    key, value = line.split("=", 1)
+                    os.environ[key.strip()] = value.strip().strip("'").strip('"')
+    elif env_file and logger:
+        logger(f"Warning: Environment file '{env_file}' not found.")
+
+
+def filter_global_imports(
+    global_imports: Iterable[str],
+    used_configs: Set[str],
+    ignored_prefixes: Optional[Iterable[str]] = None,
+) -> list[str]:
+    """
+    Filters global import lines, keeping only those that are used or generic.
+    Renames imports if they contain multiple items but only some are used.
+    """
+    filtered = []
+    ignored = ignored_prefixes or []
+    for line in global_imports:
+        if " import " in line and line.strip().startswith("from "):
+            parts = line.split(" import ")
+            prefix = parts[0]
+            imported_vars = [v.strip() for v in parts[1].split(",")]
+
+            # Check for intersection with used configs
+            common = [v for v in imported_vars if v in used_configs]
+
+            if common:
+                # If we use some of the imported vars, keep only those
+                filtered.append(f"{prefix} import {', '.join(common)}")
+            else:
+                # If we use NONE, check if it looks like a config import we should drop
+                if any(p in prefix for p in ignored):
+                    continue
+                # Otherwise keep generic imports (e.g. typing, os, etc if they appeared here)
+                filtered.append(line)
+        else:
+            filtered.append(line)
+    return filtered
+
+
+class DatacubeRegistry:
+    """
+    Encapsulates parsing logic for the datacube_registry.yaml.
+    """
+
+    def __init__(self, config_dict: MutableMapping[str, Any]):
+        self.config = config_dict
+        self.tables = self.config.setdefault("tables", {})
+        self.global_imports = self.config.get("global_imports", [])
+
+        # Resolve Mappings
+        raw_mappings = self.config.get("type_mappings", {})
+        self.processed_mappings = {
+            label: get_sa_type_classes(types) for label, types in raw_mappings.items()
+        }
+
+        # Defaults
+        # Defaults
+        self.valid_paths = self.config.get("valid_libpaths")
+        if not self.valid_paths:
+            # Make it empty/None? Or raise? Or let it be provided by caller?
+            # For now, let's just make it empty list so it fails secure check if not provided
+            self.valid_paths = []
+
+        self.valid_fieldmap_paths = self.config.get("valid_fieldmap_paths") or []
+        self.default_connection_obj = self.config.get(
+            "default_connection_obj", self.config.get("default_config_obj")
+        )
+        self.default_base_class = self.config.get("default_base_class")
+        self.default_base_import = self.config.get("default_base_import")
+        self.default_backend = self.config.get("default_backend", "sqlalchemy")
+        self.class_suffix = self.config.get("class_suffix", "Dc")
+
+    def get_table_details(self, table_name: str) -> dict[str, Any]:
+        return self.tables.get(table_name, {})
+
+    def group_tables_by_file(self) -> dict[str, list[tuple]]:
+        """
+        Groups tables by their target file path.
+        Returns: Dict[path_str, List[(table_name, conf_obj, base_cls, base_imp, cls_name)]]
+        """
+        file_groups: dict[str, list[tuple]] = {}
+        for table_name, details in self.tables.items():
+            target = details.get("save_to_path", details.get("path"))
+            if not target:
+                continue
+
+            conf_obj = details.get(
+                "connection_obj", details.get("config_obj", self.default_connection_obj)
+            )
+            base_cls = details.get("base_class", self.default_base_class)
+            base_imp = details.get("import_from", self.default_base_import)
+            cls_name = details.get("class_name")
+
+            file_groups.setdefault(target, []).append(
+                (table_name, conf_obj, base_cls, base_imp, cls_name)
+            )
+        return file_groups
+
+    def merge_discovered(
+        self, discovered_tables: Mapping[str, Mapping[str, Any]]
+    ) -> None:
+        """
+        Merges discovered tables into the existing configuration.
+        """
+        for table, new_details in discovered_tables.items():
+            if table in self.tables:
+                existing = self.tables[table]
+                for k, v in new_details.items():
+                    if v is not None:
+                        existing[k] = v
+                    elif k not in existing:
+                        existing[k] = v
+
+                if "class_name" not in existing:
+                    existing["class_name"] = new_details.get("class_name")
+
+                self.tables[table] = existing
+            else:
+                self.tables[table] = new_details
+
+    def prune_tables(self, keep_tables: Iterable[str]) -> list[str]:
+        """
+        Removes tables from the registry that are not in the keep_tables list.
+        """
+        to_remove = [t for t in self.tables if t not in keep_tables]
+        for t in to_remove:
+            del self.tables[t]
+        return to_remove
+
+    def to_class_name(self, table_name: str) -> str:
+        """
+        Converts snake_case table name to CamelCase + Suffix.
+        """
+        return (
+            "".join(w.capitalize() for w in table_name.split("_")) + self.class_suffix
+        )
+
+    def normalize_aliases(self, ignored_prefixes: Optional[list[str]] = None) -> None:
+        pass
+
+    def discover(self, allowed_paths: Optional[list[str]] = None) -> list[str]:
+        return []
+
+
+def perform_discovery(
+    all_table_names: Iterable[str],
+    match_rule_callback: Callable[[str], Optional[Mapping[str, Any]]],
+    registry: DatacubeRegistry,
+    whitelist: Optional[Mapping[str, Mapping[str, Any]]] = None,
+    template: Optional[Mapping[str, Any]] = None,
+    field_map_template: Optional[str] = None,
+    default_connection: str = "default_db_conf",
+) -> dict[str, dict[str, Any]]:
+    """
+    Executes the discovery logic:
+    1. Filter by whitelist (if provided).
+    2. Match tables against rules.
+    3. Construct entries using (Template + Rule + Whitelist Override).
+    """
+    discovered = {}
+
+    candidates = all_table_names
+    if whitelist is not None:
+        candidates = [t for t in all_table_names if t in whitelist]
+
+    for table in candidates:
+        rule = match_rule_callback(table)
+        if not rule:
+            continue
+
+        path = rule["path"]
+        class_name = registry.to_class_name(table)
+
+        # entry = (template or {}).copy()
+        # Refactoring Note: dict() creates a copy and supports Mapping
+        entry = dict(template or {})
+
+        if rule.get("domain") and field_map_template:
+            # e.g. "solutions.conf.transforms.fields.{domain}.{table}.field_map"
+            entry["field_map"] = field_map_template.format(
+                domain=rule["domain"], table=table
+            )
+
+        entry.update(
+            {
+                "class_name": class_name,
+                "connection_obj": rule.get("connection_obj") or default_connection,
+                "save_to_path": path,
+            }
+        )
+
+        if whitelist and table in whitelist:
+            overrides = whitelist[table]
+            if overrides:
+                entry.update(overrides)
+
+        discovered[table] = entry
+
+    return discovered
+
+
+def generate_field_map_files(
+    discovered_entries: Mapping[str, Mapping[str, Any]],
+    inspector: Any,
+    root_path: Path,
+    force: bool = False,
+    logger: Optional[Any] = None,
+    allowed_paths: Optional[Iterable[str]] = None,
+) -> None:
+    """
+    Generates or updates Python field map files for discovered tables.
+    If file exists and not force: performs smart sync (comment out stale, add new).
+    If force: overwrites completely.
+    """
+    key_pattern = re.compile(r'^\s*(?P<comment>#\s*)?[\'"](?P<key>\w+)[\'"]\s*:')
+
+    for table, entry in discovered_entries.items():
+        field_map = entry.get("field_map")
+        if not field_map:
+            continue
+
+        rule = entry  # entry has domain if matched via discovery
+        if "domain" not in rule and "field_map" in rule:
+            # Try to infer target dir from field_map string??
+            # Convention: solutions.conf.transforms.fields.<domain>.<table_name>.field_map
+            parts = field_map.split(".")
+            if len(parts) > 4:
+                # domain is parts[-3] if standard convention
+                pass
+
+        # We assume root_path is passed correctly (e.g. solutions/conf/transforms/fields)
+        # And we append domain/table.py
+        # But we really should trust the rule['domain'] if available
+        domain = rule.get("domain", "common")
+        target_dir = root_path / domain
+        target_file = target_dir / f"{table}.py"
+
+        # Security Check
+        if allowed_paths:
+            if not is_secure_path(str(target_file), allowed_paths):
+                if logger:
+                    logger(
+                        f"[red]Blocked Field Map Generation: {target_file} is outside allowed paths.[/red]"
+                    )
+                continue
+
+        target_dir.mkdir(parents=True, exist_ok=True)
+
+        if not (target_dir / "__init__.py").exists():
+            (target_dir / "__init__.py").touch()
+
+        try:
+            cols = inspector.get_columns(table)
+            db_col_names = {c["name"] for c in cols}
+            col_list = [c["name"] for c in cols]  # Preserve order
+
+            if target_file.exists() and not force:
+                # --- Smart Sync ---
+                with open(target_file, "r") as f:
+                    lines = f.readlines()
+
+                new_lines = []
+                existing_keys = set()
+
+                for line in lines:
+                    match = key_pattern.match(line)
+                    if match:
+                        key = match.group("key")
+                        existing_keys.add(key)
+                        is_commented = bool(match.group("comment"))
+
+                        if key not in db_col_names:
+                            if not is_commented:
+                                line = f"# {line}"
+                    new_lines.append(line)
+
+                missing_cols = db_col_names - existing_keys
+
+                insert_idx = len(new_lines)
+                for i in range(len(new_lines) - 1, -1, -1):
+                    if "}" in new_lines[i]:
+                        insert_idx = i
+                        break
+
+                to_insert = []
+                for col in sorted(missing_cols):
+                    to_insert.append(
+                        f'    "{col}": "{col}", # TODO: Translate to English\n'
+                    )
+
+                new_lines[insert_idx:insert_idx] = to_insert
+
+                with open(target_file, "w") as f:
+                    f.writelines(new_lines)
+
+                if logger:
+                    logger(
+                        f"[green]Synced {target_file} (Added {len(missing_cols)}, Checked Stale)[/green]"
+                    )
+            else:
+                # --- Fresh Generation ---
+                lines = ["field_map = {"]
+                for col in col_list:
+                    lines.append(f'    "{col}": "{col}", # TODO: Translate to English')
+                lines.append("}")
+
+                with open(target_file, "w") as f:
+                    f.write("\n".join(lines))
+
+                if logger:
+                    logger(f"[green]Generated {target_file}[/green]")
+
+        except (OSError, sa.exc.SQLAlchemyError) as e:
+            if logger:
+                logger(f"[red]Failed to generate/sync fields for {table}: {e}[/red]")
+
+
+def generate_datacube_module_code(
+    items: Iterable[tuple],
+    registry: DatacubeRegistry,
+    get_db_url_callback: Callable[[str], str],
+    logger: Optional[Any] = None,
+) -> Tuple[list[str], list[str]]:
+    """
+    Generates the code components (imports, class definitions) for a single DataCube module file.
+    Handles import aliasing to prevent variable shadowing.
+    """
+    # Refactoring Note: using concrete Dict in string literal is fine as it's code generation
+    imports: Set[str] = {
+        "from typing import Optional, List, ClassVar, Dict",
+        "import dask.dataframe as dd",
+    }
+    classes_code = []
+
+    for table_name, conf_obj, base_cls, base_imp, cls_name in items:
+        imports.add(base_imp)
+
+        details = registry.get_table_details(table_name)
+        field_map_str = details.get("field_map")
+        sticky_filters = details.get("sticky_filters")
+
+        field_map_var = None
+        legacy = False
+
+        field_map_dict = None
+        if field_map_str:
+            if "." in field_map_str:
+                mod, var = field_map_str.rsplit(".", 1)
+                # FIX: Create a unique alias for the import to avoid shadowing
+                alias = f"{var}_{table_name}"
+                imports.add(f"from {mod} import {var} as {alias}")
+                field_map_var = alias
+                legacy = True
+
+                # Try to load the actual dict for alias application during introspection
+                try:
+                    module = importlib.import_module(mod)
+                    field_map_dict = getattr(module, var)
+                    if not isinstance(field_map_dict, dict):
+                        if logger:
+                            logger(
+                                f"[yellow]Warning: {var} in {mod} is not a dict. Alias mapping skipped.[/yellow]"
+                            )
+                        field_map_dict = None
+                except (ImportError, AttributeError) as e:
+                    if logger:
+                        logger(
+                            f"[yellow]Warning: Could not load field_map {var} from {mod}: {e}. Alias mapping skipped.[/yellow]"
+                        )
+                    field_map_dict = None
+            else:
+                if logger:
+                    logger(
+                        f"[yellow]Warning: Invalid field_map format '{field_map_str}' for {table_name}. Expected module.path.variable[/yellow]"
+                    )
+
+        try:
+            db_url = get_db_url_callback(conf_obj)
+            engine = sa.create_engine(db_url)
+            inspector = inspect(engine)
+
+            classes_code.append(
+                generate_class_code(
+                    inspector,
+                    table_name,
+                    conf_obj,
+                    base_cls,
+                    registry.processed_mappings,
+                    class_name=cls_name,
+                    class_suffix=registry.class_suffix,
+                    backend=registry.default_backend,
+                    field_map_var=field_map_var,
+                    legacy_filters=legacy,
+                    field_map_dict=field_map_dict,
+                    sticky_filters=sticky_filters,
+                )
+            )
+        except sa.exc.SQLAlchemyError as e:
+            if logger:
+                logger(
+                    f"[yellow]Warning: Could not introspect {table_name} ({e}). Generating shell class.[/yellow]"
+                )
+            classes_code.append(
+                generate_class_code(
+                    None,
+                    table_name,
+                    conf_obj,
+                    base_cls,
+                    registry.processed_mappings,
+                    class_name=cls_name,
+                    class_suffix=registry.class_suffix,
+                    backend=registry.default_backend,
+                    field_map_var=field_map_var,
+                )
+            )
+
+    return list(imports), classes_code
+
+
+def dump_db_schema(
+    engine: Any,
+    db_name: str,
+    output_dir: Path,
+    whitelist: Optional[Mapping[str, Any]] = None,
+    logger: Optional[Any] = None,
+) -> None:
+    """
+    Introspects the database and dumps a comprehensive schema YAML file.
+    Structure:
+    <db_name>:
+       <table_name>:
+          columns: ...
+          indexes: ...
+          pk: ...
+          foreign_keys: ...
+    """
+    inspector = inspect(engine)
+    all_tables = inspector.get_table_names()
+
+    if whitelist:
+        tables_to_process = sorted([t for t in all_tables if t in whitelist])
+    else:
+        tables_to_process = sorted(all_tables)
+
+    schema_data = {}
+
+    for table in tables_to_process:
+        try:
+            # Columns
+            columns = []
+            for col in inspector.get_columns(table):
+                columns.append({"name": col["name"], "type": str(col["type"])})
+
+            schema_data[table] = {"columns": columns}
+        except sa.exc.SQLAlchemyError as e:
+            if logger:
+                logger(f"[red]Error dumping schema for table {table}: {e}[/red]")
+            schema_data[table] = {"error": str(e)}
+
+    # Top level key is db_name
+    final_output = {db_name: schema_data}
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    out_file = output_dir / f"schema_{db_name}.yaml"
+
+    import yaml
+
+    with open(out_file, "w") as f:
+        yaml.dump(final_output, f, sort_keys=False, default_flow_style=False)
+
+    if logger:
+        logger(
+            f"[green]Dumped schema to {out_file} ({len(tables_to_process)} tables)[/green]"
+        )