feat(knowledge-bases): add waitForDatabase polling helper

src/index.ts

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

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

@@ -34,6 +34,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);
@@ -135,6 +136,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);
+  }
 }
 
 /**
@@ -422,6 +448,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,173 @@
+// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+import { Gradient, type ClientOptions } from '../../client';
+import { GradientError } from '../../core/error';
+import { sleep } from '../../internal/utils/sleep';
+import { KnowledgeBaseRetrieveResponse } from './knowledge-bases';
+
+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;
+
+  /**
+   * Additional request options to pass through to the knowledge base retrieval request.
+   */
+  requestOptions?: import('../../internal/request-options').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();
+ * 
+ * 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');
+ *   }
+ * }
+ * ```
+ *
+ * @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
+ */
+export async function waitForDatabase(
+  client: Gradient,
+  uuid: string,
+  options: WaitForDatabaseOptions = {}
+): Promise<KnowledgeBaseRetrieveResponse> {
+  const {
+    interval = 5000,
+    timeout = 600000,
+    requestOptions,
+  } = options;
+
+  const startTime = Date.now();
+
+  while (true) {
+    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();
+    });
+  });
+});