buster/packages/data-source/tests/integration/data-source.test.ts

import { afterEach, describe, expect, it } from 'vitest';
import { DataSource, QueryRouter } from '../../src/data-source';
import type { DataSourceConfig } from '../../src/data-source';
import { DataSourceType } from '../../src/types/credentials';
import type { MySQLCredentials, PostgreSQLCredentials } from '../../src/types/credentials';
import { TEST_TIMEOUT, hasCredentials, testConfig } from '../setup';

// Helper function to create PostgreSQL credentials with proper validation
function createPostgreSQLCredentials(): PostgreSQLCredentials {
  if (
    !testConfig.postgresql.database ||
    !testConfig.postgresql.username ||
    !testConfig.postgresql.password
  ) {
    throw new Error(
      'TEST_POSTGRES_DATABASE, TEST_POSTGRES_USERNAME, and TEST_POSTGRES_PASSWORD are required for this test'
    );
  }

  return {
    type: DataSourceType.PostgreSQL,
    host: testConfig.postgresql.host,
    port: testConfig.postgresql.port,
    database: testConfig.postgresql.database,
    username: testConfig.postgresql.username,
    password: testConfig.postgresql.password,
    schema: testConfig.postgresql.schema,
    ssl: testConfig.postgresql.ssl,
  };
}

// Helper function to create MySQL credentials with proper validation
function createMySQLCredentials(): MySQLCredentials {
  if (!testConfig.mysql.database || !testConfig.mysql.username || !testConfig.mysql.password) {
    throw new Error(
      'TEST_MYSQL_DATABASE, TEST_MYSQL_USERNAME, and TEST_MYSQL_PASSWORD are required for this test'
    );
  }

  return {
    type: DataSourceType.MySQL,
    host: testConfig.mysql.host,
    port: testConfig.mysql.port,
    database: testConfig.mysql.database,
    username: testConfig.mysql.username,
    password: testConfig.mysql.password,
    ssl: testConfig.mysql.ssl,
  };
}

describe('DataSource Integration', () => {
  let dataSource: DataSource;

  afterEach(async () => {
    if (dataSource) {
      await dataSource.close();
    }
  });

  describe('single data source configuration', () => {
    it('should initialize with PostgreSQL data source', async () => {
      if (!hasCredentials('postgresql')) {
        return; // Skip if no credentials
      }

      const dataSources: DataSourceConfig[] = [
        {
          name: 'test-postgres',
          type: DataSourceType.PostgreSQL,
          credentials: createPostgreSQLCredentials(),
        },
      ];

      dataSource = new DataSource({ dataSources });

      const dataSourceNames = dataSource.getDataSources();
      expect(dataSourceNames).toEqual(['test-postgres']);
    });

    it(
      'should execute query on default data source',
      async () => {
        if (!hasCredentials('postgresql')) {
          return; // Skip if no credentials
        }

        const dataSources: DataSourceConfig[] = [
          {
            name: 'test-postgres',
            type: DataSourceType.PostgreSQL,
            credentials: createPostgreSQLCredentials(),
          },
        ];

        dataSource = new DataSource({ dataSources });

        const result = await dataSource.execute({
          sql: "SELECT 1 as test_value, 'hello' as message",
        });

        expect(result.success).toBe(true);
        expect(result.rows).toHaveLength(1);
        expect(result.rows[0]).toEqual({ test_value: 1, message: 'hello' });
        expect(result.warehouse).toBe('test-postgres');
      },
      TEST_TIMEOUT
    );
  });

  describe('multiple data source configuration', () => {
    it('should initialize with multiple data sources', async () => {
      const dataSources: DataSourceConfig[] = [];

      if (hasCredentials('postgresql')) {
        dataSources.push({
          name: 'test-postgres',
          type: DataSourceType.PostgreSQL,
          credentials: createPostgreSQLCredentials(),
        });
      }

      if (hasCredentials('mysql')) {
        dataSources.push({
          name: 'test-mysql',
          type: DataSourceType.MySQL,
          credentials: createMySQLCredentials(),
        });
      }

      if (dataSources.length === 0) {
        return; // Skip if no credentials available
      }

      const firstDataSource = dataSources[0];
      if (!firstDataSource) {
        throw new Error('Expected at least one data source');
      }

      dataSource = new DataSource({
        dataSources,
        defaultDataSource: firstDataSource.name,
      });

      const dataSourceNames = dataSource.getDataSources();
      expect(dataSourceNames).toHaveLength(dataSources.length);
    });

    it(
      'should route query to specific data source',
      async () => {
        if (!hasCredentials('postgresql')) {
          return; // Skip if no credentials
        }

        const dataSources: DataSourceConfig[] = [
          {
            name: 'test-postgres',
            type: DataSourceType.PostgreSQL,
            credentials: createPostgreSQLCredentials(),
          },
        ];

        dataSource = new DataSource({ dataSources });

        const result = await dataSource.execute({
          sql: 'SELECT 1 as test_value',
          warehouse: 'test-postgres',
        });

        expect(result.success).toBe(true);
        expect(result.warehouse).toBe('test-postgres');
      },
      TEST_TIMEOUT
    );
  });

  describe('data source management', () => {
    it(
      'should add data source dynamically',
      async () => {
        if (!hasCredentials('postgresql')) {
          return; // Skip if no credentials
        }

        dataSource = new DataSource({ dataSources: [] });

        await dataSource.addDataSource({
          name: 'dynamic-postgres',
          type: DataSourceType.PostgreSQL,
          credentials: createPostgreSQLCredentials(),
        });

        const dataSourceNames = dataSource.getDataSources();
        expect(dataSourceNames).toContain('dynamic-postgres');
      },
      TEST_TIMEOUT
    );

    it('should remove data source', async () => {
      if (!hasCredentials('postgresql')) {
        return; // Skip if no credentials
      }

      const dataSources: DataSourceConfig[] = [
        {
          name: 'test-postgres',
          type: DataSourceType.PostgreSQL,
          credentials: createPostgreSQLCredentials(),
        },
      ];

      dataSource = new DataSource({ dataSources });

      expect(dataSource.getDataSources()).toContain('test-postgres');

      await dataSource.removeDataSource('test-postgres');

      expect(dataSource.getDataSources()).not.toContain('test-postgres');
    });

    it(
      'should test all data source connections',
      async () => {
        if (!hasCredentials('postgresql')) {
          return; // Skip if no credentials
        }

        const dataSources: DataSourceConfig[] = [
          {
            name: 'test-postgres',
            type: DataSourceType.PostgreSQL,
            credentials: createPostgreSQLCredentials(),
          },
        ];

        dataSource = new DataSource({ dataSources });

        const results = await dataSource.testAllDataSources();

        expect(results).toHaveProperty('test-postgres');
        expect(results['test-postgres']).toBe(true);
      },
      TEST_TIMEOUT
    );
  });

  describe('introspection capabilities', () => {
    it(
      'should get databases from data source',
      async () => {
        if (!hasCredentials('postgresql')) {
          return; // Skip if no credentials
        }

        const dataSources: DataSourceConfig[] = [
          {
            name: 'test-postgres',
            type: DataSourceType.PostgreSQL,
            credentials: createPostgreSQLCredentials(),
          },
        ];

        dataSource = new DataSource({ dataSources });

        const databases = await dataSource.getDatabases('test-postgres');
        expect(Array.isArray(databases)).toBe(true);
        // PostgreSQL should have at least the test database
        expect(databases.length).toBeGreaterThan(0);
      },
      TEST_TIMEOUT
    );

    it(
      'should get schemas from data source',
      async () => {
        if (!hasCredentials('postgresql')) {
          return; // Skip if no credentials
        }

        const dataSources: DataSourceConfig[] = [
          {
            name: 'test-postgres',
            type: DataSourceType.PostgreSQL,
            credentials: createPostgreSQLCredentials(),
          },
        ];

        dataSource = new DataSource({ dataSources });

        const schemas = await dataSource.getSchemas('test-postgres');
        expect(Array.isArray(schemas)).toBe(true);
        // PostgreSQL should have at least the public schema
        expect(schemas.length).toBeGreaterThan(0);
      },
      TEST_TIMEOUT
    );

    it(
      'should get introspector instance',
      async () => {
        if (!hasCredentials('postgresql')) {
          return; // Skip if no credentials
        }

        const dataSources: DataSourceConfig[] = [
          {
            name: 'test-postgres',
            type: DataSourceType.PostgreSQL,
            credentials: createPostgreSQLCredentials(),
          },
        ];

        dataSource = new DataSource({ dataSources });

        const introspector = await dataSource.introspect('test-postgres');
        expect(introspector.getDataSourceType()).toBe(DataSourceType.PostgreSQL);
      },
      TEST_TIMEOUT
    );
  });

  describe('error handling', () => {
    it(
      'should handle query errors gracefully',
      async () => {
        if (!hasCredentials('postgresql')) {
          return; // Skip if no credentials
        }

        const dataSources: DataSourceConfig[] = [
          {
            name: 'test-postgres',
            type: DataSourceType.PostgreSQL,
            credentials: createPostgreSQLCredentials(),
          },
        ];

        dataSource = new DataSource({ dataSources });

        const result = await dataSource.execute({
          sql: 'SELECT * FROM non_existent_table',
        });

        expect(result.success).toBe(false);
        expect(result.error).toBeDefined();
        expect(result.error?.code).toBe('QUERY_EXECUTION_ERROR');
      },
      TEST_TIMEOUT
    );

    it('should throw error for non-existent data source', async () => {
      dataSource = new DataSource({ dataSources: [] });

      await expect(
        dataSource.execute({
          sql: 'SELECT 1',
          warehouse: 'non-existent',
        })
      ).rejects.toThrow("Specified data source 'non-existent' not found");
    });
  });
});

// Test backward compatibility with QueryRouter alias
describe('QueryRouter Backward Compatibility', () => {
  let router: DataSource;

  afterEach(async () => {
    if (router) {
      await router.close();
    }
  });

  it('should work with QueryRouter alias', async () => {
    if (!hasCredentials('postgresql')) {
      return; // Skip if no credentials
    }

    const dataSources: DataSourceConfig[] = [
      {
        name: 'test-postgres',
        type: DataSourceType.PostgreSQL,
        credentials: createPostgreSQLCredentials(),
      },
    ];

    router = new QueryRouter({ dataSources });

    const dataSourceNames = router.getDataSources();
    expect(dataSourceNames).toEqual(['test-postgres']);
  });
});