buster/packages/data-source/src/adapters/bigquery.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest';
import { DataSourceType } from '../types/credentials';
import type { BigQueryCredentials } from '../types/credentials';
import { BigQueryAdapter } from './bigquery';

// Create mock BigQuery instance
const mockBigQuery = {
  dataset: vi.fn().mockReturnValue({
    table: vi.fn().mockReturnValue({
      query: vi.fn(),
    }),
  }),
  getDatasets: vi.fn(),
  createQueryJob: vi.fn(),
};

// Mock @google-cloud/bigquery module
vi.mock('@google-cloud/bigquery', () => ({
  BigQuery: vi.fn(() => mockBigQuery),
}));

describe('BigQueryAdapter', () => {
  let adapter: BigQueryAdapter;

  beforeEach(() => {
    vi.clearAllMocks();
    adapter = new BigQueryAdapter();
  });

  describe('initialization', () => {
    it('should initialize with project ID and service account key', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        service_account_key: JSON.stringify({
          type: 'service_account',
          project_id: 'test-project',
          private_key: 'test-key',
          client_email: 'test@test.iam.gserviceaccount.com',
        }),
      };

      await adapter.initialize(credentials);

      const { BigQuery } = await import('@google-cloud/bigquery');
      expect(BigQuery).toHaveBeenCalledWith({
        projectId: 'test-project',
        credentials: {
          type: 'service_account',
          project_id: 'test-project',
          private_key: 'test-key',
          client_email: 'test@test.iam.gserviceaccount.com',
        },
      });
    });

    it('should initialize with project ID and key file path', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        key_file_path: '/path/to/key.json',
      };

      await adapter.initialize(credentials);

      const { BigQuery } = await import('@google-cloud/bigquery');
      expect(BigQuery).toHaveBeenCalledWith({
        projectId: 'test-project',
        keyFilename: '/path/to/key.json',
      });
    });

    it('should handle invalid JSON in service account key', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        service_account_key: 'invalid json',
      };

      // BigQuery will treat invalid JSON as a file path, so initialization should succeed
      await expect(adapter.initialize(credentials)).resolves.not.toThrow();
    });

    it('should handle service account key as parsed object', async () => {
      const serviceAccountObject = {
        type: 'service_account',
        project_id: 'test-project',
        private_key_id: 'key123',
        private_key: '-----BEGIN PRIVATE KEY-----\ntest\n-----END PRIVATE KEY-----',
        client_email: 'test@test-project.iam.gserviceaccount.com',
      };

      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        service_account_key: serviceAccountObject,
      };

      await adapter.initialize(credentials);

      const { BigQuery } = await import('@google-cloud/bigquery');
      expect(BigQuery).toHaveBeenCalledWith(
        expect.objectContaining({
          projectId: 'test-project',
          credentials: serviceAccountObject,
        })
      );
    });

    it('should allow initialization without credentials (uses ADC)', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
      };

      // BigQuery can use Application Default Credentials
      await expect(adapter.initialize(credentials)).resolves.not.toThrow();
    });

    it('should set default dataset when provided', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        default_dataset: 'my-dataset',
        service_account_key: JSON.stringify({ type: 'service_account' }),
      };

      await adapter.initialize(credentials);

      const { BigQuery } = await import('@google-cloud/bigquery');
      expect(BigQuery).toHaveBeenCalledWith(
        expect.objectContaining({
          projectId: 'test-project',
        })
      );
    });

    it('should set location when provided', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        location: 'us-central1',
        service_account_key: JSON.stringify({ type: 'service_account' }),
      };

      await adapter.initialize(credentials);

      const { BigQuery } = await import('@google-cloud/bigquery');
      expect(BigQuery).toHaveBeenCalledWith(
        expect.objectContaining({
          projectId: 'test-project',
        })
      );
    });

    it('should throw error with invalid credentials type', async () => {
      const credentials = {
        type: DataSourceType.PostgreSQL,
        host: 'localhost',
        database: 'testdb',
        username: 'testuser',
        password: 'testpass',
      };

      await expect(adapter.initialize(credentials)).rejects.toThrow(
        'Invalid credentials type. Expected bigquery, got postgres'
      );
    });
  });

  describe('query execution', () => {
    const credentials: BigQueryCredentials = {
      type: DataSourceType.BigQuery,
      project_id: 'test-project',
      service_account_key: JSON.stringify({ type: 'service_account' }),
    };

    beforeEach(async () => {
      await adapter.initialize(credentials);
    });

    it('should execute simple query without parameters', async () => {
      const mockRows = [{ id: 1, name: 'Test' }];
      const mockMetadata = {
        schema: {
          fields: [
            { name: 'id', type: 'INT64', mode: 'NULLABLE' },
            { name: 'name', type: 'STRING', mode: 'NULLABLE' },
          ],
        },
      };
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([mockRows, null, mockMetadata]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      const result = await adapter.query('SELECT * FROM dataset.table');

      expect(mockBigQuery.createQueryJob).toHaveBeenCalledWith({
        query: 'SELECT * FROM dataset.table',
        location: undefined,
        jobTimeoutMs: 60000,
        useLegacySql: false,
      });

      expect(result).toEqual({
        rows: mockRows,
        rowCount: 1,
        fields: [
          { name: 'id', type: 'bigint', nullable: true, length: 0, precision: 0 },
          { name: 'name', type: 'text', nullable: true, length: 0, precision: 0 },
        ],
        hasMoreRows: false,
      });
    });

    it('should execute parameterized query', async () => {
      const mockRows = [{ id: 1, name: 'Test' }];
      const mockMetadata = {
        schema: {
          fields: [
            { name: 'id', type: 'INT64', mode: 'NULLABLE' },
            { name: 'name', type: 'STRING', mode: 'NULLABLE' },
          ],
        },
      };
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([mockRows, null, mockMetadata]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      const result = await adapter.query('SELECT * FROM dataset.table WHERE id = ?', [1]);

      expect(mockBigQuery.createQueryJob).toHaveBeenCalledWith({
        query: 'SELECT * FROM dataset.table WHERE id = @param0',
        params: { param0: 1 },
        location: undefined,
        jobTimeoutMs: 60000,
        useLegacySql: false,
      });
      expect(result.rows).toEqual(mockRows);
    });

    it('should handle maxRows limit', async () => {
      const mockRows = Array.from({ length: 10 }, (_, i) => ({ id: i + 1 }));
      const mockMetadata = {
        schema: {
          fields: [{ name: 'id', type: 'INT64', mode: 'NULLABLE' }],
        },
      };
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([mockRows, null, mockMetadata]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      const result = await adapter.query('SELECT * FROM dataset.table', [], 10);

      expect(mockBigQuery.createQueryJob).toHaveBeenCalledWith(
        expect.objectContaining({
          maxResults: 11, // Requests one extra to check for more rows
        })
      );
      expect(result.rows).toHaveLength(10);
      expect(result.hasMoreRows).toBe(false);
    });

    it('should detect when there are more rows', async () => {
      const mockRows = Array.from({ length: 11 }, (_, i) => ({ id: i + 1 }));
      const mockMetadata = {
        schema: {
          fields: [{ name: 'id', type: 'INT64', mode: 'NULLABLE' }],
        },
      };
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([mockRows, null, mockMetadata]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      const result = await adapter.query('SELECT * FROM dataset.table', [], 10);

      expect(result.rows).toHaveLength(10);
      expect(result.hasMoreRows).toBe(true);
    });

    it('should use custom timeout when provided', async () => {
      const mockMetadata = { schema: { fields: [] } };
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([[], null, mockMetadata]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      await adapter.query('SELECT 1', [], undefined, 5000);

      expect(mockBigQuery.createQueryJob).toHaveBeenCalledWith(
        expect.objectContaining({
          jobTimeoutMs: 5000,
        })
      );
    });

    it('should pass options to createQueryJob', async () => {
      const mockMetadata = { schema: { fields: [] } };
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([[], null, mockMetadata]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      await adapter.query('SELECT 1');

      expect(mockBigQuery.createQueryJob).toHaveBeenCalledWith({
        query: 'SELECT 1',
        location: undefined,
        jobTimeoutMs: 60000,
        useLegacySql: false,
      });
    });

    it('should handle query errors', async () => {
      mockBigQuery.createQueryJob.mockRejectedValueOnce(new Error('Query failed'));

      await expect(adapter.query('SELECT * FROM invalid_table')).rejects.toThrow(
        'BigQuery query failed: Query failed'
      );
    });

    it('should throw error when not connected', async () => {
      const disconnectedAdapter = new BigQueryAdapter();

      await expect(disconnectedAdapter.query('SELECT 1')).rejects.toThrow(
        'bigquery adapter is not connected. Call initialize() first.'
      );
    });

    it('should handle empty result sets', async () => {
      const mockMetadata = {
        schema: {
          fields: [{ name: 'id', type: 'INT64', mode: 'NULLABLE' }],
        },
      };
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([[], null, mockMetadata]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      const result = await adapter.query('SELECT * FROM dataset.table WHERE 1=0');

      expect(result.rows).toEqual([]);
      expect(result.rowCount).toBe(0);
      expect(result.fields).toEqual([
        { name: 'id', type: 'bigint', nullable: true, length: 0, precision: 0 },
      ]);
    });

    it('should handle results without schema metadata', async () => {
      const mockRows = [{ id: 1 }];
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([mockRows, null, {}]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      const result = await adapter.query('SELECT 1 as id');

      expect(result.fields).toEqual([]);
    });

    it('should handle repeated fields', async () => {
      const mockRows = [{ id: 1, tags: ['tag1', 'tag2'] }];
      const mockMetadata = {
        schema: {
          fields: [
            { name: 'id', type: 'INT64', mode: 'NULLABLE' },
            { name: 'tags', type: 'STRING', mode: 'REPEATED' },
          ],
        },
      };
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([mockRows, null, mockMetadata]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      const result = await adapter.query('SELECT * FROM dataset.table');

      expect(result.rows).toEqual(mockRows);
    });
  });

  describe('connection management', () => {
    it('should test connection successfully', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        service_account_key: JSON.stringify({ type: 'service_account' }),
      };

      await adapter.initialize(credentials);

      const mockMetadata = {};
      const mockJob = {
        getQueryResults: vi.fn().mockResolvedValueOnce([[], null, mockMetadata]),
      };

      mockBigQuery.createQueryJob.mockResolvedValueOnce([mockJob]);

      const result = await adapter.testConnection();

      expect(result).toBe(true);
      expect(mockBigQuery.createQueryJob).toHaveBeenCalledWith({
        query: 'SELECT 1 as test',
        useLegacySql: false,
      });
    });

    it('should return false when test connection fails', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        service_account_key: JSON.stringify({ type: 'service_account' }),
      };

      await adapter.initialize(credentials);

      mockBigQuery.createQueryJob.mockRejectedValueOnce(new Error('Connection test failed'));

      const result = await adapter.testConnection();

      expect(result).toBe(false);
    });

    it('should handle close (no-op for BigQuery)', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        service_account_key: JSON.stringify({ type: 'service_account' }),
      };

      await adapter.initialize(credentials);
      await adapter.close();

      // BigQuery doesn't require explicit closing, just verify it doesn't throw
      await expect(adapter.testConnection()).resolves.toBe(false);
    });
  });

  describe('introspection', () => {
    it('should return introspector', async () => {
      const credentials: BigQueryCredentials = {
        type: DataSourceType.BigQuery,
        project_id: 'test-project',
        service_account_key: JSON.stringify({ type: 'service_account' }),
      };

      await adapter.initialize(credentials);

      const introspector = adapter.introspect();

      // Just verify it returns an introspector with the correct interface
      expect(introspector).toBeDefined();
      expect(introspector.getDatabases).toBeDefined();
      expect(introspector.getSchemas).toBeDefined();
      expect(introspector.getTables).toBeDefined();
      expect(introspector.getColumns).toBeDefined();
    });

    it('should throw error when trying to introspect without connection', () => {
      expect(() => adapter.introspect()).toThrow(
        'bigquery adapter is not connected. Call initialize() first.'
      );
    });
  });

  describe('data source type', () => {
    it('should return correct data source type', () => {
      expect(adapter.getDataSourceType()).toBe(DataSourceType.BigQuery);
    });
  });
});