keboola-cli 0.63.4__py3-none-any.whl

keboola_agent_cli/services/_sync_storage.py

@@ -0,0 +1,376 @@
+"""Pull-side storage-metadata, jobs, and sample writers (from sync_service.py).
+
+Free functions that materialise a project's storage metadata (buckets/tables),
+per-config job history, and table data samples to the filesystem during
+``sync pull``. Only :func:`fetch_jobs_per_config` needs the ``SyncService`` (for
+its ``_resolve_max_workers`` helper); the rest are pure. ``_ensure_path_within``
+(the storage-write path-traversal guard) moved here with its only callers.
+"""
+
+from __future__ import annotations
+
+import csv
+import io
+import json
+import logging
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from ..constants import (
+    ENCRYPTED_COLUMN_MASK,
+    ENCRYPTED_COLUMN_PREFIX,
+    JOBS_FILENAME,
+    STORAGE_BUCKETS_FILENAME,
+    STORAGE_DIR_NAME,
+    STORAGE_SAMPLES_DIR_NAME,
+)
+from ..errors import ConfigError
+from ..sync.manifest import ManifestConfiguration
+from ..sync.naming import sanitize_path_segment
+
+if TYPE_CHECKING:
+    from .sync_service import SyncService
+
+logger = logging.getLogger(__name__)
+
+# Storage API sync-export column ceiling (avoids the API 400 on wide tables).
+_MAX_SYNC_COLUMNS = 30
+
+
+def _ensure_path_within(base_dir: Path, target: Path, what: str) -> None:
+    """Reject a write whose path escapes *base_dir* (defense-in-depth).
+
+    Mirrors ``_ensure_within_branch`` for non-config writes (storage metadata +
+    samples) whose path segments derive from API-controlled bucket ids / table
+    names (GHSA-833q-c5wv-26r7). Raises ConfigError on escape so a malformed or
+    compromised Storage response cannot write outside the sync workspace.
+    """
+    try:
+        base_resolved = base_dir.resolve()
+        target_resolved = target.resolve()
+    except OSError as exc:
+        raise ConfigError(f"Cannot resolve sync path: {exc}") from exc
+    if not target_resolved.is_relative_to(base_resolved):
+        raise ConfigError(
+            f"Storage path escapes sync workspace ({what}). Refusing to write "
+            f"outside '{base_resolved}'. This indicates a malformed or compromised "
+            f"API response or a path-sanitization regression."
+        )
+
+
+def write_storage_metadata(
+    project_root: Path,
+    buckets: list[dict[str, Any]],
+    tables: list[dict[str, Any]],
+    samples: dict[str, str],
+) -> dict[str, int]:
+    """Write storage bucket and table metadata to the filesystem.
+
+    Creates:
+        storage/buckets.json - list of all buckets
+        storage/tables/{bucket_id}/{table_name}.json - per-table metadata
+        storage/samples/{bucket}/{table}/sample.csv - data samples (if any)
+
+    Returns:
+        Dict with counts: buckets, tables, samples written.
+    """
+    storage_dir = project_root / STORAGE_DIR_NAME
+    storage_dir.mkdir(parents=True, exist_ok=True)
+
+    # Write buckets index. API may return null for tablesCount / dataSizeBytes
+    # on empty buckets; coerce to 0 (dict.get default only fires when key is
+    # missing, not when the value is explicitly null).
+    bucket_summaries = [
+        {
+            "id": b.get("id", ""),
+            "name": b.get("name", ""),
+            "stage": b.get("stage", ""),
+            "description": b.get("description", ""),
+            "tables_count": b.get("tablesCount") or 0,
+            "data_size_bytes": b.get("dataSizeBytes") or 0,
+            "metadata": b.get("metadata", []),
+        }
+        for b in buckets
+    ]
+    buckets_file = storage_dir / STORAGE_BUCKETS_FILENAME
+    buckets_file.write_text(
+        json.dumps(bucket_summaries, indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )
+
+    # Group tables by bucket
+    tables_by_bucket: dict[str, list[dict[str, Any]]] = {}
+    for t in tables:
+        bucket_id = (
+            t.get("bucket", {}).get("id", "")
+            if isinstance(t.get("bucket"), dict)
+            else t.get("bucketId", "")
+        )
+        if not bucket_id:
+            continue
+        tables_by_bucket.setdefault(bucket_id, []).append(t)
+
+    tables_written = 0
+    tables_dir = storage_dir / "tables"
+    for bucket_id, bucket_tables in tables_by_bucket.items():
+        # Sanitize bucket_id for filesystem. sanitize_path_segment first kills
+        # traversal (`/`, `..`, absolute paths); the trailing replace keeps the
+        # legacy `in.c-foo` -> `in-c-foo` directory naming for legitimate ids
+        # (GHSA-833q-c5wv-26r7).
+        safe_bucket = sanitize_path_segment(bucket_id).replace(".", "-")
+        bucket_dir = tables_dir / safe_bucket
+        _ensure_path_within(storage_dir, bucket_dir, f"bucket_id={bucket_id!r}")
+        bucket_dir.mkdir(parents=True, exist_ok=True)
+
+        for t in bucket_tables:
+            table_name = t.get("name", "unknown")
+            table_meta = {
+                "id": t.get("id", ""),
+                "name": table_name,
+                "primary_key": t.get("primaryKey", []),
+                "columns": t.get("columns", []),
+                # API may return null for rowsCount / dataSizeBytes on
+                # newly-created or empty tables; coerce to 0 explicitly
+                # (dict.get default only fires when the key is missing).
+                "rows_count": t.get("rowsCount") or 0,
+                "data_size_bytes": t.get("dataSizeBytes") or 0,
+                "last_import_date": t.get("lastImportDate", ""),
+                "last_change_date": t.get("lastChangeDate", ""),
+                "description": t.get("description", ""),
+                "metadata": t.get("metadata", []),
+                "column_metadata": t.get("columnMetadata", {}),
+            }
+            # The table name comes from the API; sanitize it for the filename
+            # and assert containment so a crafted name cannot escape the bucket
+            # dir (GHSA-833q-c5wv-26r7). The original name stays verbatim in the
+            # metadata body above.
+            safe_table = sanitize_path_segment(table_name)
+            table_file = bucket_dir / f"{safe_table}.json"
+            _ensure_path_within(storage_dir, table_file, f"table={table_name!r}")
+            table_file.write_text(
+                json.dumps(table_meta, indent=2, ensure_ascii=False),
+                encoding="utf-8",
+            )
+            tables_written += 1
+
+    # Write samples
+    samples_written = 0
+    if samples:
+        samples_dir = storage_dir / STORAGE_SAMPLES_DIR_NAME
+        for table_id, csv_data in samples.items():
+            # table_id format: "in.c-bucket.table" -> samples/in-c-bucket/table/
+            # Every segment derives from the API table_id; sanitize each and
+            # assert containment so a crafted id cannot escape (GHSA-833q).
+            parts = table_id.split(".", 2)
+            if len(parts) >= 3:
+                safe_bucket = sanitize_path_segment(f"{parts[0]}-{parts[1]}")
+                safe_table = sanitize_path_segment(parts[2])
+            else:
+                safe_bucket = sanitize_path_segment(table_id.replace(".", "-"))
+                safe_table = "data"
+            sample_dir = samples_dir / safe_bucket / safe_table
+            _ensure_path_within(storage_dir, sample_dir, f"table_id={table_id!r}")
+            sample_dir.mkdir(parents=True, exist_ok=True)
+
+            masked_csv = mask_encrypted_columns(csv_data)
+            (sample_dir / "sample.csv").write_text(masked_csv, encoding="utf-8")
+            samples_written += 1
+
+    return {
+        "buckets": len(buckets),
+        "tables": tables_written,
+        "samples": samples_written,
+    }
+
+
+def fetch_jobs_per_config(
+    service: SyncService,
+    client: Any,
+    components: list[dict[str, Any]],
+    job_limit: int,
+) -> list[dict[str, Any]]:
+    """Fetch jobs per config via /search/jobs in parallel.
+
+    Used as fallback when the grouped-jobs API cannot return all configs in a
+    single call (jobsPerGroup * limit <= 500 constraint). Returns data in the
+    same format as ``list_jobs_grouped()`` so :func:`write_per_config_jobs`
+    works unchanged.
+    """
+    config_pairs: list[tuple[str, str]] = []
+    for comp in components:
+        comp_id = comp.get("id", "")
+        for cfg in comp.get("configurations", []):
+            cfg_id = str(cfg.get("id", ""))
+            if comp_id and cfg_id:
+                config_pairs.append((comp_id, cfg_id))
+
+    if not config_pairs:
+        return []
+
+    results: list[dict[str, Any]] = []
+    lock = threading.Lock()
+    max_workers = min(len(config_pairs), service._resolve_max_workers())
+
+    def _fetch_one(pair: tuple[str, str]) -> None:
+        comp_id, cfg_id = pair
+        try:
+            jobs = client.list_jobs(
+                component_id=comp_id,
+                config_id=cfg_id,
+                limit=job_limit,
+            )
+            if jobs:
+                with lock:
+                    results.append(
+                        {
+                            "group": {"componentId": comp_id, "configId": cfg_id},
+                            "jobs": jobs,
+                        }
+                    )
+        except Exception:
+            logger.debug("Failed to fetch jobs for %s/%s", comp_id, cfg_id, exc_info=True)
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [executor.submit(_fetch_one, pair) for pair in config_pairs]
+        for future in as_completed(futures):
+            future.result()  # propagate unexpected errors
+
+    return results
+
+
+def write_per_config_jobs(
+    branch_dir: Path,
+    configurations: list[ManifestConfiguration],
+    jobs_grouped: list[dict[str, Any]],
+) -> int:
+    """Write ``_jobs.jsonl`` files next to each configuration.
+
+    Matches grouped jobs to configs by componentId+configId, then writes a JSONL
+    file with light job records. Returns the number of files written.
+    """
+    jobs_by_config: dict[tuple[str, str], list[dict[str, Any]]] = {}
+    for group in jobs_grouped:
+        group_key = group.get("group", {})
+        component_id = group_key.get("componentId", "")
+        config_id = group_key.get("configId", "")
+        if component_id and config_id:
+            jobs_by_config[(component_id, config_id)] = group.get("jobs", [])
+
+    files_written = 0
+    for cfg in configurations:
+        key = (cfg.component_id, cfg.id)
+        jobs = jobs_by_config.get(key)
+        if not jobs:
+            continue
+
+        config_dir = branch_dir / cfg.path
+        config_dir.mkdir(parents=True, exist_ok=True)
+        jobs_file = config_dir / JOBS_FILENAME
+
+        lines: list[str] = []
+        for job in jobs:
+            light_job: dict[str, Any] = {
+                "id": str(job.get("id", "")),
+                "status": job.get("status", ""),
+                "start_time": job.get("startTime", ""),
+                "end_time": job.get("endTime", ""),
+                "duration_seconds": job.get("durationSeconds", 0),
+            }
+            if job.get("mode") and job["mode"] != "run":
+                light_job["mode"] = job["mode"]
+            # Include error message for failed/warning jobs
+            status = job.get("status", "")
+            if status in ("error", "warning", "terminated", "cancelled"):
+                result = job.get("result", {})
+                if isinstance(result, dict) and result.get("message"):
+                    light_job["error_message"] = result["message"]
+            lines.append(json.dumps(light_job, ensure_ascii=False))
+
+        jobs_file.write_text("\n".join(lines) + "\n", encoding="utf-8")
+        files_written += 1
+
+    return files_written
+
+
+def fetch_samples(
+    client: Any,
+    tables: list[dict[str, Any]],
+    sample_limit: int,
+    max_samples: int,
+) -> dict[str, str]:
+    """Fetch CSV data previews for tables, respecting limits.
+
+    Selects tables sorted by rowsCount descending (largest first), limited to
+    ``max_samples`` tables. Returns a dict mapping table_id -> CSV string.
+    """
+
+    # Storage API may return rowsCount=None for empty/newly-created tables on
+    # some backends; dict.get() does not coerce None to the default value.
+    def _rows(t: dict[str, Any]) -> int:
+        return t.get("rowsCount") or 0
+
+    sorted_tables = sorted(
+        [t for t in tables if _rows(t) > 0],
+        key=_rows,
+        reverse=True,
+    )[:max_samples]
+
+    samples: dict[str, str] = {}
+    for t in sorted_tables:
+        table_id = t.get("id", "")
+        if not table_id:
+            continue
+        try:
+            # Limit columns to avoid the API 400 error on wide tables.
+            all_columns = t.get("columns", [])
+            columns = (
+                all_columns[:_MAX_SYNC_COLUMNS] if len(all_columns) > _MAX_SYNC_COLUMNS else None
+            )
+            csv_data = client.get_table_data_preview(table_id, limit=sample_limit, columns=columns)
+            samples[table_id] = csv_data
+        except Exception:
+            logger.warning("Failed to fetch sample for %s", table_id, exc_info=True)
+
+    return samples
+
+
+def mask_encrypted_columns(csv_data: str) -> str:
+    """Mask encrypted column values in CSV data.
+
+    Encrypted columns in Keboola start with '#' in the column name; their values
+    are replaced with the masked placeholder.
+    """
+    if not csv_data:
+        return csv_data
+
+    lines = csv_data.split("\n")
+    if not lines:
+        return csv_data
+
+    reader = csv.reader(io.StringIO(lines[0]))
+    try:
+        header = next(reader)
+    except StopIteration:
+        return csv_data
+
+    encrypted_indices = [
+        i for i, col in enumerate(header) if col.startswith(ENCRYPTED_COLUMN_PREFIX)
+    ]
+    if not encrypted_indices:
+        return csv_data
+
+    output = io.StringIO()
+    writer = csv.writer(output)
+    full_reader = csv.reader(io.StringIO(csv_data))
+    for row_idx, row in enumerate(full_reader):
+        if row_idx == 0:
+            writer.writerow(row)  # header unchanged
+        else:
+            for idx in encrypted_indices:
+                if idx < len(row):
+                    row[idx] = ENCRYPTED_COLUMN_MASK
+            writer.writerow(row)
+
+    return output.getvalue()

keboola_agent_cli/services/_sync_writeback.py

@@ -0,0 +1,167 @@
+"""Manifest + local-file writeback after a push (extracted from sync_service.py).
+
+Free functions that record a freshly-created config/row in the manifest, write
+the API-assigned ULID + encrypted secrets back into the local ``_config.yml``,
+and propagate ``KBC.*`` metadata. Only :func:`writeback_after_push` needs the
+``SyncService`` (for its ``_write_config_file`` helper); the rest are pure.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from ..errors import KeboolaApiError
+from ..sync.manifest import ManifestConfigRow, ManifestConfiguration
+from ._encryption import apply_encrypted_to_local
+from ._sync_models import WritebackResult
+
+if TYPE_CHECKING:
+    from ..sync.manifest import Manifest
+    from .sync_service import SyncService
+
+logger = logging.getLogger(__name__)
+
+
+def writeback_create_config_in_manifest(
+    *,
+    manifest: Manifest,
+    component_id: str,
+    branch_id: int | None,
+    config_path_str: str,
+    new_id: str,
+    file_hash: str,
+    cfg_hash: str,
+) -> WritebackResult:
+    """Record a freshly-created config in the manifest.
+
+    If a placeholder entry already exists at ``(branch_id, component_id, path)``
+    -- the FIIA / scaffold emit pattern -- update it in place, preserving any
+    user-declared metadata (e.g. ``KBC.configuration.folderName``) and
+    refreshing only the bookkeeping hashes. Otherwise append a new entry.
+
+    Matching includes ``branch_id`` because a single manifest can hold entries
+    from multiple branches in git-branching mode; matching on
+    ``(component_id, path)`` alone would risk updating the wrong branch's entry
+    when the same logical path exists under two branches.
+
+    Returns a :class:`WritebackResult` carrying the entry and its pre-overwrite
+    ``previous_id`` so the create pass can remap any child row parents /
+    transformation variable links from the placeholder id to the freshly-assigned
+    ULID.
+    """
+    target_branch = branch_id or 0
+    for entry in manifest.configurations:
+        if (
+            entry.branch_id == target_branch
+            and entry.component_id == component_id
+            and entry.path == config_path_str
+        ):
+            previous_id = entry.id
+            entry.id = new_id
+            entry.metadata["pull_hash"] = file_hash
+            entry.metadata["pull_config_hash"] = cfg_hash
+            return WritebackResult(entry=entry, previous_id=previous_id)
+    new_entry = ManifestConfiguration(
+        branchId=target_branch,
+        componentId=component_id,
+        id=new_id,
+        path=config_path_str,
+        metadata={"pull_hash": file_hash, "pull_config_hash": cfg_hash},
+    )
+    manifest.configurations.append(new_entry)
+    return WritebackResult(entry=new_entry, previous_id="")
+
+
+def writeback_create_row_in_manifest(
+    *,
+    parent: ManifestConfiguration,
+    row_path_str: str,
+    new_row_id: str,
+    file_hash: str,
+    cfg_hash: str,
+) -> ManifestConfigRow:
+    """Record a freshly-created row under its parent in the manifest.
+
+    Mirrors :func:`writeback_create_config_in_manifest` for rows: update any
+    placeholder row entry in place, otherwise append.
+    """
+    for row in parent.rows:
+        if row.path == row_path_str:
+            row.id = new_row_id
+            row.metadata["pull_hash"] = file_hash
+            row.metadata["pull_config_hash"] = cfg_hash
+            return row
+    new_row = ManifestConfigRow(
+        id=new_row_id,
+        path=row_path_str,
+        metadata={"pull_hash": file_hash, "pull_config_hash": cfg_hash},
+    )
+    parent.rows.append(new_row)
+    return new_row
+
+
+def propagate_kbc_metadata(
+    client: Any,
+    entry: ManifestConfiguration,
+    branch_id: int | None,
+) -> str | None:
+    """POST any ``KBC.*`` keys from the manifest entry to the metadata API.
+
+    Bookkeeping keys (``pull_hash``, ``pull_config_hash``, ...) live in the same
+    metadata dict but are filtered by the ``KBC.`` prefix. Called only on CREATE;
+    updates use ``kbagent config set-metadata`` explicitly. The metadata API
+    stores configuration-level annotations only -- this is **not** a secret
+    store; do not place tokens or passwords under ``KBC.*`` keys.
+
+    Returns ``None`` on success (or when there are no KBC.* keys to propagate).
+    Returns the API error message on a non-fatal write failure: the config is
+    already created on the remote and the manifest writeback is complete, so a
+    single failed metadata POST is reported back to the push loop as an
+    accumulated error rather than aborting the rest of the push.
+    """
+    entries = [(key, str(value)) for key, value in entry.metadata.items() if key.startswith("KBC.")]
+    if not entries:
+        return None
+    try:
+        client.set_config_metadata(
+            component_id=entry.component_id,
+            config_id=entry.id,
+            entries=entries,
+            branch_id=branch_id,
+        )
+    except KeboolaApiError as exc:
+        logger.warning(
+            "Failed to propagate KBC.* metadata for %s/%s: %s",
+            entry.component_id,
+            entry.id,
+            exc,
+        )
+        return exc.message
+    return None
+
+
+def writeback_after_push(
+    service: SyncService,
+    local_data: dict[str, Any],
+    config_dir: Path,
+    config_id: str,
+    pushed_configuration: dict[str, Any],
+) -> None:
+    """Update local ``_config.yml`` after a successful push.
+
+    Writes back the API-assigned ``_keboola.config_id`` (on first create) and
+    the encrypted secret values (so local matches remote state).
+    """
+    keboola_meta = local_data.setdefault("_keboola", {})
+    if config_id:
+        keboola_meta["config_id"] = config_id
+
+    pushed_params = pushed_configuration.get("parameters", {})
+    local_params = local_data.get("parameters", {})
+    if pushed_params and local_params:
+        apply_encrypted_to_local(local_params, pushed_params)
+
+    service._write_config_file(config_dir, local_data)
+    logger.debug("Updated local config at %s after push", config_dir)