sidra-sql 1.3.0__py3-none-any.whl

sidra_sql/scaffold.py

@@ -0,0 +1,248 @@
+import subprocess
+import tomllib
+from pathlib import Path
+
+
+def _slugify(name: str) -> str:
+    return name.lower().replace("-", "_").replace(" ", "_")
+
+
+def _fetch_toml_template() -> str:
+    return (
+        "# Busca dados da API SIDRA (IBGE).\n"
+        "# Encontre IDs de tabelas em https://sidra.ibge.gov.br\n"
+        "#\n"
+        '# tabela_sidra    — ID da tabela no SIDRA (ex: "839")\n'
+        '# variables       — ["allxp"] para todas, ou IDs específicos: '
+        '["109", "216"]\n'
+        "# territories     — {6 = []} todos os municípios; {3 = []} todos os estados\n"
+        '# classifications — {81 = ["allxt"]} todas as categorias '
+        "(descomente se precisar)\n"
+        "# split_variables — true para enviar uma requisição por variável\n"
+        "\n"
+        "[[tabelas]]\n"
+        'tabela_sidra = "XXXX"        # substitua pelo ID da tabela\n'
+        'variables    = ["allxp"]\n'
+        "territories  = {6 = []}      # nível 6 = municípios\n"
+        '# classifications = {81 = ["allxt"]}\n'
+    )
+
+
+def _transform_toml_template(slug: str) -> str:
+    return (
+        "# Cada [[table]] declara uma saída do pipeline.\n"
+        "# Para múltiplas saídas, adicione mais blocos [[table]] e crie um\n"
+        "# arquivo .sql correspondente para cada um.\n"
+        "\n"
+        "[[table]]\n"
+        f'name        = "{slug}"\n'
+        'schema      = "analytics"\n'
+        'strategy    = "replace"        # "replace" ou "view"\n'
+        f'sql         = "{slug}.sql"\n'
+        'description = "Descrição da tabela de saída"\n'
+    )
+
+
+def _transform_sql_template() -> str:
+    return (
+        "-- Adapte esta query para o seu pipeline.\n"
+        "-- Tabelas normalizadas disponíveis:\n"
+        "--   IBGE_SIDRA.DADOS  — valores brutos\n"
+        "--   PERIODO           — dimensão temporal (ano, mês, trimestre...)\n"
+        "--   DIMENSAO          — variáveis e classificações\n"
+        "--   LOCALIDADE        — unidades territoriais\n"
+        "\n"
+        "SELECT\n"
+        "    P.ANO           AS ANO,\n"
+        "    L.D1C           AS ID_MUNICIPIO,\n"
+        "    L.D1N           AS NOME_MUNICIPIO,\n"
+        "    DIM.D2N         AS VARIAVEL,\n"
+        "    DIM.MN          AS UNIDADE,\n"
+        "    CASE WHEN D.V ~ '^-?[0-9]' THEN D.V::NUMERIC END AS VALOR\n"
+        "FROM\n"
+        "    IBGE_SIDRA.DADOS D\n"
+        "    JOIN PERIODO    P   ON D.PERIODO_ID    = P.ID\n"
+        "    JOIN DIMENSAO   DIM ON D.DIMENSAO_ID   = DIM.ID\n"
+        "    JOIN LOCALIDADE L   ON D.LOCALIDADE_ID = L.ID\n"
+        "WHERE\n"
+        "    D.tabela_sidra_ID IN ('XXXX')  -- substitua pelo(s) ID(s) "
+        "da(s) tabela(s)\n"
+        "    AND D.ATIVO = TRUE;\n"
+    )
+
+
+class PluginScaffolder:
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        version: str,
+        output_dir: Path,
+        git_init: bool,
+    ):
+        self.name = name
+        self.slug = _slugify(name)
+        self.description = description
+        self.version = version
+        self.plugin_dir = Path(output_dir) / name
+        self.git_init = git_init
+
+    def create(self) -> Path:
+        if self.plugin_dir.exists():
+            raise FileExistsError(f"Directory '{self.plugin_dir}' already exists.")
+
+        self.plugin_dir.mkdir(parents=True)
+        pipeline_dir = self.plugin_dir / self.slug
+        pipeline_dir.mkdir()
+
+        self._write(self.plugin_dir / "manifest.toml", self._manifest())
+        self._write(self.plugin_dir / "README.md", self._readme())
+        self._write(pipeline_dir / "fetch.toml", _fetch_toml_template())
+        self._write(
+            pipeline_dir / "transform.toml",
+            _transform_toml_template(self.slug),
+        )
+        self._write(pipeline_dir / f"{self.slug}.sql", _transform_sql_template())
+
+        if self.git_init:
+            self._write(self.plugin_dir / ".gitignore", self._gitignore())
+            self._run_git_init()
+
+        return self.plugin_dir
+
+    def _write(self, path: Path, content: str) -> None:
+        path.write_text(content, encoding="utf-8")
+
+    def _manifest(self) -> str:
+        return (
+            f'name        = "{self.name}"\n'
+            f'description = "{self.description}"\n'
+            f'version     = "{self.version}"\n'
+            f"\n"
+            f"[[pipeline]]\n"
+            f'id          = "{self.slug}"\n'
+            f'description = "Descrição do pipeline"\n'
+            f'path        = "{self.slug}"\n'
+        )
+
+    def _readme(self) -> str:
+        return (
+            f"# {self.name}\n"
+            "\n"
+            f"{self.description or 'Descrição do plugin.'}\n"
+            "\n"
+            "## Instalação\n"
+            "\n"
+            "```bash\n"
+            "sidra-sql plugin install <git-url>\n"
+            "```\n"
+            "\n"
+            "## Pipelines\n"
+            "\n"
+            "| ID | Descrição | Path |\n"
+            "|---|---|---|\n"
+            f"| {self.slug} | Descrição do pipeline | {self.slug}/ |\n"
+            "\n"
+            "## Desenvolvimento\n"
+            "\n"
+            "1. Encontre a tabela desejada em https://sidra.ibge.gov.br\n"
+            f"2. Edite `{self.slug}/fetch.toml` com o ID da tabela e variáveis\n"
+            f"3. Ajuste `{self.slug}/{self.slug}.sql` para a transformação desejada\n"
+            f"4. Atualize `{self.slug}/transform.toml` com o nome da tabela de saída\n"
+            "5. Adicione mais pipelines em `manifest.toml` conforme necessário\n"
+            "\n"
+            "### Territórios disponíveis\n"
+            "\n"
+            "| Código | Nível |\n"
+            "|---|---|\n"
+            "| 1 | Brasil |\n"
+            "| 2 | Grandes Regiões |\n"
+            "| 3 | Unidades da Federação |\n"
+            "| 6 | Municípios |\n"
+            "| 7 | Regiões Metropolitanas |\n"
+        )
+
+    def _gitignore(self) -> str:
+        return "__pycache__/\n*.py[cod]\n.env\n.DS_Store\n"
+
+    def _run_git_init(self) -> None:
+        cwd = str(self.plugin_dir)
+        try:
+            subprocess.run(["git", "init"], cwd=cwd, check=True, capture_output=True)
+            subprocess.run(
+                ["git", "add", "."], cwd=cwd, check=True, capture_output=True
+            )
+            subprocess.run(
+                ["git", "commit", "-m", "chore: initial scaffold"],
+                cwd=cwd,
+                check=True,
+                capture_output=True,
+            )
+        except FileNotFoundError as e:
+            raise RuntimeError(
+                "git não encontrado. Instale o Git ou use --no-git-init."
+            ) from e
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(
+                f"Falha ao inicializar repositório Git: {e.stderr.decode().strip()}"
+            ) from e
+
+
+class PipelineAdder:
+    def __init__(
+        self,
+        pipeline_id: str,
+        description: str,
+        path: str,
+        plugin_dir: Path,
+    ):
+        self.pipeline_id = pipeline_id
+        self.slug = _slugify(pipeline_id)
+        self.description = description
+        self.path = path or self.slug
+        self.plugin_dir = Path(plugin_dir)
+        self.manifest_path = self.plugin_dir / "manifest.toml"
+        self.pipeline_dir = self.plugin_dir / self.path
+
+    def add(self) -> Path:
+        if not self.manifest_path.exists():
+            raise FileNotFoundError(
+                f"manifest.toml não encontrado em '{self.plugin_dir}'. "
+                "Execute o comando dentro do diretório do plugin ou use --plugin-dir."
+            )
+
+        with open(self.manifest_path, "rb") as f:
+            manifest = tomllib.load(f)
+        existing_ids = {p["id"] for p in manifest.get("pipeline", [])}
+        if self.pipeline_id in existing_ids:
+            raise ValueError(
+                f"Pipeline '{self.pipeline_id}' já existe no manifest.toml."
+            )
+
+        if self.pipeline_dir.exists():
+            raise FileExistsError(f"Diretório '{self.pipeline_dir}' já existe.")
+
+        self.pipeline_dir.mkdir(parents=True)
+        self.pipeline_dir.joinpath("fetch.toml").write_text(
+            _fetch_toml_template(), encoding="utf-8"
+        )
+        self.pipeline_dir.joinpath("transform.toml").write_text(
+            _transform_toml_template(self.slug), encoding="utf-8"
+        )
+        self.pipeline_dir.joinpath(f"{self.slug}.sql").write_text(
+            _transform_sql_template(), encoding="utf-8"
+        )
+
+        self._append_to_manifest()
+        return self.pipeline_dir
+
+    def _append_to_manifest(self) -> None:
+        entry = (
+            "\n"
+            "[[pipeline]]\n"
+            f'id          = "{self.pipeline_id}"\n'
+            f'description = "{self.description}"\n'
+            f'path        = "{self.path}"\n'
+        )
+        with open(self.manifest_path, "a", encoding="utf-8") as f:
+            f.write(entry)

sidra_sql/sidra.py

@@ -0,0 +1,362 @@
+"""Utilities to fetch and store SIDRA tables.
+
+This module provides a `Fetcher` context-managed helper that wraps the
+`sidra_fetcher` client to download SIDRA tables as CSV-backed pandas
+DataFrames and write them to the project's data directory. It also
+contains a helper `unnest_classificacoes` which expands nested
+classification/category combinations into flat dictionaries suitable for
+request parameters.
+
+Public API
+- `Fetcher`: context-managed client for downloading SIDRA tables.
+- `unnest_classificacoes`: yields classification/category mappings.
+"""
+
+import logging
+import threading
+from collections.abc import Callable, Generator
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Any
+
+import httpx
+from sidra_fetcher.agregados import Agregado, Classificacao
+from sidra_fetcher.fetcher import SidraClient
+from sidra_fetcher.sidra import Formato, Parametro, Precisao
+
+from .config import Config
+from .storage import Storage
+
+logger = logging.getLogger(__name__)
+
+_MAX_RETRIES = 5
+_RETRY_BASE_DELAY = 5  # seconds; doubles on each attempt (5, 10, 20, 40, 80)
+
+# Transient network conditions that warrant a retry
+_TRANSIENT_ERRORS = (
+    httpx.ReadTimeout,
+    httpx.ConnectTimeout,
+    httpx.ConnectError,
+    httpx.RemoteProtocolError,
+    httpx.NetworkError,
+)
+
+
+class Fetcher:
+    """Helper to download SIDRA tables and save them locally.
+
+    This class wraps a `SidraClient` to provide higher-level operations
+    to download all periods of a given SIDRA table, write each period's
+    result to disk and return the written file paths.
+
+    Usage example::
+
+        with Fetcher() as f:
+            files = f.download_table(...)
+
+    Attributes:
+        sidra_client: An instance of `SidraClient` used to perform HTTP
+            requests to the SIDRA API.
+        storage: `Storage` repository where downloaded files are written.
+        max_workers: Maximum number of concurrent period downloads.
+    """
+
+    def __init__(
+        self,
+        config: Config,
+        max_workers: int = 4,
+        storage: Storage | None = None,
+    ):
+        self.sidra_client = SidraClient(timeout=600)
+        self.storage = storage if storage is not None else Storage.default(config)
+        self.max_workers = max_workers
+        self._cancel = threading.Event()
+
+    def plan_periods(
+        self,
+        tabela_sidra: str,
+        territories: dict[str, list[str]],
+        variables: list[str] | None = None,
+        classifications: dict[str, list[str]] | None = None,
+    ) -> list[tuple[Parametro, str]]:
+        """Build (Parametro, modification) tuples for every period of a table.
+
+        Pure planning — no downloads. Use ``download_periods`` to fetch
+        a flat plan concurrently across many tables.
+
+        Args:
+            tabela_sidra: SIDRA table code (numeric string accepted).
+            territories: Mapping of territory type codes to lists of
+                territory identifiers.
+            variables: Optional list of variable codes. Defaults to ["all"].
+            classifications: Optional classification → category mapping.
+                If omitted, defaults to empty list for each declared
+                classification (read from cached or fetched metadata).
+
+        Returns:
+            List of (Parametro, modification_iso_string) tuples — one per
+            period of the requested table.
+        """
+        if variables is None:
+            variables = ["all"]
+
+        # Use cached metadata when available — avoids redundant round-trips
+        # after load_metadata has already fetched and stored the Agregado.
+        metadata_path = self.storage.get_metadata_filepath(tabela_sidra)
+        if metadata_path.exists():
+            metadados = self.storage.read_metadata(tabela_sidra)
+        else:
+            metadados = self.sidra_client.get_agregado_metadados(int(tabela_sidra))
+
+        if classifications is None:
+            classifications = {str(c.id): [] for c in metadados.classificacoes}
+
+        periodos = getattr(
+            metadados, "periodos", None
+        ) or self.sidra_client.get_agregado_periodos(agregado_id=int(tabela_sidra))
+
+        period_params: list[tuple[Parametro, str]] = []
+        for periodo in periodos:
+            parameter = Parametro(
+                agregado=tabela_sidra,
+                territorios=territories,
+                variaveis=variables,
+                periodos=[periodo.id],
+                classificacoes=classifications,
+                decimais={"": Precisao.M},  # Precisão: Máxima
+                formato=Formato.A,
+            )
+            period_params.append((parameter, periodo.modificacao.isoformat()))
+        return period_params
+
+    def download_periods(
+        self,
+        plan: list[tuple[Any, Parametro, str]],
+        on_file_done: Callable[[Any], None] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Download many periods concurrently from a flat plan.
+
+        Submits every (Parametro, modification) entry of ``plan`` to a
+        single ``ThreadPoolExecutor`` capped at ``self.max_workers``,
+        regardless of which source table each entry came from.
+
+        Args:
+            plan: Tuples of (key, parameter, modification). ``key`` is
+                opaque metadata returned alongside each result so callers
+                can correlate downloads back to their originating request.
+            on_file_done: Optional callback fired once per completed
+                download (success or failure), useful for progress bars.
+
+        Returns:
+            List of dicts with keys "key", "filepath", "modificacao", in
+            completion order. Raises the first download error after all
+            futures complete.
+        """
+        results: list[dict[str, Any]] = []
+        errors: list[Exception] = []
+        executor = ThreadPoolExecutor(max_workers=self.max_workers)
+        try:
+            future_to_meta = {
+                executor.submit(self._download_period, parameter, modification): (
+                    key,
+                    modification,
+                )
+                for key, parameter, modification in plan
+            }
+            for future in as_completed(future_to_meta):
+                key, modification = future_to_meta[future]
+                try:
+                    results.append(
+                        {
+                            "key": key,
+                            "filepath": future.result(),
+                            "modificacao": modification,
+                        }
+                    )
+                except Exception as e:
+                    logger.error("Period download failed: %s", e)
+                    errors.append(e)
+                if on_file_done is not None:
+                    on_file_done(key)
+        except KeyboardInterrupt:
+            self._cancel.set()
+            executor.shutdown(wait=True, cancel_futures=True)
+            raise
+        else:
+            executor.shutdown(wait=True)
+        if errors:
+            raise errors[0]
+        return results
+
+    def download_table(
+        self,
+        tabela_sidra: str,
+        territories: dict[str, list[str]],
+        variables: list[str] | None = None,
+        classifications: dict[str, list[str]] | None = None,
+        on_file_done: Callable[[], None] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Download all periods of a single SIDRA table and save them to disk.
+
+        Convenience wrapper around ``plan_periods`` + ``download_periods``
+        for callers that only need to fetch one table at a time. To
+        parallelize across many tables, build a combined plan via
+        ``plan_periods`` and submit it through ``download_periods``.
+
+        Returns:
+            A list of dicts with keys "filepath" (Path) and "modificacao" (str).
+        """
+        plan = [
+            (None, parameter, modification)
+            for parameter, modification in self.plan_periods(
+                tabela_sidra=tabela_sidra,
+                territories=territories,
+                variables=variables,
+                classifications=classifications,
+            )
+        ]
+        results = self.download_periods(plan, on_file_done=on_file_done)
+        return [
+            {"filepath": r["filepath"], "modificacao": r["modificacao"]}
+            for r in results
+        ]
+
+    def fetch_metadata(self, tabela_sidra: str) -> Agregado:
+        """Fetch full metadata for a SIDRA table including localidades and periodos."""
+        agregado = self.sidra_client.get_agregado_metadados(int(tabela_sidra))
+
+        all_niveis = (
+            agregado.nivel_territorial.administrativo
+            + agregado.nivel_territorial.ibge
+            + agregado.nivel_territorial.especial
+        )
+
+        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            loc_futures = [
+                executor.submit(
+                    self.sidra_client.get_agregado_localidades,
+                    agregado_id=int(tabela_sidra),
+                    localidades_nivel=nivel,
+                )
+                for nivel in all_niveis
+            ]
+
+        localidades = []
+        for f in loc_futures:
+            localidades.extend(f.result())
+
+        agregado.localidades = localidades
+        agregado.periodos = self.sidra_client.get_agregado_periodos(int(tabela_sidra))
+        return agregado
+
+    def _download_period(
+        self,
+        parameter: Parametro,
+        modification: str,
+    ) -> Path:
+        """Download a single period and save it; return the destination path."""
+        if self._cancel.is_set():
+            raise InterruptedError("cancelled")
+        if self.storage.exists(parameter, modification):
+            filepath = self.storage.get_data_filepath(parameter, modification)
+            logger.debug("File already exists (cache hit): %s", filepath)
+            return filepath
+        logger.info(
+            "Downloading %s",
+            self.storage.get_data_filepath(parameter, modification).name,
+        )
+        data = self.get_table(parameter)
+        return self.storage.write_data(
+            data=data, parameter=parameter, modification=modification
+        )
+
+    def get_table(self, parameter: Parametro) -> dict:
+        """Request a SIDRA table and return it as a dictionary.
+
+        Retries up to `_MAX_RETRIES` times on transient network errors
+        using exponential backoff (5 s, 10 s, 20 s, …).  Raises the
+        underlying exception once all attempts are exhausted.
+
+        Args:
+            parameter: A `Parametro` instance with the desired request
+                configuration.
+
+        Returns:
+            A `dict` constructed from the JSON response.
+        """
+        url = parameter.url()
+        for attempt in range(_MAX_RETRIES):
+            if self._cancel.is_set():
+                raise InterruptedError("cancelled")
+            try:
+                return self.sidra_client.get(url)
+            except _TRANSIENT_ERRORS as e:
+                if attempt >= _MAX_RETRIES - 1:
+                    raise
+                delay = _RETRY_BASE_DELAY * (2**attempt)
+                logger.error("%s while fetching data: %s", type(e).__name__, e)
+                logger.info(
+                    "Retrying in %d s (attempt %d/%d)…",
+                    delay,
+                    attempt + 1,
+                    _MAX_RETRIES,
+                )
+                if self._cancel.wait(delay):
+                    raise InterruptedError("cancelled") from None
+
+    def __enter__(self):
+        """Enter the context manager and return this `Fetcher`."""
+        self.sidra_client.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Close resources held by the fetcher.
+
+        Delegates to the `SidraClient` context manager to ensure any
+        network resources are cleaned up. Arguments are forwarded from
+        the context manager protocol.
+        """
+        self.sidra_client.__exit__(exc_type, exc_value, traceback)
+
+
+def unnest_classificacoes(
+    classificacoes: list[Classificacao],
+    data: dict[str, list[str]] | None = None,
+) -> Generator[dict[str, list[str]], None, None]:
+    """Recursively enumerate classification/category combinations.
+
+    SIDRA classifications can be nested. This generator produces a flat
+    sequence of mappings suitable to pass as the ``classificacoes``
+    parameter when requesting aggregated data: each yielded dict maps a
+    classification id (string) to a single-element list containing a
+    category id (string).
+
+    The function skips categories with id "0" which usually represent
+    an undefined or "all" category.
+
+    Args:
+        classificacoes: List of `Classificacao` objects (from
+            `sidra_fetcher`) to expand.
+        data: Internal accumulator used by recursion; callers should
+            normally omit this argument.
+
+    Yields:
+        Dictionaries mapping classification id to a singleton list of
+        category ids, representing one combination of categories across
+        the provided classifications.
+    """
+    if data is None:
+        data = {}
+    if not classificacoes:
+        return
+    classificacao = classificacoes[0]
+    classificacao_id = str(classificacao.id)
+    for categoria in classificacao.categorias:
+        categoria_id = str(categoria.id)
+        if categoria_id == "0":
+            continue
+        new_data = {**data, classificacao_id: [categoria_id]}
+        if len(classificacoes) == 1:
+            yield new_data
+        else:
+            yield from unnest_classificacoes(classificacoes[1:], new_data)