diff --git a/src/api/api.test.ts b/src/api/api.test.ts index 44d45cc..5800607 100644 --- a/src/api/api.test.ts +++ b/src/api/api.test.ts @@ -1225,4 +1225,141 @@ describe("TinybirdApi", () => { expect(authorizationHeader).toBe("Bearer p.override-token"); }); }); + + describe("sampleDatasource", () => { + it("sends default payload (max_files=1, full_export=false)", async () => { + let body: unknown; + let contentType: string | null = null; + + server.use( + http.post(`${BASE_URL}/v0/datasources/events/sample`, async ({ request }) => { + contentType = request.headers.get("Content-Type"); + body = await request.json(); + return HttpResponse.json({ job_id: "job-1", status: "waiting" }); + }) + ); + + const api = createTinybirdApi({ + baseUrl: BASE_URL, + token: "p.default-token", + }); + + const result = await api.sampleDatasource("events"); + + expect(result).toEqual({ job_id: "job-1", status: "waiting" }); + expect(contentType).toBe("application/json"); + expect(body).toEqual({ max_files: 1, full_export: false }); + }); + + it("forwards maxFiles for blob storage connectors", async () => { + let body: unknown; + + server.use( + http.post(`${BASE_URL}/v0/datasources/events/sample`, async ({ request }) => { + body = await request.json(); + return HttpResponse.json({ job_id: "job-1" }); + }) + ); + + const api = createTinybirdApi({ + baseUrl: BASE_URL, + token: "p.default-token", + }); + + await api.sampleDatasource("events", { maxFiles: 3 }); + + expect(body).toEqual({ max_files: 3, full_export: false }); + }); + + it("forwards DynamoDB rows option", async () => { + let body: unknown; + + server.use( + http.post(`${BASE_URL}/v0/datasources/ddb_ds/sample`, async ({ request }) => { + body = await request.json(); + return HttpResponse.json({ job_id: "job-1" }); + }) + ); + + const api = createTinybirdApi({ + baseUrl: BASE_URL, + token: "p.default-token", + }); + + await api.sampleDatasource("ddb_ds", { rows: 100000 }); + + expect(body).toEqual({ max_files: 1, full_export: false, rows: 100000 }); + }); + + it("forwards DynamoDB maxBytes option", async () => { + let body: unknown; + + server.use( + http.post(`${BASE_URL}/v0/datasources/ddb_ds/sample`, async ({ request }) => { + body = await request.json(); + return HttpResponse.json({ job_id: "job-1" }); + }) + ); + + const api = createTinybirdApi({ + baseUrl: BASE_URL, + token: "p.default-token", + }); + + await api.sampleDatasource("ddb_ds", { maxBytes: "1GB" }); + + expect(body).toEqual({ max_files: 1, full_export: false, max_bytes: "1GB" }); + }); + + it("forwards DynamoDB fullExport option", async () => { + let body: unknown; + + server.use( + http.post(`${BASE_URL}/v0/datasources/ddb_ds/sample`, async ({ request }) => { + body = await request.json(); + return HttpResponse.json({ job_id: "job-1" }); + }) + ); + + const api = createTinybirdApi({ + baseUrl: BASE_URL, + token: "p.default-token", + }); + + await api.sampleDatasource("ddb_ds", { fullExport: true }); + + expect(body).toEqual({ max_files: 1, full_export: true }); + }); + + it("throws when rows and maxBytes are both provided", async () => { + const api = createTinybirdApi({ + baseUrl: BASE_URL, + token: "p.default-token", + }); + + await expect( + api.sampleDatasource("ddb_ds", { rows: 10, maxBytes: "1GB" }) + ).rejects.toThrow("'rows' and 'maxBytes' are mutually exclusive; pass only one"); + }); + + it("allows per-request token override", async () => { + let authorizationHeader: string | null = null; + + server.use( + http.post(`${BASE_URL}/v0/datasources/events/sample`, ({ request }) => { + authorizationHeader = request.headers.get("Authorization"); + return HttpResponse.json({ job_id: "job-1" }); + }) + ); + + const api = createTinybirdApi({ + baseUrl: BASE_URL, + token: "p.default-token", + }); + + await api.sampleDatasource("events", {}, { token: "p.override-token" }); + + expect(authorizationHeader).toBe("Bearer p.override-token"); + }); + }); }); diff --git a/src/api/api.ts b/src/api/api.ts index 53b9230..bbd53a5 100644 --- a/src/api/api.ts +++ b/src/api/api.ts @@ -8,6 +8,8 @@ import type { IngestResult, QueryOptions, QueryResult, + SampleOptions, + SampleResult, TruncateOptions, TruncateResult, TinybirdErrorResponse, @@ -86,6 +88,15 @@ export interface TinybirdApiTruncateOptions extends TruncateOptions { token?: string; } +export interface TinybirdApiSampleOptions { + /** Optional token override for this request */ + token?: string; + /** Request timeout in milliseconds */ + timeout?: number; + /** AbortController signal for cancellation */ + signal?: AbortSignal; +} + /** * Scope definition for token creation APIs */ @@ -505,6 +516,61 @@ export class TinybirdApi { return this.parseOptionalJson(response); } + /** + * Start a bounded sample import job for an S3/GCS/DynamoDB connected datasource. + * + * For blob storage (S3/GCS) connectors, `maxFiles` bounds how many files are + * imported (default 1, max 10). For DynamoDB, the sample is bounded by either + * `rows` or `maxBytes` (mutually exclusive), or `fullExport` triggers a full + * PITR export of the whole table instead of a bounded scan. + */ + async sampleDatasource( + datasourceName: string, + options: SampleOptions = {}, + apiOptions: TinybirdApiSampleOptions = {} + ): Promise { + if (options.rows !== undefined && options.maxBytes !== undefined) { + throw new Error("'rows' and 'maxBytes' are mutually exclusive; pass only one"); + } + + const url = new URL( + `/v0/datasources/${encodeURIComponent(datasourceName)}/sample`, + `${this.baseUrl}/` + ); + + const payload: Record = { + max_files: options.maxFiles ?? 1, + full_export: options.fullExport ?? false, + }; + + if (options.rows !== undefined) { + payload.rows = options.rows; + } + + if (options.maxBytes !== undefined) { + payload.max_bytes = options.maxBytes; + } + + const response = await this.request(url.toString(), { + method: "POST", + token: apiOptions.token, + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(payload), + signal: this.createAbortSignal( + options.timeout ?? apiOptions.timeout, + options.signal ?? apiOptions.signal + ), + }); + + if (!response.ok) { + await this.handleErrorResponse(response); + } + + return (await response.json()) as SampleResult; + } + /** * Create a token using Tinybird Token API. * Supports both static and JWT token payloads. diff --git a/src/client/base.test.ts b/src/client/base.test.ts index 697aca9..cf8a8a2 100644 --- a/src/client/base.test.ts +++ b/src/client/base.test.ts @@ -235,7 +235,7 @@ describe("TinybirdClient", () => { expect(client.datasources).toBeDefined(); }); - it("datasources namespace has ingest/append/replace/delete/truncate methods", () => { + it("datasources namespace has ingest/append/replace/delete/truncate/sample methods", () => { const client = createClient({ baseUrl: "https://api.tinybird.co", token: "test-token", @@ -246,6 +246,7 @@ describe("TinybirdClient", () => { expect(typeof client.datasources.replace).toBe("function"); expect(typeof client.datasources.delete).toBe("function"); expect(typeof client.datasources.truncate).toBe("function"); + expect(typeof client.datasources.sample).toBe("function"); }); it("datasources conforms to DatasourcesNamespace interface", () => { @@ -261,6 +262,7 @@ describe("TinybirdClient", () => { expect(typeof datasources.replace).toBe("function"); expect(typeof datasources.delete).toBe("function"); expect(typeof datasources.truncate).toBe("function"); + expect(typeof datasources.sample).toBe("function"); }); }); }); diff --git a/src/client/base.ts b/src/client/base.ts index 637c0c8..058c355 100644 --- a/src/client/base.ts +++ b/src/client/base.ts @@ -14,6 +14,8 @@ import type { IngestResult, QueryOptions, IngestOptions, + SampleOptions, + SampleResult, TruncateOptions, TruncateResult, } from "./types.js"; @@ -113,6 +115,12 @@ export class TinybirdClient { ): Promise => { return this.truncateDatasource(datasourceName, options); }, + sample: ( + datasourceName: string, + options: SampleOptions = {} + ): Promise => { + return this.sampleDatasource(datasourceName, options); + }, }; // Initialize tokens namespace @@ -219,6 +227,44 @@ export class TinybirdClient { } } + /** + * Start a bounded sample import job for an S3/GCS/DynamoDB connected datasource + * + * For blob storage (S3/GCS) connectors, `maxFiles` bounds how many files are + * imported (default 1, max 10). For DynamoDB, the sample is bounded by either + * `rows` or `maxBytes` (mutually exclusive — passing both throws), or + * `fullExport` triggers a full PITR export of the whole table instead of a + * bounded scan. + * + * @param datasourceName - Name of the datasource + * @param options - Sample options (maxFiles, rows, maxBytes, fullExport) + * @returns Sample job result + * + * @example + * ```ts + * // Sample up to 3 files from an S3/GCS connected datasource + * await client.datasources.sample('s3_events', { maxFiles: 3 }); + * + * // Bound a DynamoDB sample by row count + * await client.datasources.sample('ddb_events', { rows: 100_000 }); + * + * // Trigger a full DynamoDB PITR export + * await client.datasources.sample('ddb_events', { fullExport: true }); + * ``` + */ + private async sampleDatasource( + datasourceName: string, + options: SampleOptions = {} + ): Promise { + const token = await this.getToken(); + + try { + return await this.getApi(token).sampleDatasource(datasourceName, options); + } catch (error) { + this.rethrowApiError(error); + } + } + /** * Ingest a single event to a datasource * diff --git a/src/client/types.ts b/src/client/types.ts index e7b093b..678c7b4 100644 --- a/src/client/types.ts +++ b/src/client/types.ts @@ -291,6 +291,54 @@ export interface TruncateResult { status?: string; } +/** + * Options for starting a bounded sample import job on an S3/GCS/DynamoDB + * connected datasource. + * + * `rows` and `maxBytes` are mutually exclusive (DynamoDB only) and only one + * of them may be provided. + */ +export interface SampleOptions { + /** + * Maximum number of files to import for blob storage (S3/GCS) connectors. + * Defaults to 1 (max 10). + */ + maxFiles?: number; + /** + * DynamoDB only: maximum number of rows to scan and import. + * Mutually exclusive with `maxBytes`. + */ + rows?: number; + /** + * DynamoDB only: maximum approximate JSONEachRow bytes to import, + * e.g. `"500MB"`. Mutually exclusive with `rows`. + */ + maxBytes?: string; + /** + * DynamoDB only: trigger a full PITR export of the whole table instead of + * a bounded scan. Defaults to false. + */ + fullExport?: boolean; + /** Request timeout in milliseconds */ + timeout?: number; + /** AbortController signal for cancellation */ + signal?: AbortSignal; +} + +/** + * Result of starting a datasource sample job + */ +export interface SampleResult { + /** Sample job ID */ + id?: string; + /** Same value as id */ + job_id?: string; + /** Job status URL */ + job_url?: string; + /** Job status */ + status?: string; +} + /** * Datasources namespace interface for raw client */ @@ -309,4 +357,6 @@ export interface DatasourcesNamespace { delete(datasourceName: string, options: DeleteOptions): Promise; /** Truncate all rows from a datasource */ truncate(datasourceName: string, options?: TruncateOptions): Promise; + /** Start a bounded sample import job for an S3/GCS/DynamoDB connected datasource */ + sample(datasourceName: string, options?: SampleOptions): Promise; } diff --git a/src/index.ts b/src/index.ts index b8bab18..2914c8e 100644 --- a/src/index.ts +++ b/src/index.ts @@ -257,6 +257,8 @@ export type { IngestResult, QueryOptions, IngestOptions, + SampleOptions, + SampleResult, TruncateOptions, TruncateResult, ColumnMeta, @@ -280,6 +282,7 @@ export type { TinybirdApiAppendOptions, TinybirdApiDeleteOptions, TinybirdApiTruncateOptions, + TinybirdApiSampleOptions, TinybirdApiRequestInit, TinybirdApiTokenScope, TinybirdApiCreateTokenRequest,