From 016d59e4ca9e89e5ef6f54c2fa0d655fdb7e599c Mon Sep 17 00:00:00 2001 From: Tommy Healy Date: Thu, 1 Oct 2026 17:34:16 +0200 Subject: [PATCH] Add schema inference from real sample data (tb datasource analyze parity) TinybirdApi.analyze() wraps POST /v0/analyze, inferring column names, ClickHouse types, and CSV dialect from a local file or URL, matching `tb datasource analyze `. tinybird_sdk.infer.parse_analyze_response() turns the response into a structured AnalyzedSchema, and generate_schema_code() emits ready-to-use define_datasource(...) source with t.* validators, reusing the same type-mapping codegen already used for reverse codegen (live workspace -> Python). Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 7 ++ README.md | 23 +++++++ src/tinybird_sdk/api/api.py | 59 ++++++++++++++++ src/tinybird_sdk/infer/__init__.py | 10 +++ src/tinybird_sdk/infer/analyze.py | 104 +++++++++++++++++++++++++++++ tests/test_api_analyze.py | 94 ++++++++++++++++++++++++++ tests/test_infer_analyze.py | 50 ++++++++++++++ 7 files changed, 347 insertions(+) create mode 100644 src/tinybird_sdk/infer/analyze.py create mode 100644 tests/test_api_analyze.py create mode 100644 tests/test_infer_analyze.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c1ecb0..c07d0e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,13 @@ 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.analyze()` wraps the Tinybird Analyze API (`POST /v0/analyze`) to infer column names, ClickHouse types, and (for CSV) dialect from a local file or URL, matching `tb datasource analyze `. +- `tinybird_sdk.infer.parse_analyze_response()` and `tinybird_sdk.infer.generate_schema_code()` turn an analyze response into a structured `AnalyzedSchema` and generate ready-to-use `define_datasource(...)` source with `t.*` validators, so a datasource can be scaffolded from real sample data before it's declared in Python. + ## [0.4.0] - 2026-06-29 ### Added diff --git a/README.md b/README.md index e42f34b..ea13e54 100644 --- a/README.md +++ b/README.md @@ -966,6 +966,29 @@ params_schema = infer_params_schema(top_pages) output_schema = infer_output_schema(top_pages) ``` +It can also infer a schema from real sample data, before a datasource is even defined, matching `tb datasource analyze `: + +```python +from tinybird_sdk import TinybirdApi +from tinybird_sdk.infer import generate_schema_code, parse_analyze_response + +api = TinybirdApi({"base_url": "https://api.tinybird.co", "token": token}) + +response = api.analyze({"file": "./events.csv"}) +# or: api.analyze({"url": "https://example.com/events.ndjson"}) + +analyzed = parse_analyze_response(response) +print(generate_schema_code(analyzed, "events")) +# events = define_datasource('events', { +# 'schema': { +# 'id': t.int64(), +# 'name': t.string().nullable(), +# }, +# }) +# +# EventsRow = dict +``` + ## License MIT diff --git a/src/tinybird_sdk/api/api.py b/src/tinybird_sdk/api/api.py index 8ee18b5..60af2db 100644 --- a/src/tinybird_sdk/api/api.py +++ b/src/tinybird_sdk/api/api.py @@ -287,6 +287,65 @@ def append_datasource( self._raise_for_error(response.status_code, response.text) return response.json() + def analyze( + self, + options: dict[str, Any], + api_options: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Analyze a local file or URL and return inferred schema/dialect info. + + Wraps the Tinybird Analyze API (``POST /v0/analyze``), matching + ``tb datasource analyze ``: infers column names, types, + and (for CSV) dialect before a data source is even defined. + """ + api_options = api_options or {} + source_url = options.get("url") + file_path = options.get("file") + + if not source_url and not file_path: + raise ValueError("Either 'url' or 'file' must be provided in options") + if source_url and file_path: + raise ValueError("Only one of 'url' or 'file' can be provided, not both") + + timeout = options.get("timeout", api_options.get("timeout")) + + if source_url: + if not isinstance(source_url, str): + raise ValueError("'url' must be a valid string") + query: dict[str, str] = {"url": source_url} + detected_format = detect_data_format(source_url) + if detected_format: + query["format"] = detected_format + response = self.request( + f"/v0/analyze?{urlencode(query)}", + method="POST", + token=api_options.get("token"), + timeout=timeout, + ) + else: + if not isinstance(file_path, str): + raise ValueError("'file' must be a valid string path") + with open(file_path, "rb") as fp: + file_content = fp.read() + detected_format = detect_data_format(file_path) + query = {"format": detected_format} if detected_format else {} + content_type, multipart = create_multipart_body( + files=[("csv", file_path, file_content, None)], + ) + request_path = f"/v0/analyze?{urlencode(query)}" if query else "/v0/analyze" + response = self.request( + request_path, + method="POST", + token=api_options.get("token"), + headers={"Content-Type": content_type}, + body=multipart, + timeout=timeout, + ) + + if not response.ok: + self._raise_for_error(response.status_code, response.text) + return response.json() + def delete_datasource( self, datasource_name: str, diff --git a/src/tinybird_sdk/infer/__init__.py b/src/tinybird_sdk/infer/__init__.py index 4aa1da1..8d67c43 100644 --- a/src/tinybird_sdk/infer/__init__.py +++ b/src/tinybird_sdk/infer/__init__.py @@ -1,3 +1,9 @@ +from .analyze import ( + AnalyzedColumn, + AnalyzedSchema, + generate_schema_code, + parse_analyze_response, +) from .index import ( infer_row_schema, infer_params_schema, @@ -12,4 +18,8 @@ "infer_output_schema", "infer_materialized_target", "is_materialized_pipe", + "AnalyzedColumn", + "AnalyzedSchema", + "parse_analyze_response", + "generate_schema_code", ] diff --git a/src/tinybird_sdk/infer/analyze.py b/src/tinybird_sdk/infer/analyze.py new file mode 100644 index 0000000..e0d3054 --- /dev/null +++ b/src/tinybird_sdk/infer/analyze.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from ..codegen.type_mapper import clickhouse_type_to_validator +from ..codegen.utils import to_pascal_case, to_snake_case + + +@dataclass(frozen=True, slots=True) +class AnalyzedColumn: + name: str + clickhouse_type: str + present_pct: float | None = None + + +@dataclass(frozen=True, slots=True) +class AnalyzedSchema: + columns: tuple[AnalyzedColumn, ...] + format: str | None = None + dialect: dict[str, Any] = field(default_factory=dict) + + +def parse_analyze_response(response: dict[str, Any]) -> AnalyzedSchema: + """Turn a raw `/v0/analyze` API response into an `AnalyzedSchema`. + + Tolerates the two shapes the Analyze API is known to return: a + `analysis.columns` list of `{name, recommended_type}` entries, or a + `analysis.schema`/`schema` ClickHouse DDL-like string (e.g. + `"col1 Int64, col2 Nullable(String)"`) as a fallback. + """ + analysis = response.get("analysis") or {} + raw_columns = analysis.get("columns") + columns: list[AnalyzedColumn] = [] + + if isinstance(raw_columns, list) and raw_columns: + for col in raw_columns: + if not isinstance(col, dict): + continue + name = col.get("name") + ch_type = col.get("recommended_type") or col.get("type") + if not name or not ch_type: + continue + columns.append( + AnalyzedColumn( + name=str(name), + clickhouse_type=str(ch_type), + present_pct=col.get("present_pct"), + ) + ) + else: + schema_str = analysis.get("schema") or response.get("schema") + if isinstance(schema_str, str) and schema_str.strip(): + columns.extend(_parse_schema_ddl(schema_str)) + + if not columns: + raise ValueError( + "Could not find column information in the analyze response. Expected " + "'analysis.columns' (a list of {'name', 'recommended_type'} entries) or an " + "'analysis.schema'/'schema' ClickHouse DDL string." + ) + + dialect = analysis.get("dialect") or response.get("dialect") or {} + detected_format = analysis.get("format") or response.get("format") + + return AnalyzedSchema( + columns=tuple(columns), + format=detected_format if isinstance(detected_format, str) else None, + dialect=dict(dialect) if isinstance(dialect, dict) else {}, + ) + + +def _parse_schema_ddl(schema: str) -> list[AnalyzedColumn]: + columns: list[AnalyzedColumn] = [] + for part in schema.split(","): + part = part.strip() + if not part: + continue + name, _, ch_type = part.partition(" ") + name = name.strip("`") + ch_type = ch_type.strip() + if name and ch_type: + columns.append(AnalyzedColumn(name=name, clickhouse_type=ch_type)) + return columns + + +def generate_schema_code(analyzed: AnalyzedSchema, name: str) -> str: + """Generate `define_datasource(...)` source code from an analyzed schema. + + Mirrors `codegen.generate_datasource_code`'s output shape so analyzed and + reverse-codegen'd datasources read the same way. + """ + var_name = to_snake_case(name) + type_name = to_pascal_case(name) + lines: list[str] = [f"{var_name} = define_datasource({name!r}, {{", " 'schema': {"] + for column in analyzed.columns: + lines.append( + f" {column.name!r}: {clickhouse_type_to_validator(column.clickhouse_type)}," + ) + lines.append(" },") + lines.append("})") + lines.append("") + lines.append(f"{type_name}Row = dict") + return "\n".join(lines) diff --git a/tests/test_api_analyze.py b/tests/test_api_analyze.py new file mode 100644 index 0000000..8f5afe8 --- /dev/null +++ b/tests/test_api_analyze.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +import tinybird_sdk.api.api as api_module +from tinybird_sdk.api.api import TinybirdApi, TinybirdApiError + + +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 _make_api() -> TinybirdApi: + return TinybirdApi({"base_url": "https://api.tinybird.co", "token": "p.test"}) + + +def test_analyze_requires_either_url_or_file() -> None: + api = _make_api() + with pytest.raises(ValueError, match="Either 'url' or 'file'"): + api.analyze({}) + + +def test_analyze_rejects_both_url_and_file() -> None: + api = _make_api() + with pytest.raises(ValueError, match="Only one of 'url' or 'file'"): + api.analyze({"url": "https://x.y/file.csv", "file": "./file.csv"}) + + +def test_analyze_with_url(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def fake_fetch(url: str, **kwargs: Any) -> _FakeResponse: + captured["url"] = url + captured["method"] = kwargs.get("method") + return _FakeResponse(200, {"analysis": {"schema": "id Int64, name String"}}) + + monkeypatch.setattr(api_module, "tinybird_fetch", fake_fetch) + + result = _make_api().analyze({"url": "https://x.y/events.csv"}) + + assert captured["method"] == "POST" + assert captured["url"].startswith("https://api.tinybird.co/v0/analyze?") + assert "url=https" in captured["url"] + assert "format=csv" in captured["url"] + assert result["analysis"]["schema"] == "id Int64, name String" + + +def test_analyze_with_local_file(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + captured: dict[str, Any] = {} + + def fake_fetch(url: str, **kwargs: Any) -> _FakeResponse: + captured["url"] = url + captured["headers"] = kwargs.get("headers") + captured["body"] = kwargs.get("body") + return _FakeResponse( + 200, + {"analysis": {"columns": [{"name": "id", "recommended_type": "Int64"}]}}, + ) + + monkeypatch.setattr(api_module, "tinybird_fetch", fake_fetch) + + local_file = tmp_path / "events.csv" + local_file.write_text("id\n1\n", encoding="utf-8") + + result = _make_api().analyze({"file": str(local_file)}) + + assert captured["url"].endswith("/v0/analyze?format=csv") + assert captured["headers"]["Content-Type"].startswith("multipart/form-data;") + assert b"id\n1\n" in captured["body"] + assert result["analysis"]["columns"][0]["name"] == "id" + + +def test_analyze_raises_on_error_response(monkeypatch: pytest.MonkeyPatch) -> None: + def fake_fetch(_url: str, **_kwargs: Any) -> _FakeResponse: + return _FakeResponse(422, {"error": "could not analyze file"}) + + monkeypatch.setattr(api_module, "tinybird_fetch", fake_fetch) + + with pytest.raises(TinybirdApiError, match="could not analyze file"): + _make_api().analyze({"url": "https://x.y/events.csv"}) diff --git a/tests/test_infer_analyze.py b/tests/test_infer_analyze.py new file mode 100644 index 0000000..2a06fe2 --- /dev/null +++ b/tests/test_infer_analyze.py @@ -0,0 +1,50 @@ +import pytest + +from tinybird_sdk.infer import generate_schema_code, parse_analyze_response + + +def test_parse_analyze_response_from_columns_list() -> None: + response = { + "analysis": { + "columns": [ + {"name": "id", "recommended_type": "Int64", "present_pct": 100.0}, + {"name": "name", "recommended_type": "Nullable(String)"}, + ], + "format": "csv", + "dialect": {"delimiter": ",", "new_line": "\\n"}, + } + } + + analyzed = parse_analyze_response(response) + + assert [c.name for c in analyzed.columns] == ["id", "name"] + assert analyzed.columns[0].clickhouse_type == "Int64" + assert analyzed.columns[0].present_pct == 100.0 + assert analyzed.columns[1].clickhouse_type == "Nullable(String)" + assert analyzed.format == "csv" + assert analyzed.dialect == {"delimiter": ",", "new_line": "\\n"} + + +def test_parse_analyze_response_falls_back_to_schema_ddl_string() -> None: + response = {"analysis": {"schema": "id Int64, name Nullable(String), ts DateTime"}} + + analyzed = parse_analyze_response(response) + + assert [c.name for c in analyzed.columns] == ["id", "name", "ts"] + assert analyzed.columns[2].clickhouse_type == "DateTime" + + +def test_parse_analyze_response_raises_without_recognizable_columns() -> None: + with pytest.raises(ValueError, match="Could not find column information"): + parse_analyze_response({"analysis": {}}) + + +def test_generate_schema_code_emits_t_validators() -> None: + analyzed = parse_analyze_response({"analysis": {"schema": "id Int64, name Nullable(String)"}}) + + code = generate_schema_code(analyzed, "events") + + assert "events = define_datasource('events', {" in code + assert "'id': t.int64()," in code + assert "'name': t.string().nullable()," in code + assert "EventsRow = dict" in code