Effect.tsはストリームを扱う目的でも便利だよ、という話 with BigQuery Storage Write API

import { Chunk, Console, Effect, Stream } from "effect";

function* naturals(): Iterablebigint> {
  let n = 0n;
  while (true) {
    yield n++;
  }
}
const stream = Stream.fromIterable(naturals()).pipe(
  Stream.filter((n) => n % 2n === 0n),
  Stream.take(10),
  Stream.map((n) => n * n),
);

stream.pipe(
  Stream.runCollect,
  Effect.map(Chunk.toArray),
  Effect.andThen(Console.log), 
  Effect.runPromise,
);

import { Chunk, Console, Effect, Stream } from "effect";

const readableStream = new ReadableStreamstring>({
  start(controller) {
    controller.enqueue("Hello, ");
    controller.enqueue("world!");
    controller.close();
  },
});
const streamFromReadable = Stream.fromReadableStream({
  evaluate: () => readableStream,
  onError: (error) => new Error(`Stream error: ${error}`),
});

streamFromReadable.pipe(
  Stream.runCollect,
  Effect.map(Chunk.toArray),
  Effect.andThen(Console.log), 
  Effect.runPromise,
);

import { Sink, Stream, Data, Effect, Console, Chunk } from "effect";


const sink = Sink.fold(
  Data.struct({ last: false, count: 0 }),
  ({ last }) => !last,
  ({ count }, e: boolean) =>
    Data.struct({
      last: e,
      count: count + 1,
    }),
).pipe(Sink.map(({ last, count }) => (last ? count - 1 : count)));

const stream = Stream.fromIterable([
  false,
  false,
  false,
  true,
  false,
  false,
  true,
  false,
  false,
  false,
  false,
  false,
  false,
]);

stream.pipe(
  Stream.transduce(sink),
  Stream.runCollect,
  Effect.andThen(Chunk.toArray),
  Effect.andThen(Console.log), 
  Effect.runPromiseExit,
);

import type { Table } from '@google-cloud/bigquery';
import { adapt, managedwriter } from '@google-cloud/bigquery-storage';
import type { google } from '@google-cloud/bigquery-storage/build/protos/protos.js';
import {
  Chunk,
  Config,
  Data,
  Effect,
  Fiber,
  Option,
  pipe,
  Schema,
  Sink,
  Stream,
} from 'effect';
import { Status } from 'google-gax';

export class StreamSetupError extends Data.TaggedError('StreamSetupError'){
  message: string;
}> {
  toString() {
    return `${this.name}: ${this.message}`;
  }
}
export class AppendRowsGoogleInternalError extends Data.TaggedError(
  'AppendRowsGoogleInternalError',
){
  readonly googleInternalRpcStatus: google.rpc.IStatus;
}> {
  toString() {
    return `${this.name}: ${this.googleInternalRpcStatus.message} (code: ${this.googleInternalRpcStatus.code})`;
  }
}
export class AppendRowsError extends Data.TaggedError(
  'AppendRowsGoogleInternalError',
){
  readonly rowErrors: google.cloud.bigquery.storage.v1.IRowError[];
}> {
  toString() {
    return `${this.name}: ${this.rowErrors
      .slice(0, 10)
      .map((e) => `${String(e.index)}: ${e.message}`)
      .join(', ')}`;
  }
}


export const BQ_APPEND_ROWS_MAX_BYTES = 10 * 1024 * 1024;
const bqAppendRowsCapBytesConfig = Schema.Config(
  'BQ_APPEND_ROWS_CAP_BYTES',
  Schema.NumberFromString.pipe(
    Schema.int(),
    Schema.between(0, BQ_APPEND_ROWS_MAX_BYTES),
  ).annotations({
    title: 'BigQuery (StorageWrite API) / Append Rows / Cap Bytes',
    description:
      'BigQuery Storage Write APIのAppend Rowsで1回に送信できる最大バイト数',
  }),
).pipe(Config.withDefault( 3 * 1024 * 1024));
const bqAppendRowsBufferBytesConfig = Schema.Config(
  'BQ_APPEND_ROWS_BUFFER_BYTES',
  Schema.NumberFromString.pipe(
    Schema.int(),
    Schema.between(0, BQ_APPEND_ROWS_MAX_BYTES),
  ).annotations({
    title: 'BigQuery (StorageWrite API) / Append Rows / Cap Bytes',
    description:
      'BigQuery Storage Write APIのAppend Rowsの合算リミットに対しもたせておく余裕サイズ',
  }),
).pipe(Config.withDefault( 1 * 1024 * 1024));
const bqAppendRowsOverheadBytesConfig = Schema.Config(
  'BQ_APPEND_ROWS_OVERHEAD_BYTES',
  Schema.NumberFromString.pipe(
    Schema.int(),
    Schema.between(0, BQ_APPEND_ROWS_MAX_BYTES),
  ).annotations({
    title:
      'BigQuery (StorageWrite API) / Append Rows / Overhead Bytes',
    description:
      'BigQuery Storage Write APIのAppend Rowsで1回の送信にかかると想定するオーバーヘッド',
  }),
).pipe(Config.withDefault( 3 * 1024));

const handleAppendResult = ({
  appendResult,
  error,
  rowErrors,
}: google.cloud.bigquery.storage.v1.IAppendRowsResponse) =>
  Effect.gen(function* () {
    if (error != null) {
      
      
      if (error.code === Status.ALREADY_EXISTS) {
        return Option.fromNullable(appendResult);
      }
      return yield* Effect.fail(
        new AppendRowsGoogleInternalError({ googleInternalRpcStatus: error }),
      );
    }

    if (rowErrors != null && rowErrors.length > 0) {
      return yield* Effect.fail(new AppendRowsError({ rowErrors }));
    }
    return Option.fromNullable(appendResult);
  });

export const estimateBigQueryValueBytes = (
  value: AppendRowsInput | AppendRowsInputItem | AppendRowsInputItemValue,
): number => {
  
  return new Blob([JSON.stringify(value)]).size;
};

export const writeStream = (
  table: Table,
  stream: Stream.StreamAppendRowsInputItem>,
) =>
  pipe(
    Effect.gen(function* () {
      const writeClient = new managedwriter.WriterClient();
      const projectId = table.dataset.bigQuery.projectId;
      const datasetId = table.dataset.id;
      const tableId = table.id;
      const destinationTable = `projects/${projectId}/datasets/${datasetId}/tables/${tableId}`;
      yield* Effect.log(`Destination Table: ${destinationTable}`);
      const writeStream = yield* Effect.tryPromise(() =>
        writeClient.createWriteStreamFullResponse({
          streamType: managedwriter.PendingStream,
          destinationTable,
        }),
      ).pipe(
        Effect.mapError(
          (e) => new StreamSetupError({ message: String(e.error) }),
        ),
      );

      const tableSchema = writeStream.tableSchema;
      if (tableSchema == null)
        return yield* Effect.fail(
          new StreamSetupError({
            message: `テーブルスキーマ(tableSchema)が取得できませんでした。テーブルID: ${tableId}`,
          }),
        );
      const protoDescriptor = adapt.convertStorageSchemaToProto2Descriptor(
        tableSchema,
        'root',
      );

      const writeStreamName = writeStream.name;
      if (writeStreamName == null)
        return yield* Effect.fail(
          new StreamSetupError({
            message: `テーブルID '${tableId}' のwriteStream.nameが取得できませんでした`,
          }),
        );

      yield* Effect.log('writeStreamのセットアップ完了', {
        writeStreamName,
      });

      const connection = yield* Effect.tryPromise(() =>
        writeClient.createStreamConnection({
          streamId: writeStreamName,
        }),
      ).pipe(
        Effect.mapError((e) => new StreamSetupError({ message: e.message })),
      );

      const writer = new managedwriter.JSONWriter({
        connection,
        protoDescriptor,
      });
      yield* Effect.addFinalizer((exit) =>
        Effect.gen(function* () {
          yield* Effect.log('StorageWriteApiのストリームを閉じます', {
            exit: exit.toString(),
            writeStreamName,
          });
          writer.close();
        }),
      );

      yield* Effect.log(
        `セットアップ完了: streamId=${connection.getStreamId()}`,
      );

      
      
      const fiber = yield* Effect.fork(
        Effect.gen(function* () {
          while (true) {
            yield* Effect.async((resume) => {
              connection.once('error', (...args) => {
                resume(
                  Effect.log(`Ignoring connection error`, {
                    args,
                  }),
                );
              });
            });
          }
        }),
      );
      yield* Effect.addFinalizer(() => Fiber.interrupt(fiber));

      return [
        connection,
        writer,
        writeClient,
        destinationTable,
        writeStreamName,
      ] as const;
    }),
    Effect.andThen(
      ([connection, writer, writeClient, destinationTable, streamId]) =>
        Effect.gen(function* () {
          const bqAppendRowsCapBytes = yield* bqAppendRowsCapBytesConfig;
          const bqAppendRowsBufferBytes = yield* bqAppendRowsBufferBytesConfig;
          const bqAppendRowsOverheadBytes =
            yield* bqAppendRowsOverheadBytesConfig;

          const semaphorePermits =
            BQ_APPEND_ROWS_MAX_BYTES - bqAppendRowsBufferBytes;
          const semaphore = yield* Effect.makeSemaphore(semaphorePermits);
          const scope = yield* Effect.scope;

          yield* stream.pipe(
            
            Stream.map((value) =>
              Data.struct({
                value,
                dataBytes: estimateBigQueryValueBytes(value),
              }),
            ),

            
            
            Stream.transduce(
              Sink.foldWeighted({
                initial: Data.struct({
                  cost: bqAppendRowsOverheadBytes,
                  chunk: Chunk.emptyStream.Stream.Successtypeof stream>>(),
                }),
                maxCost: bqAppendRowsCapBytes,
                cost: (_acc, { dataBytes }) => dataBytes,
                body: (acc, { value, dataBytes }) =>
                  Data.struct({
                    cost: acc.cost + dataBytes,
                    chunk: Chunk.append(acc.chunk, value),
                  }),
              }),
            ),

            
            Stream.zipWithIndex,
            Stream.map(([{ chunk, cost }, chunkIndex]) =>
              Data.struct({
                chunk,
                cost,
                chunkIndex,
              }),
            ),

            
            Stream.mapAccum(0, (accum: number, v) => [
              accum + Chunk.size(v.chunk),
              Data.struct({
                ...v,
                offset: accum,
              }),
            ]),

            
            Stream.mapEffect(({ chunk, cost, chunkIndex, offset }) =>
              Effect.gen(function* () {
                yield* Effect.log('appendRows権の取得完了(semaphore)', {
                  AppendRowsを実行する: {
                    chunkIndex,
                    行数: Chunk.size(chunk),
                    cost,
                    offset,
                  },
                });
                return yield* Effect.tryPromise(async () => {
                  const pw = writer.appendRows(
                    chunk.pipe(Chunk.toArray),
                    offset,
                  );
                  return await pw.getResult();
                }).pipe(Effect.andThen(handleAppendResult));
              }).pipe(
                Effect.tap(() =>
                  Effect.log(
                    `AppendRowsを実行しました: chunkIndex=${chunkIndex}`,
                  ),
                ),
                semaphore.withPermits(Math.min(cost, semaphorePermits)),
                Effect.forkIn(scope),
                Effect.withSpan('appendRows'),
                Effect.withSpan(`chunkIndex=${chunkIndex}`),
              ),
            ),
            Stream.runCollect,
            Effect.andThen(Fiber.joinAll),
          );

          yield* Effect.log(
            '全てのAppendRowsが完了しました。connection.finalizeを実行します。',
          );
          const rowCount = Option.fromNullable(
            yield* Effect.tryPromise(() => connection.finalize()),
          ).pipe(Option.map((e) => String(e.rowCount)));
          yield* Effect.log(`connection.finalize完了`, { rowCount });

          yield* Effect.log(
            `writeStreamをコミットします。destinationTable: ${destinationTable}, streamId: ${streamId}`,
          );
          const response = yield* Effect.tryPromise(() =>
            writeClient.batchCommitWriteStream({
              parent: destinationTable,
              writeStreams: [streamId],
            }),
          );
          yield* Effect.log('batchCommitWriteStream完了', {
            response,
          });
        }),
    ),
    Effect.withSpan('BigQuery StorageWrite API / writeStream'),
    Effect.withSpan(`対象テーブル=${table.id}`),
    Effect.scoped,
  );