From 8efe05e0ca2e09cfde0ac2d496ec813e4262bb67 Mon Sep 17 00:00:00 2001
From: Jon <jon@wearesupereffective.com>
Date: Sat, 22 Jun 2024 13:28:39 -0400
Subject: [PATCH] feat(js): add support for listing, reading, and updating
 dataset versions

---
 js/src/client.ts  | 122 ++++++++++++++++++++++++++++++++++++++++++++++
 js/src/schemas.ts |   7 +++
 2 files changed, 129 insertions(+)

diff --git a/js/src/client.ts b/js/src/client.ts
index 15be0d209..74590fc86 100644
--- a/js/src/client.ts
+++ b/js/src/client.ts
@@ -7,6 +7,7 @@ import {
   Dataset,
   DatasetDiffInfo,
   DatasetShareSchema,
+  DatasetVersion,
   Example,
   ExampleCreate,
   ExampleUpdate,
@@ -1991,6 +1992,127 @@ export class Client {
     return (await response.json()) as Dataset;
   }
 
+  /**
+   * Get a list of versions of a dataset
+   * @param props The dataset and pagination details
+   * @returns Async iterable list of dataset versions
+   */
+  public async *listDatasetVersions({
+    datasetId,
+    datasetName,
+    search,
+    limit = 100,
+    offset = 0,
+  }: {
+    datasetId?: string;
+    datasetName?: string;
+    search?: string;
+    limit?: number;
+    offset?: number;
+  }): AsyncIterable<DatasetVersion> {
+    if (!datasetId && !datasetName) {
+      throw new Error("Must provide either datasetName or datasetId");
+    }
+    const _datasetId =
+      datasetId ?? (await this.readDataset({ datasetName })).id;
+    assertUuid(_datasetId);
+
+    const path = `/datasets/${_datasetId}/versions`;
+    const params = new URLSearchParams({
+      limit: limit.toString(),
+      offset: offset.toString(),
+    });
+    if (search !== undefined) {
+      params.append("search", search);
+    }
+
+    for await (const datasets of this._getPaginated<DatasetVersion>(
+      path,
+      params
+    )) {
+      yield* datasets;
+    }
+  }
+
+  /**
+   * Get a single dataset version
+   * @param props The dataset version details
+   * @returns The dataset version
+   */
+  public async readDatasetVersion(props: {
+    datasetId?: string;
+    datasetName?: string;
+    asOf?: Date | string;
+    tag?: string;
+  }): Promise<DatasetVersion> {
+    const { datasetId, datasetName, asOf, tag } = props;
+    if (!datasetId && !datasetName) {
+      throw new Error("Must provide either datasetName or datasetId");
+    }
+    if ((asOf && tag) || (asOf && tag)) {
+      throw new Error("Must specify either asOf or tag (but not both)");
+    }
+    const _datasetId =
+      datasetId ?? (await this.readDataset({ datasetName })).id;
+    assertUuid(_datasetId);
+
+    const params = new URLSearchParams();
+    if (asOf) {
+      params.append(
+        "as_of",
+        typeof asOf === "string" ? asOf : asOf.toISOString()
+      );
+    }
+    if (tag) {
+      params.append("tag", tag);
+    }
+
+    const path = `/datasets/${_datasetId}/version`;
+    return await this._get<DatasetVersion>(path, params);
+  }
+
+  /**
+   * Set a tag on a dataset version.
+   *
+   * If the tag is already assigned to a different version of this dataset,
+   * the tag will be moved to the new version.
+   * @param props The dataset version details
+   */
+  public async updateDatasetTag(props: {
+    datasetId?: string;
+    datasetName?: string;
+    asOf: Date | string;
+    tag: string;
+  }): Promise<void> {
+    const { datasetId, datasetName, asOf, tag } = props;
+    if (!datasetId && !datasetName) {
+      throw new Error("Must provide either datasetName or datasetId");
+    }
+    const _datasetId =
+      datasetId ?? (await this.readDataset({ datasetName })).id;
+    assertUuid(_datasetId);
+
+    const response = await this.caller.call(
+      fetch,
+      `${this.apiUrl}/datasets/${_datasetId}/tags`,
+      {
+        method: "PUT",
+        headers: { ...this.headers, "Content-Type": "application/json" },
+        body: JSON.stringify({
+          as_of: typeof asOf === "string" ? asOf : asOf.toISOString(),
+          tag,
+        }),
+        signal: AbortSignal.timeout(this.timeout_ms),
+        ...this.fetchOptions,
+      }
+    );
+    if (!response.ok) {
+      throw new Error(
+        `Failed to update dataset ${_datasetId}: ${response.status} ${response.statusText}`
+      );
+    }
+  }
+
   public async deleteDataset({
     datasetId,
     datasetName,
diff --git a/js/src/schemas.ts b/js/src/schemas.ts
index 6cba693ef..bf2a40f3e 100644
--- a/js/src/schemas.ts
+++ b/js/src/schemas.ts
@@ -268,6 +268,13 @@ export interface DatasetShareSchema {
   url: string;
 }
 
+export interface DatasetVersion {
+  /** A list of tags for the version */
+  tags: string[],
+  /** The ISO date of the version */
+  as_of: string;
+}
+
 export interface FeedbackSourceBase {
   type: string;
   metadata?: KVMap;