feat(knowledge-bases): add waitForDatabase polling helper (#8)

Add helper function to poll knowledge base database status until ONLINE state.
Includes error handling for failed states, timeout protection, and configurable
polling intervals. Fixes the need for manual polling loops.

Co-authored-by: Ben Batha <[email protected]>

Author: yashranaway

src/index.ts

@@ -20,6 +20,14 @@ export {
   PermissionDeniedError,
   UnprocessableEntityError,
 } from './core/error';
+
+// Export knowledge base helpers
+export {
+  waitForDatabase,
+  type WaitForDatabaseOptions,
+  WaitForDatabaseTimeoutError,
+  WaitForDatabaseFailedError,
+} from './resources/knowledge-bases/wait-for-database';
 export {
   IndexingJobAbortedError,
   IndexingJobNotFoundError,

src/resources/knowledge-bases/knowledge-bases.ts

@@ -36,6 +36,7 @@ import {
 import { APIPromise } from '../../core/api-promise';
 import { RequestOptions } from '../../internal/request-options';
 import { path } from '../../internal/utils/path';
+import { waitForDatabase } from './wait-for-database';
 
 export class KnowledgeBases extends APIResource {
   dataSources: DataSourcesAPI.DataSources = new DataSourcesAPI.DataSources(this._client);
@@ -137,6 +138,31 @@ export class KnowledgeBases extends APIResource {
       ...options,
     });
   }
+
+  /**
+   * Polls for knowledge base database creation to complete.
+   *
+   * This helper method polls the knowledge base status until the database is ONLINE,
+   * handling various error states and providing configurable timeout and polling intervals.
+   *
+   * @example
+   * ```ts
+   * const kb = await client.knowledgeBases.waitForDatabase('123e4567-e89b-12d3-a456-426614174000');
+   * console.log('Database is ready:', kb.database_status); // 'ONLINE'
+   * ```
+   *
+   * @param uuid - The knowledge base UUID to poll for
+   * @param options - Configuration options for polling behavior
+   * @returns Promise<KnowledgeBaseRetrieveResponse> - The knowledge base with ONLINE database status
+   * @throws WaitForDatabaseTimeoutError - If polling times out
+   * @throws WaitForDatabaseFailedError - If the database enters a failed state
+   */
+  async waitForDatabase(
+    uuid: string,
+    options?: import('./wait-for-database').WaitForDatabaseOptions,
+  ): Promise<KnowledgeBaseRetrieveResponse> {
+    return waitForDatabase(this._client, uuid, options || {});
+  }
 }
 
 /**
@@ -424,6 +450,13 @@ export interface KnowledgeBaseListParams {
 KnowledgeBases.DataSources = DataSources;
 KnowledgeBases.IndexingJobs = IndexingJobs;
 
+export {
+  waitForDatabase,
+  WaitForDatabaseOptions,
+  WaitForDatabaseTimeoutError,
+  WaitForDatabaseFailedError,
+} from './wait-for-database';
+
 export declare namespace KnowledgeBases {
   export {
     type APIKnowledgeBase as APIKnowledgeBase,

src/resources/knowledge-bases/wait-for-database.ts

@@ -0,0 +1,197 @@
+// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+import { Gradient } from '../../client';
+import { GradientError } from '../../core/error';
+import { sleep } from '../../internal/utils/sleep';
+import { KnowledgeBaseRetrieveResponse } from './knowledge-bases';
+import { RequestOptions } from '../../internal/request-options';
+
+export interface WaitForDatabaseOptions {
+  /**
+   * The polling interval in milliseconds. Defaults to 5000 (5 seconds).
+   */
+  interval?: number;
+
+  /**
+   * The maximum time to wait in milliseconds. Defaults to 600000 (10 minutes).
+   */
+  timeout?: number;
+
+  /**
+   * AbortSignal to cancel the polling operation.
+   */
+  signal?: AbortSignal;
+
+  /**
+   * Additional request options to pass through to the knowledge base retrieval request.
+   */
+  requestOptions?: RequestOptions;
+}
+
+/**
+ * Database status values that indicate a successful deployment.
+ */
+const ONLINE_STATUSES = ['ONLINE'] as const;
+
+/**
+ * Database status values that indicate a failed deployment.
+ */
+const FAILED_STATUSES = ['DECOMMISSIONED', 'UNHEALTHY'] as const;
+
+/**
+ * Database status values that indicate the deployment is still in progress.
+ */
+const PENDING_STATUSES = [
+  'CREATING',
+  'POWEROFF',
+  'REBUILDING',
+  'REBALANCING',
+  'FORKING',
+  'MIGRATING',
+  'RESIZING',
+  'RESTORING',
+  'POWERING_ON',
+] as const;
+
+export class WaitForDatabaseTimeoutError extends GradientError {
+  constructor(message: string, kbId?: string, timeout?: number) {
+    super(message);
+    this.name = 'WaitForDatabaseTimeoutError';
+    if (kbId) {
+      (this as any).knowledgeBaseId = kbId;
+      (this as any).timeout = timeout;
+    }
+  }
+}
+
+export class WaitForDatabaseFailedError extends GradientError {
+  constructor(message: string, kbId?: string, status?: string) {
+    super(message);
+    this.name = 'WaitForDatabaseFailedError';
+    if (kbId) {
+      (this as any).knowledgeBaseId = kbId;
+      (this as any).databaseStatus = status;
+    }
+  }
+}
+
+/**
+ * Polls for knowledge base database creation to complete.
+ *
+ * This helper function polls the knowledge base status until the database is ONLINE,
+ * handling various error states and providing configurable timeout and polling intervals.
+ *
+ * @example
+ * ```ts
+ * import Gradient from '@digitalocean/gradient';
+ *
+ * const client = new Gradient();
+ *
+ * // Basic usage
+ * try {
+ *   const kb = await client.knowledgeBases.waitForDatabase('123e4567-e89b-12d3-a456-426614174000');
+ *   console.log('Database is ready:', kb.database_status); // 'ONLINE'
+ * } catch (error) {
+ *   if (error instanceof WaitForDatabaseTimeoutError) {
+ *     console.log('Polling timed out');
+ *   } else if (error instanceof WaitForDatabaseFailedError) {
+ *     console.log('Database deployment failed');
+ *   }
+ * }
+ *
+ * // With AbortSignal
+ * const controller = new AbortController();
+ * setTimeout(() => controller.abort(), 30000); // Cancel after 30 seconds
+ *
+ * try {
+ *   const kb = await client.knowledgeBases.waitForDatabase('123e4567-e89b-12d3-a456-426614174000', {
+ *     signal: controller.signal
+ *   });
+ * } catch (error) {
+ *   if (error.message === 'Operation was aborted') {
+ *     console.log('Operation was cancelled');
+ *   }
+ * }
+ * ```
+ *
+ * @param client - The Gradient client instance
+ * @param uuid - The knowledge base UUID to poll for
+ * @param options - Configuration options for polling behavior
+ * @returns Promise<KnowledgeBaseRetrieveResponse> - The knowledge base with ONLINE database status
+ * @throws WaitForDatabaseTimeoutError - If polling times out
+ * @throws WaitForDatabaseFailedError - If the database enters a failed state
+ * @throws Error - If the operation is aborted via AbortSignal
+ */
+export async function waitForDatabase(
+  client: Gradient,
+  uuid: string,
+  options: WaitForDatabaseOptions,
+): Promise<KnowledgeBaseRetrieveResponse> {
+  const { interval = 5000, timeout = 600000, signal, requestOptions } = options;
+
+  const startTime = Date.now();
+
+  while (true) {
+    // Check if operation was aborted
+    if (signal?.aborted) {
+      throw new Error('Operation was aborted');
+    }
+
+    const elapsed = Date.now() - startTime;
+
+    if (elapsed > timeout) {
+      throw new WaitForDatabaseTimeoutError(
+        `Knowledge base database ${uuid} did not become ONLINE within ${timeout}ms`,
+        uuid,
+        timeout,
+      );
+    }
+
+    try {
+      const response = await client.knowledgeBases.retrieve(uuid, requestOptions);
+      const status = response.database_status;
+
+      if (!status) {
+        // If database_status is not present, continue polling
+        await sleep(interval);
+        continue;
+      }
+
+      // Check for successful completion
+      if (ONLINE_STATUSES.includes(status as any)) {
+        return response;
+      }
+
+      // Check for failed states
+      if (FAILED_STATUSES.includes(status as any)) {
+        throw new WaitForDatabaseFailedError(
+          `Knowledge base database ${uuid} entered failed state: ${status}`,
+          uuid,
+          status,
+        );
+      }
+
+      // Check if still in progress
+      if (PENDING_STATUSES.includes(status as any)) {
+        await sleep(interval);
+        continue;
+      }
+
+      // Unknown status - treat as error for safety
+      throw new WaitForDatabaseFailedError(
+        `Knowledge base database ${uuid} entered unknown state: ${status}`,
+        uuid,
+        status,
+      );
+    } catch (error) {
+      // If it's our custom error, re-throw it
+      if (error instanceof WaitForDatabaseFailedError || error instanceof WaitForDatabaseTimeoutError) {
+        throw error;
+      }
+
+      // For other errors (network issues, etc.), try waiting a bit longer before retrying
+      await sleep(Math.min(interval * 2, 30000)); // Max 30 seconds between retries
+      continue;
+    }
+  }
+}

tests/resources/knowledge-bases/wait-for-database.test.ts

@@ -0,0 +1,81 @@
+// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+import Gradient from '@digitalocean/gradient';
+import {
+  WaitForDatabaseFailedError,
+  WaitForDatabaseTimeoutError,
+  waitForDatabase,
+} from '../../../src/resources/knowledge-bases/wait-for-database';
+
+const client = new Gradient({
+  accessToken: 'My Access Token',
+  modelAccessKey: 'My Model Access Key',
+  agentAccessKey: 'My Agent Access Key',
+  baseURL: process.env['TEST_API_BASE_URL'] ?? 'http://127.0.0.1:4010',
+});
+
+describe('waitForDatabase', () => {
+  const kbUuid = '123e4567-e89b-12d3-a456-426614174000';
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+  });
+
+  describe('error classes', () => {
+    it('should create WaitForDatabaseTimeoutError with correct properties', () => {
+      const error = new WaitForDatabaseTimeoutError('Test timeout', kbUuid, 1000);
+
+      expect(error.name).toBe('WaitForDatabaseTimeoutError');
+      expect(error.message).toBe('Test timeout');
+      expect(error).toBeInstanceOf(Error);
+    });
+
+    it('should create WaitForDatabaseFailedError with correct properties', () => {
+      const error = new WaitForDatabaseFailedError('Test failure', kbUuid, 'DECOMMISSIONED');
+
+      expect(error.name).toBe('WaitForDatabaseFailedError');
+      expect(error.message).toBe('Test failure');
+      expect(error).toBeInstanceOf(Error);
+    });
+  });
+
+  describe('function parameters', () => {
+    it('should accept correct parameters', () => {
+      expect(typeof waitForDatabase).toBe('function');
+      expect(waitForDatabase.length).toBe(3);
+    });
+
+    it('should use default options when none provided', () => {
+      const options = {};
+      expect(waitForDatabase).toBeDefined();
+      // Function should exist and be callable (will fail at runtime due to mocking, but should compile)
+    });
+  });
+
+  describe('status constants', () => {
+    it('should handle different database status values', () => {
+      // Test that status strings are handled correctly
+      const statuses = [
+        'ONLINE',
+        'CREATING',
+        'REBUILDING',
+        'RESIZING',
+        'POWERING_ON',
+        'DECOMMISSIONED',
+        'UNHEALTHY',
+      ];
+
+      statuses.forEach((status) => {
+        expect(typeof status).toBe('string');
+      });
+    });
+  });
+
+  describe('integration test placeholder', () => {
+    it('should integrate with knowledge base retrieval', async () => {
+      // This is a placeholder test - actual integration tests would require a running mock server
+      expect(client.knowledgeBases).toBeDefined();
+      expect(client.knowledgeBases.retrieve).toBeDefined();
+    });
+  });
+});