sidra-sql 1.3.0__py3-none-any.whl

sidra_sql/storage.py

@@ -0,0 +1,156 @@
+"""Storage repository for SIDRA data files.
+
+This module provides the `Storage` class, which centralizes all
+filesystem operations for SIDRA tables: constructing deterministic file
+paths from a `Parametro`, checking whether a file already exists, and
+reading/writing JSON data files.
+"""
+
+import logging
+from pathlib import Path
+
+import orjson
+from sidra_fetcher.agregados import Agregado
+from sidra_fetcher.reader import load_agregado, save_agregado
+from sidra_fetcher.sidra import Parametro
+
+from .config import Config
+
+logger = logging.getLogger(__name__)
+
+
+class Storage:
+    def __init__(self, data_dir: Path | str):
+        self.data_dir = Path(data_dir)
+
+    @classmethod
+    def default(cls, config: Config) -> "Storage":
+        """Create a Storage rooted at the data directory from config."""
+        data_dir = config.data_dir
+        data_dir.mkdir(exist_ok=True, parents=True)
+        return cls(data_dir)
+
+    @staticmethod
+    def build_data_filename(parameter: Parametro, modification: str) -> str:
+        """Build a deterministic filename for a SIDRA `Parametro`.
+
+        The filename encodes the table id, periods, format, territorial
+        levels, variables and classifications so that each unique request
+        maps to a unique file name. The returned filename ends with
+        ``@{modification}.json`` where ``modification`` is typically the
+        period modification timestamp used by the API.
+
+        Args:
+            parameter: A `Parametro` instance containing request
+                configuration (table, periods, territorios, variaveis,
+                classificacoes, formato).
+            modification: A string representing the modification
+                timestamp to append to the filename.
+
+        Returns:
+            A string suitable for use as a JSON filename.
+        """
+        sidra_table = parameter.agregado
+        periods = ",".join(parameter.periodos)
+        formato = parameter.formato.value
+        name = f"t-{sidra_table}_p-{periods}_f-{formato}"
+        for territorial_level in parameter.territorios:
+            territorial_codes = ",".join(
+                str(code) for code in parameter.territorios[territorial_level]
+            )
+            if territorial_codes == "":
+                territorial_codes = "all"
+            name += f"_n{territorial_level}-{territorial_codes}"
+        if parameter.variaveis:
+            variables = ",".join(str(var) for var in parameter.variaveis)
+            name += f"_v-{variables}"
+        for classification, categories in parameter.classificacoes.items():
+            str_categories = ",".join(categories)
+            name += f"_c{classification}-{str_categories}"
+        name += f"@{modification}"
+        name += ".json"
+        return name
+
+    def get_metadata_filepath(self, agregado: int | str) -> Path:
+        """Return the full path for a table's metadata JSON file."""
+        return self.data_dir / f"t-{agregado}" / "metadados.json"
+
+    def get_data_filepath(
+        self,
+        parameter: Parametro,
+        modification: str,
+    ) -> Path:
+        """Return the full path for the given parameter and modification."""
+        filename = self.build_data_filename(parameter, modification)
+        return self.data_dir / f"t-{parameter.agregado}" / filename
+
+    def exists(self, parameter: Parametro, modification: str) -> bool:
+        """Return True if the file for the given parameter already exists."""
+        return self.get_data_filepath(parameter, modification).exists()
+
+    def write_data(
+        self,
+        data: dict,
+        parameter: Parametro,
+        modification: str,
+    ) -> Path:
+        """Write *data* to disk as JSON and return the destination path."""
+        filepath = self.get_data_filepath(parameter, modification)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        logger.info("Writing file %s", filepath)
+        with filepath.open("w", encoding="utf-8") as f:
+            f.write(orjson.dumps(data, option=orjson.OPT_INDENT_2).decode())
+        return filepath
+
+    def read_data(self, filepath: Path) -> list[dict]:
+        """Read a JSON file previously written by `write`.
+
+        Args:
+            filepath: Path to the JSON file to read.
+
+        Returns:
+            A list of dicts containing the table data.
+        """
+        logger.info("Reading file %s", filepath)
+        with filepath.open("rb") as f:
+            data = orjson.loads(f.read())
+
+        if len(data) > 1:
+            rows = data[1:]
+            for row in rows:
+                for k, v in row.items():
+                    if v in ("...", "-"):
+                        row[k] = None
+            return rows
+        return []
+
+    def write_metadata(self, agregado: Agregado) -> Path:
+        """Write *agregado* metadata to disk as JSON and return the destination path."""
+        filepath = self.get_metadata_filepath(agregado.id)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        logger.info("Writing file %s", filepath)
+        save_agregado(agregado, filepath)
+        return filepath
+
+    def read_metadata(self, agregado: int | str) -> Agregado:
+        """Read metadata from disk as JSON and return the destination path."""
+        filepath = self.get_metadata_filepath(agregado)
+        agregado = load_agregado(filepath)
+        return agregado
+
+    def read_data_dir(self, dirpath: Path) -> list[dict]:
+        # Group files by base name (before the @modification suffix) and keep
+        # only the file with the latest modification per parameter combination.
+        latest: dict[str, tuple[Path, str]] = {}
+        for f in dirpath.glob("*.json"):
+            stem = f.stem
+            if "@" in stem:
+                base, mod = stem.rsplit("@", 1)
+            else:
+                base, mod = stem, ""
+            if base not in latest or mod > latest[base][1]:
+                latest[base] = (f, mod)
+        data = []
+        for f, _ in latest.values():
+            data.extend(self.read_data(f))
+        return data

sidra_sql/toml_runner.py

@@ -0,0 +1,374 @@
+"""Pipeline runner driven by a TOML configuration file.
+
+``TomlScript`` reads a TOML file that declares which SIDRA tables to fetch
+and drives the full ETL pipeline: metadata → download → database load.
+
+TOML schema
+-----------
+Each ``[[tabelas]]`` entry maps directly to a `sidra.Fetcher.download_table`
+call.  Two optional boolean flags extend the static format:
+
+``unnest_classifications = true``
+    Fetch the table's metadata at runtime and expand every classification /
+    category combination with `sidra.unnest_classificacoes`.  The
+    ``classifications`` key, if present, is ignored.
+
+``unnest_classifications = ["id1", "id2"]``
+    Same expansion, but only for the listed classification IDs (integers also
+    accepted).  The ``classifications`` key, if present, is merged as static
+    defaults — unnested entries take precedence on key conflicts.
+
+``split_variables = true``
+    Issue one request per variable listed in ``variables`` instead of a
+    single request with all variables.
+
+Example TOML
+~~~~~~~~~~~~
+::
+
+    [[tabelas]]
+    tabela_sidra = "5938"
+    variables = ["37", "498"]
+    territories = {6 = ["all"]}
+
+    [[tabelas]]
+    tabela_sidra = "1613"
+    variables = ["allxp"]
+    territories = {6 = []}
+    unnest_classifications = true
+
+    [[tabelas]]
+    tabela_sidra = "5938"
+    variables = ["allxp"]
+    territories = {6 = []}
+    classifications = {81 = ["allxt"]}
+    unnest_classifications = ["87"]
+
+    [[tabelas]]
+    tabela_sidra = "1002"
+    variables = ["109", "216", "214", "112"]
+    split_variables = true
+    territories = {6 = []}
+    classifications = {81 = ["allxt"]}
+"""
+
+import logging
+import tomllib
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+import sqlalchemy as sa
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TaskID,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+from rich.table import Table
+from rich.text import Text
+
+from . import database, models, sidra
+from .config import Config
+from .storage import Storage
+
+logger = logging.getLogger(__name__)
+
+
+class _MainOnlyTimeElapsedColumn(TimeElapsedColumn):
+    def render(self, task):
+        if not task.fields.get("main"):
+            return Text("")
+        return super().render(task)
+
+
+class _MainOnlyTimeRemainingColumn(TimeRemainingColumn):
+    def render(self, task):
+        if not task.fields.get("main"):
+            return Text("")
+        return super().render(task)
+
+
+def _make_progress(console: Console | None) -> Progress:
+    return Progress(
+        SpinnerColumn(finished_text="[green]✓[/green]"),
+        TextColumn("[progress.description]{task.description}", table_column=None),
+        BarColumn(bar_width=28),
+        TextColumn("[progress.percentage]{task.percentage:>3.0f}%", style="grey70"),
+        _MainOnlyTimeElapsedColumn(),
+        _MainOnlyTimeRemainingColumn(),
+        console=console,
+        transient=False,
+        disable=console is None,
+    )
+
+
+def _make_download_progress(console: Console | None) -> Progress:
+    return Progress(
+        SpinnerColumn(finished_text="[green]✓[/green]"),
+        TextColumn("[progress.description]{task.description}", table_column=None),
+        BarColumn(bar_width=28),
+        MofNCompleteColumn(),
+        _MainOnlyTimeElapsedColumn(),
+        _MainOnlyTimeRemainingColumn(),
+        console=console,
+        transient=False,
+        disable=console is None,
+    )
+
+
+class TomlScript:
+    """ETL pipeline runner that loads table definitions from a TOML file.
+
+    Drives the full pipeline:
+
+    1. Create ORM tables if they don't exist.
+    2. Fetch and save metadata (tabela_sidra, localidade).
+    3. Download all data files.
+    4. Load data rows into the dados table (also upserts dimensions).
+    """
+
+    def __init__(
+        self,
+        config: Config,
+        toml_path: Path,
+        max_workers: int = 4,
+        force_metadata: bool = False,
+        force_load: bool = False,
+        console: Console | None = None,
+    ):
+        self.config = config
+        self.toml_path = toml_path
+        self.force_metadata = force_metadata
+        self.force_load = force_load
+        self.console = console
+        self.storage = Storage.default(config)
+        self.fetcher = sidra.Fetcher(
+            config, max_workers=max_workers, storage=self.storage
+        )
+
+    def get_tabelas(self) -> Iterable[dict[str, Any]]:
+        """Read the TOML file and return an expanded list of table request dicts."""
+        with open(self.toml_path, "rb") as f:
+            data = tomllib.load(f)
+
+        result: list[dict[str, Any]] = []
+        for entry in data.get("tabelas", []):
+            entry = dict(entry)
+            unnest = entry.pop("unnest_classifications", False)
+            split_vars = entry.pop("split_variables", False)
+
+            if unnest is not False:
+                metadados = self.fetcher.sidra_client.get_agregado_metadados(
+                    entry["tabela_sidra"]
+                )
+                if unnest is True:
+                    entry.pop("classifications", None)
+                    to_unnest = metadados.classificacoes
+                    static_cls: dict[str, list[str]] = {}
+                else:
+                    unnest_ids = {str(u) for u in unnest}
+                    to_unnest = [
+                        c for c in metadados.classificacoes if str(c.id) in unnest_ids
+                    ]
+                    static_cls = entry.pop("classifications", {})
+                    for c in metadados.classificacoes:
+                        cls_id = str(c.id)
+                        if cls_id not in unnest_ids and cls_id not in static_cls:
+                            static_cls[cls_id] = ["all"]
+                for cls_combo in sidra.unnest_classificacoes(to_unnest):
+                    result.append(
+                        {
+                            **entry,
+                            "classifications": {**static_cls, **cls_combo},
+                        }
+                    )
+            elif split_vars:
+                variables = entry.pop("variables")
+                for var in variables:
+                    result.append({**entry, "variables": [var]})
+            else:
+                result.append(entry)
+
+        return result
+
+    def download(self, tabelas: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Download all tables and return a list of data-file descriptors.
+
+        Builds a flat plan across every requested table and submits it
+        through a single ``ThreadPoolExecutor`` so downloads parallelize
+        across table boundaries, not just across periods of one table.
+        """
+        plan: list[tuple[dict[str, Any], Any, str]] = []
+        for tabela in tabelas:
+            for parameter, modification in self.fetcher.plan_periods(**tabela):
+                plan.append((tabela, parameter, modification))
+        results = self.fetcher.download_periods(plan)
+        return [
+            r["key"] | {"filepath": r["filepath"], "modificacao": r["modificacao"]}
+            for r in results
+        ]
+
+    def load_metadata(self, engine: sa.Engine, tabelas: Iterable[dict[str, Any]]):
+        """Fetch and persist metadata for all unique SIDRA tables."""
+        seen: set[str] = set()
+        for tabela in tabelas:
+            tabela_sidra_id = tabela["tabela_sidra"]
+            if tabela_sidra_id in seen:
+                continue
+            seen.add(tabela_sidra_id)
+
+            metadata_filepath = self.storage.get_metadata_filepath(tabela_sidra_id)
+            if metadata_filepath.exists() and not self.force_metadata:
+                logger.info("Reading cached metadata for table %s", tabela_sidra_id)
+                agregado = self.storage.read_metadata(tabela_sidra_id)
+            else:
+                logger.info("Fetching metadata for table %s", tabela_sidra_id)
+                agregado = self.fetcher.fetch_metadata(tabela_sidra_id)
+                self.storage.write_metadata(agregado)
+
+            logger.info("Saving metadata to database for table %s", tabela_sidra_id)
+            database.save_agregado(engine, agregado)
+
+    def run(self):
+        """Execute the full fetch-and-load pipeline."""
+        engine = database.get_engine(self.config)
+        with engine.begin() as conn:
+            conn.exec_driver_sql(
+                f'CREATE SCHEMA IF NOT EXISTS "{self.config.db_schema}"'
+            )
+        models.Base.metadata.create_all(engine)
+        database.ensure_vintage_schema(engine)
+        try:
+            self._run(engine)
+        except KeyboardInterrupt:
+            if self.console is not None:
+                self.console.print("\n[yellow]Interrompido.[/yellow]")
+            raise SystemExit(1) from None
+
+    def _run(self, engine: sa.Engine):
+        tabelas = list(self.get_tabelas())
+        n_meta = len({t["tabela_sidra"] for t in tabelas})
+        s_meta = "tabela" if n_meta == 1 else "tabelas"
+
+        with self.fetcher:
+            with _make_progress(self.console) as progress:
+                meta_task = progress.add_task(
+                    f"Metadados ({n_meta} {s_meta})", total=None, main=True
+                )
+                self.load_metadata(engine, tabelas)
+                progress.update(
+                    meta_task,
+                    total=1,
+                    completed=1,
+                    description=f"Metadados ({n_meta} {s_meta})",
+                )
+
+            plan: list[tuple[dict[str, Any], Any, str]] = []
+            for tabela in tabelas:
+                for parameter, modification in self.fetcher.plan_periods(**tabela):
+                    plan.append((tabela, parameter, modification))
+
+            n_plan = len(plan)
+            if self.console is not None:
+                info = Table.grid(padding=(0, 2))
+                info.add_column(style="bold")
+                info.add_column()
+                info.add_row("Pipeline", str(self.toml_path))
+                info.add_row("Tabelas", f"{n_meta} {s_meta}")
+                info.add_row("Arquivos", str(n_plan))
+                info.add_row("Threads", str(self.fetcher.max_workers))
+                info.add_row(
+                    "Banco",
+                    f"{self.config.db_host}:{self.config.db_port}/{self.config.db_name}"
+                    f"  schema={self.config.db_schema}",
+                )
+                info.add_row("Storage", str(self.config.data_dir))
+                self.console.print(info)
+                self.console.print()
+
+            files_per_table: dict[str, int] = {}
+            for tabela, _, _ in plan:
+                sid = tabela["tabela_sidra"]
+                files_per_table[sid] = files_per_table.get(sid, 0) + 1
+
+            with _make_download_progress(self.console) as progress:
+                global_task = progress.add_task("Download", total=n_plan, main=True)
+                task_by_table: dict[str, TaskID] = {}
+                if len(files_per_table) > 1:
+                    for sid, count in files_per_table.items():
+                        task_by_table[sid] = progress.add_task(
+                            f"Tabela {sid}", total=count
+                        )
+
+                def _on_done(key: dict[str, Any]) -> None:
+                    sub = task_by_table.get(key["tabela_sidra"])
+                    if sub is not None:
+                        progress.advance(sub)
+                    progress.advance(global_task)
+
+                results = self.fetcher.download_periods(plan, on_file_done=_on_done)
+                data_files = [
+                    r["key"]
+                    | {
+                        "filepath": r["filepath"],
+                        "modificacao": r["modificacao"],
+                    }
+                    for r in results
+                ]
+                for sid, task_id in task_by_table.items():
+                    progress.update(task_id, description=f"Tabela {sid} ✓")
+                progress.update(global_task, description="Download concluído ✓")
+
+        if not self.force_load:
+            loaded = database.get_loaded_filenames(
+                engine, {f["filepath"].name for f in data_files}
+            )
+            data_files = [f for f in data_files if f["filepath"].name not in loaded]
+        if not data_files:
+            if self.console is not None:
+                self.console.print(
+                    "[green]Todos os arquivos já foram carregados.[/green]"
+                )
+            return
+
+        with _make_download_progress(self.console) as progress:
+            db_files_per_table: dict[str, int] = {}
+            for d in data_files:
+                sid = str(d["tabela_sidra"])
+                db_files_per_table[sid] = db_files_per_table.get(sid, 0) + 1
+            n_db_files = sum(db_files_per_table.values())
+            db_global_task = progress.add_task(
+                "Carregando no banco de dados", total=n_db_files * 2, main=True
+            )
+            db_task_by_table: dict[str, TaskID] = {}
+            if len(db_files_per_table) > 1:
+                for sid, count in db_files_per_table.items():
+                    db_task_by_table[sid] = progress.add_task(
+                        f"Tabela {sid}", total=count * 2
+                    )
+
+            def _on_db_file_done(sid: str) -> None:
+                sub = db_task_by_table.get(sid)
+                if sub is not None:
+                    progress.advance(sub)
+                progress.advance(db_global_task)
+
+            def _on_db_table_done(sid: str) -> None:
+                sub = db_task_by_table.get(sid)
+                if sub is not None:
+                    progress.update(sub, description=f"Tabela {sid} ✓")
+
+            database.load_dados(
+                engine,
+                self.storage,
+                data_files,
+                on_file_done=_on_db_file_done,
+                on_table_done=_on_db_table_done,
+            )
+            progress.update(db_global_task, description="Carregamento concluído ✓")

sidra_sql/transform_runner.py

@@ -0,0 +1,153 @@
+"""Pipeline runner for SQL transformations defined by ``transform.toml``.
+
+A pipeline directory contains one ``transform.toml`` declaring one or more
+output tables/views. Each ``[[table]]`` entry references its own ``.sql``
+file in the same directory.
+
+TOML schema
+-----------
+::
+
+    [[table]]
+    name        = "ipca"
+    schema      = "analytics"
+    strategy    = "replace"     # "replace" or "view"
+    sql         = "ipca.sql"
+    description = "IPCA - série detalhada"
+
+    [[table]]
+    name        = "ipca_resumo"
+    schema      = "analytics"
+    strategy    = "view"
+    sql         = "ipca_resumo.sql"
+
+Required fields per entry: ``name``, ``schema``, ``strategy``, ``sql``.
+Optional: ``description``, ``primary_key``, ``indexes``.
+
+Strategies
+~~~~~~~~~~
+``replace``
+    ``DROP TABLE IF EXISTS`` followed by ``CREATE TABLE ... AS``.
+    Full refresh — best for batch imports into Power BI, Excel, etc.
+
+``view``
+    ``CREATE OR REPLACE VIEW``.  Zero storage cost, always up-to-date,
+    best for live database connections.
+
+Each entry runs in its own transaction; if one fails, previously
+materialized outputs from the same pipeline persist.
+
+The SQL queries use unqualified table names (``dados``, ``dimensao``,
+``localidade``, ``tabela_sidra``).  They resolve via the database
+``search_path`` set by ``get_engine`` from ``config.ini``.
+"""
+
+import logging
+import tomllib
+from pathlib import Path
+
+from rich.console import Console
+from rich.progress import (
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+
+from . import database
+from .config import Config
+
+logger = logging.getLogger(__name__)
+
+
+class TransformRunner:
+    """Run SQL transformations declared in a ``transform.toml`` file."""
+
+    def __init__(self, config: Config, toml_path: Path, console: Console | None = None):
+        self.config = config
+        self.toml_path = toml_path
+        self.console = console
+
+    def run(self):
+        with open(self.toml_path, "rb") as f:
+            data = tomllib.load(f)
+
+        tables = data.get("table")
+        if not isinstance(tables, list) or not tables:
+            raise ValueError(
+                f"{self.toml_path}: esperado um ou mais [[table]] (array). "
+                "O schema [table] singular foi removido — migre para [[table]] "
+                "com campo 'sql' explícito por entrada."
+            )
+
+        engine = database.get_engine(self.config)
+
+        with Progress(
+            SpinnerColumn(finished_text="[green]✓[/green]"),
+            TextColumn("[progress.description]{task.description}"),
+            TimeElapsedColumn(),
+            console=self.console,
+            transient=False,
+            disable=self.console is None,
+        ) as progress:
+            for entry in tables:
+                self._materialize(engine, entry, progress)
+
+    def _materialize(self, engine, entry: dict, progress: Progress) -> None:
+        missing = [f for f in ("name", "schema", "strategy", "sql") if f not in entry]
+        if missing:
+            raise ValueError(
+                f"{self.toml_path}: [[table]] sem campo(s) obrigatório(s): "
+                f"{', '.join(missing)}"
+            )
+
+        name = entry["name"]
+        schema = entry["schema"]
+        strategy = entry["strategy"]
+        sql_rel = entry["sql"]
+        primary_key = entry.get("primary_key")
+        indexes = entry.get("indexes", [])
+
+        sql_path = self.toml_path.parent / sql_rel
+        if not sql_path.exists():
+            raise FileNotFoundError(
+                f"{self.toml_path}: arquivo SQL '{sql_rel}' não encontrado em "
+                f"{self.toml_path.parent}"
+            )
+        query = sql_path.read_text(encoding="utf-8").strip().replace("%", "%%")
+
+        qualified = f'"{schema}"."{name}"'
+        strategy_label = {"replace": "tabela", "view": "view"}.get(strategy, strategy)
+        task = progress.add_task(
+            f"{qualified} [dim][{strategy_label}][/dim]", total=None
+        )
+
+        with engine.begin() as conn:
+            conn.exec_driver_sql(f'CREATE SCHEMA IF NOT EXISTS "{schema}"')
+
+            if strategy == "view":
+                conn.exec_driver_sql(f"CREATE OR REPLACE VIEW {qualified} AS\n{query}")
+            elif strategy == "replace":
+                conn.exec_driver_sql(f"DROP TABLE IF EXISTS {qualified}")
+                conn.exec_driver_sql(f"CREATE TABLE {qualified} AS\n{query}")
+
+                if primary_key:
+                    pk_cols = ", ".join(f'"{c}"' for c in primary_key)
+                    conn.exec_driver_sql(
+                        f"ALTER TABLE {qualified} ADD PRIMARY KEY ({pk_cols})"
+                    )
+
+                for idx in indexes:
+                    idx_name = idx["name"]
+                    idx_cols = ", ".join(f'"{c}"' for c in idx["columns"])
+                    unique = "UNIQUE" if idx.get("unique") else ""
+                    conn.exec_driver_sql(
+                        f'CREATE {unique} INDEX "{idx_name}" '
+                        f"ON {qualified} ({idx_cols})"
+                    )
+            else:
+                raise ValueError(
+                    f"Unknown strategy {strategy!r} for {qualified} in {self.toml_path}"
+                )
+
+        progress.update(task, total=1, completed=1)