Function setupBigQuery

packages/bigquery/src/client.ts:38–115

({
  dataset,
  logger,
}: {
  dataset?: string
  logger: Logger
})

}

export async function setupBigQuery({
  dataset,
  logger,
}: {
  dataset?: string
  logger: Logger
}) {
  if (client) {
    return
  }
  const resolvedDataset = dataset ?? DATASET
  try {
    client = new BigQuery()

    // Ensure dataset exists
    const [ds] = await client.dataset(resolvedDataset).get({ autoCreate: true })

    // Ensure tables exist
    await ds.table(TRACES_TABLE).get({
      autoCreate: true,
      schema: TRACES_SCHEMA,
      timePartitioning: {
        type: 'MONTH',
        field: 'created_at',
      },
      clustering: {
        fields: ['user_id', 'agent_step_id'],
      },
    })
    await ds.table(RELABELS_TABLE).get({
      autoCreate: true,
      schema: RELABELS_SCHEMA,
      timePartitioning: {
        type: 'MONTH',
        field: 'created_at',
      },
      clustering: {
        fields: ['user_id', 'agent_step_id'],
      },
    })
    await ds.table(MESSAGE_TABLE).get({
      autoCreate: true,
      schema: MESSAGE_SCHEMA,
      timePartitioning: {
        type: 'MONTH',
        field: 'finished_at',
      },
      clustering: {
        fields: ['user_id'],
      },
    })
    await ds.table(CHAT_COMPLETION_TRACES_TABLE).get({
      autoCreate: true,
      schema: CHAT_COMPLETION_TRACES_SCHEMA,
      timePartitioning: {
        type: 'MONTH',
        field: 'created_at',
      },