Add TinybirdApi.sample_datasource() for bounded S3/GCS/DynamoDB sampling

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+
+- `TinybirdApi.sample_datasource()` starts a sample import job for S3/GCS/DynamoDB connected data sources via `POST /v0/datasources/{name}/sample`. For blob storage connectors `max_files` bounds the number of imported files; for DynamoDB the sample is bounded by either `rows` or `max_bytes` (mutually exclusive), or `full_export` triggers a full PITR export of the whole table instead of a bounded scan. This lets cloud branches and local import bounded DynamoDB samples, avoiding slow exports and unnecessary egress costs on branches.
+
 ## [0.4.0] - 2026-06-29
 
 ### Added

src/tinybird_sdk/api/api.py

@@ -339,6 +339,57 @@ def truncate_datasource(
         except json.JSONDecodeError:
             return {}
 
+    def sample_datasource(
+        self,
+        datasource_name: str,
+        options: dict[str, Any] | None = None,
+        api_options: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Start a sample import job for an S3/GCS/DynamoDB connected data source.
+
+        For blob storage (S3/GCS) connectors, ``max_files`` bounds how many files
+        are imported (default 1, max 10). For DynamoDB, the sample is bounded by
+        either ``rows`` or ``max_bytes`` (mutually exclusive), or ``full_export``
+        triggers a full PITR export of the whole table instead of a bounded scan.
+
+        Options:
+            max_files: Maximum number of files to import for blob storage connectors.
+            rows: For DynamoDB, the maximum number of rows to scan and import
+                (mutually exclusive with ``max_bytes``).
+            max_bytes: For DynamoDB, the maximum approximate JSONEachRow bytes to
+                import, e.g. ``"500MB"`` (mutually exclusive with ``rows``).
+            full_export: For DynamoDB, trigger a full PITR export instead of a
+                bounded scan.
+        """
+        options = options or {}
+        api_options = api_options or {}
+
+        rows = options.get("rows")
+        max_bytes = options.get("max_bytes")
+        if rows is not None and max_bytes is not None:
+            raise ValueError("'rows' and 'max_bytes' are mutually exclusive; pass only one")
+
+        payload: dict[str, Any] = {
+            "max_files": options.get("max_files", 1),
+            "full_export": bool(options.get("full_export", False)),
+        }
+        if rows is not None:
+            payload["rows"] = rows
+        if max_bytes is not None:
+            payload["max_bytes"] = max_bytes
+
+        response = self.request(
+            f"/v0/datasources/{datasource_name}/sample",
+            method="POST",
+            token=api_options.get("token"),
+            headers={"Content-Type": "application/json"},
+            body=json.dumps(payload),
+            timeout=options.get("timeout", api_options.get("timeout")),
+        )
+        if not response.ok:
+            self._raise_for_error(response.status_code, response.text)
+        return response.json()
+
     def create_token(
         self,
         body: dict[str, Any],

tests/test_api_datasource_sample.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+import pytest
+
+import tinybird_sdk.api.api as api_module
+from tinybird_sdk.api.api import TinybirdApi
+
+
+class _FakeResponse:
+    def __init__(self, status_code: int, payload: dict[str, Any]):
+        self.status_code = status_code
+        self._payload = payload
+        self.text = json.dumps(payload)
+
+    @property
+    def ok(self) -> bool:
+        return 200 <= self.status_code < 300
+
+    def json(self) -> dict[str, Any]:
+        return self._payload
+
+
+def _capture_fetch(captured: dict[str, Any]) -> Any:
+    def fake_fetch(url: str, **kwargs: Any) -> _FakeResponse:
+        captured["url"] = url
+        captured["method"] = kwargs.get("method")
+        captured["headers"] = kwargs.get("headers")
+        captured["body"] = kwargs.get("body")
+        return _FakeResponse(200, {"job_id": "job-1", "status": "waiting"})
+
+    return fake_fetch
+
+
+def _make_api() -> TinybirdApi:
+    return TinybirdApi({"base_url": "https://api.tinybird.co", "token": "p.test"})
+
+
+def test_sample_datasource_defaults(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured: dict[str, Any] = {}
+    monkeypatch.setattr(api_module, "tinybird_fetch", _capture_fetch(captured))
+
+    result = _make_api().sample_datasource("events")
+
+    assert result == {"job_id": "job-1", "status": "waiting"}
+    assert captured["method"] == "POST"
+    assert captured["url"].endswith("/v0/datasources/events/sample")
+    assert captured["headers"]["Content-Type"] == "application/json"
+    assert json.loads(captured["body"]) == {"max_files": 1, "full_export": False}
+
+
+def test_sample_datasource_forwards_max_files(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured: dict[str, Any] = {}
+    monkeypatch.setattr(api_module, "tinybird_fetch", _capture_fetch(captured))
+
+    _make_api().sample_datasource("events", {"max_files": 3})
+
+    assert json.loads(captured["body"]) == {"max_files": 3, "full_export": False}
+
+
+def test_sample_datasource_forwards_dynamodb_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured: dict[str, Any] = {}
+    monkeypatch.setattr(api_module, "tinybird_fetch", _capture_fetch(captured))
+
+    _make_api().sample_datasource("ddb_ds", {"rows": 100000})
+
+    assert json.loads(captured["body"]) == {"max_files": 1, "full_export": False, "rows": 100000}
+
+
+def test_sample_datasource_forwards_dynamodb_max_bytes(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured: dict[str, Any] = {}
+    monkeypatch.setattr(api_module, "tinybird_fetch", _capture_fetch(captured))
+
+    _make_api().sample_datasource("ddb_ds", {"max_bytes": "1GB"})
+
+    assert json.loads(captured["body"]) == {
+        "max_files": 1,
+        "full_export": False,
+        "max_bytes": "1GB",
+    }
+
+
+def test_sample_datasource_forwards_full_export(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured: dict[str, Any] = {}
+    monkeypatch.setattr(api_module, "tinybird_fetch", _capture_fetch(captured))
+
+    _make_api().sample_datasource("ddb_ds", {"full_export": True})
+
+    assert json.loads(captured["body"]) == {"max_files": 1, "full_export": True}
+
+
+def test_sample_datasource_rows_and_max_bytes_mutually_exclusive() -> None:
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        _make_api().sample_datasource("ddb_ds", {"rows": 10, "max_bytes": "1GB"})