@useairfoil/flight 0.3.0 → 0.4.1

package/dist/proto/Flight.proto

@@ -17,10 +17,11 @@
  */
 
 syntax = "proto3";
-import "google/protobuf/timestamp.proto";
 
 package arrow.flight.protocol;
 
+import "google/protobuf/timestamp.proto";
+
 /*
  * A flight service is an endpoint for retrieving or storing Arrow data. A
  * flight service can expose one or more predefined endpoints that can be
@@ -28,7 +29,6 @@ package arrow.flight.protocol;
  * can expose a set of actions that are available.
  */
 service FlightService {
-
   /*
    * Handshake between client and server. Depending on the server, the
    * handshake may be required to determine the token that should be used for
@@ -93,7 +93,7 @@ service FlightService {
    * GetFlightInfo this interface may generate a new flight that was not previously
    * available in ListFlights.
    */
-   rpc GetSchema(FlightDescriptor) returns (SchemaResult) {}
+  rpc GetSchema(FlightDescriptor) returns (SchemaResult) {}
 
   /*
    * Retrieve a single stream associated with a particular descriptor
@@ -138,14 +138,12 @@ service FlightService {
    * understand the capabilities of the flight service.
    */
   rpc ListActions(Empty) returns (stream ActionType) {}
-
 }
 
 /*
  * The request that a client provides to a server on handshake.
  */
 message HandshakeRequest {
-
   /*
    * A defined protocol version
    */
@@ -158,7 +156,6 @@ message HandshakeRequest {
 }
 
 message HandshakeResponse {
-
   /*
    * A defined protocol version
    */
@@ -276,12 +273,10 @@ message SchemaResult {
  * a flight or be used to expose a set of previously defined flights.
  */
 message FlightDescriptor {
-
   /*
    * Describes what type of descriptor is defined.
    */
   enum DescriptorType {
-
     // Protobuf pattern, not used.
     UNKNOWN = 0;
 
@@ -424,7 +419,6 @@ message PollInfo {
  * A particular stream or split associated with a flight.
  */
 message FlightEndpoint {
-
   /*
    * Token used to retrieve this stream.
    */
@@ -490,7 +484,6 @@ message Ticket {
  * A batch of Arrow data as part of a stream of batches.
  */
 message FlightData {
-
   /*
    * The descriptor of the data. This is only relevant when a client is
    * starting a new DoPut stream.
@@ -521,4 +514,4 @@ message FlightData {
  */
 message PutResult {
   bytes app_metadata = 1;
-}
+}

package/package.json

@@ -1,50 +1,26 @@
 {
   "name": "@useairfoil/flight",
-  "version": "0.3.0",
-  "type": "module",
+  "version": "0.4.1",
+  "repository": "useairfoil/typescript-sdk",
   "files": [
     "dist",
     "src",
     "README.md"
   ],
-  "main": "./dist/src/index.js",
-  "types": "./dist/src/index.d.ts",
+  "type": "module",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
   "exports": {
     ".": {
-      "types": "./dist/src/index.d.ts",
-      "import": "./dist/src/index.js",
-      "default": "./dist/src/index.js"
-    },
-    "./test": {
-      "types": "./dist/test/index.d.ts",
-      "import": "./dist/test/index.js",
-      "default": "./dist/test/index.js"
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js",
+      "default": "./dist/index.js"
     }
   },
   "scripts": {
     "build": "tsdown",
     "build:proto": "buf generate proto",
-    "typecheck": "tsc --noEmit",
-    "format": "biome format --write .",
-    "lint": "biome check .",
-    "lint:fix": "biome check . --write",
-    "test": "vitest",
-    "test:ci": "vitest run"
-  },
-  "peerDependencies": {
-    "effect": "^4.0.0-beta.25"
-  },
-  "optionalDependencies": {
-    "testcontainers": "^11.11.0"
-  },
-  "devDependencies": {
-    "@bufbuild/buf": "^1.57.2",
-    "@effect/vitest": "4.0.0-beta.25",
-    "@rollup/plugin-typescript": "^12.1.4",
-    "effect": "^4.0.0-beta.25",
-    "testcontainers": "^11.11.0",
-    "tsdown": "^0.15.6",
-    "vitest": "^3.2.4"
+    "typecheck": "tsc --noEmit"
   },
   "dependencies": {
     "@bufbuild/protobuf": "^2.9.0",
@@ -54,5 +30,16 @@
     "long": "^5.3.2",
     "nice-grpc": "^2.1.13",
     "nice-grpc-common": "^2.0.2"
+  },
+  "devDependencies": {
+    "@bufbuild/buf": "^1.5.7",
+    "@effect/vitest": "4.0.0-beta.83",
+    "@types/node": "^24.2.0",
+    "effect": "4.0.0-beta.83",
+    "tsdown": "^0.15.6",
+    "vitest": "^3.2.4"
+  },
+  "peerDependencies": {
+    "effect": "4.0.0-beta.83"
   }
 }

package/src/arrow-flight-client.ts

@@ -0,0 +1,168 @@
+import type { RecordBatch, Schema } from "apache-arrow";
+
+import { Config, Context, Effect, Layer, Scope, Stream } from "effect";
+import {
+  type CallOptions,
+  createClient,
+  type DefaultCallOptions,
+  type NormalizedServiceDefinition,
+} from "nice-grpc";
+
+import type { RecordBatchWithMetadata } from "./record-batch-with-metadata";
+
+import { decodeFlightDataStream, decodeSchemaFromFlightInfo } from "./arrow-utils";
+import {
+  FlightDecodeError,
+  type FlightClientErrorLike,
+  FlightProtocolError,
+  FlightTransportError,
+} from "./flight-client-error";
+import { type ClientOptions, createChannelFromConfig, type HostOrChannel } from "./proto-utils";
+import {
+  type FlightData,
+  type FlightDescriptor,
+  type FlightInfo,
+  type FlightServiceClient,
+  FlightServiceDefinition,
+  type HandshakeRequest,
+  type HandshakeResponse,
+  type PutResult,
+  type Ticket,
+} from "./proto/Flight";
+
+export type ArrowFlightClientOptions = HostOrChannel & {
+  readonly defaultCallOptions?: DefaultCallOptions<
+    NormalizedServiceDefinition<FlightServiceDefinition>
+  >;
+};
+
+export interface ArrowFlightClientService {
+  readonly handshake: (
+    request: AsyncIterable<HandshakeRequest>,
+    options?: CallOptions,
+  ) => Stream.Stream<HandshakeResponse, FlightClientErrorLike>;
+  readonly getFlightInfo: (
+    request: FlightDescriptor,
+    options?: CallOptions,
+  ) => Effect.Effect<FlightInfo, FlightClientErrorLike>;
+  readonly doGet: (
+    request: Ticket,
+    options: { readonly schema: Schema } & CallOptions,
+  ) => Stream.Stream<RecordBatchWithMetadata, FlightClientErrorLike>;
+  readonly doPut: (
+    request: AsyncIterable<FlightData>,
+    options?: CallOptions,
+  ) => AsyncIterable<PutResult>;
+  readonly executeFlightInfo: (
+    info: FlightInfo,
+    options?: CallOptions,
+  ) => Stream.Stream<RecordBatchWithMetadata, FlightClientErrorLike>;
+}
+
+export class ArrowFlightClient extends Context.Service<
+  ArrowFlightClient,
+  ArrowFlightClientService
+>()("@useairfoil/flight/ArrowFlightClient") {}
+
+const mapTransportError = (message: string, cause: unknown) =>
+  new FlightTransportError({
+    message,
+    cause,
+  });
+
+const mapDecodeError = (message: string, cause: unknown) =>
+  new FlightDecodeError({
+    message,
+    cause,
+  });
+
+const streamFromAsyncIterable = <A>(
+  iterable: AsyncIterable<A>,
+  message: string,
+): Stream.Stream<A, FlightTransportError> =>
+  Stream.fromAsyncIterable(iterable, (cause) => mapTransportError(message, cause));
+
+export const make = Effect.fnUntraced(function* (
+  options: ArrowFlightClientOptions,
+): Effect.fn.Return<ArrowFlightClientService, never, Scope.Scope> {
+  const ownsChannel = "host" in options;
+  const channel = createChannelFromConfig(options);
+  const client: FlightServiceClient = createClient(
+    FlightServiceDefinition,
+    channel,
+    options.defaultCallOptions,
+  );
+
+  if (ownsChannel) {
+    const scope = yield* Scope.Scope;
+    yield* Scope.addFinalizer(
+      scope,
+      Effect.sync(() => {
+        channel.close();
+      }),
+    );
+  }
+
+  const doGet = (
+    request: Ticket,
+    options: { readonly schema: Schema } & CallOptions,
+  ): Stream.Stream<RecordBatchWithMetadata, FlightClientErrorLike> => {
+    const { schema: expectedSchema } = options;
+    return Stream.fromAsyncIterable(
+      decodeFlightDataStream(client.doGet(request, options), { expectedSchema }),
+      (cause) => mapDecodeError("Flight batch decoding failed", cause),
+    );
+  };
+
+  const executeFlightInfo = (
+    info: FlightInfo,
+    options?: CallOptions,
+  ): Stream.Stream<RecordBatchWithMetadata, FlightClientErrorLike> => {
+    const schema = decodeSchemaFromFlightInfo(info);
+    if (!schema) {
+      return Stream.fail(
+        new FlightProtocolError({
+          message: "FlightInfo must include a schema",
+        }),
+      );
+    }
+
+    return Stream.fromIterable(info.endpoint).pipe(
+      Stream.flatMap((endpoint) =>
+        endpoint.ticket === undefined
+          ? Stream.empty
+          : doGet(endpoint.ticket, {
+              schema,
+              ...options,
+            }),
+      ),
+    );
+  };
+
+  return ArrowFlightClient.of({
+    handshake: (request, options) =>
+      streamFromAsyncIterable(client.handshake(request, options), "Flight handshake failed"),
+    getFlightInfo: (request, options) =>
+      Effect.tryPromise({
+        try: () => client.getFlightInfo(request, options),
+        catch: (cause) => mapTransportError("Flight getFlightInfo failed", cause),
+      }),
+    doGet,
+    doPut: (request, options) => client.doPut(request, options),
+    executeFlightInfo,
+  });
+});
+
+export const layer = (options: ArrowFlightClientOptions): Layer.Layer<ArrowFlightClient> =>
+  Layer.effect(ArrowFlightClient, make(options));
+
+export const layerConfig = (options: Config.Wrap<ArrowFlightClientOptions>) =>
+  Layer.effect(
+    ArrowFlightClient,
+    Effect.gen(function* () {
+      const resolved = yield* Config.unwrap(options);
+      return yield* make(resolved);
+    }),
+  );
+
+export type { CallOptions, ClientOptions, HostOrChannel, RecordBatch, RecordBatchWithMetadata };

package/src/arrow-flight-sql-client.ts

@@ -0,0 +1,147 @@
+import type { CallOptions } from "nice-grpc";
+
+import { Config, Context, Effect, Layer, Scope, Stream } from "effect";
+
+import type { FlightClientErrorLike } from "./flight-client-error";
+import type { RemoveTypeUrl } from "./proto-utils";
+import type { RecordBatchWithMetadata } from "./record-batch-with-metadata";
+
+import {
+  make as makeArrowFlightClient,
+  type ArrowFlightClientOptions,
+  type ArrowFlightClientService,
+} from "./arrow-flight-client";
+import { Any } from "./proto/any";
+import { FlightDescriptor, FlightDescriptor_DescriptorType, type FlightInfo } from "./proto/Flight";
+import {
+  CommandGetCatalogs,
+  CommandGetDbSchemas,
+  CommandGetTables,
+  CommandGetTableTypes,
+  CommandStatementQuery,
+} from "./proto/FlightSql";
+
+export type ArrowFlightSqlClientOptions = ArrowFlightClientOptions;
+
+export interface ArrowFlightSqlClientService {
+  readonly executeFlightInfo: (
+    info: FlightInfo,
+    options?: CallOptions,
+  ) => Stream.Stream<RecordBatchWithMetadata, FlightClientErrorLike>;
+  readonly getCatalogs: (
+    request: RemoveTypeUrl<CommandGetCatalogs>,
+    options?: CallOptions,
+  ) => Effect.Effect<FlightInfo, FlightClientErrorLike>;
+  readonly getDbSchemas: (
+    request: RemoveTypeUrl<CommandGetDbSchemas>,
+    options?: CallOptions,
+  ) => Effect.Effect<FlightInfo, FlightClientErrorLike>;
+  readonly getTables: (
+    request: RemoveTypeUrl<CommandGetTables>,
+    options?: CallOptions,
+  ) => Effect.Effect<FlightInfo, FlightClientErrorLike>;
+  readonly getTableTypes: (
+    request: RemoveTypeUrl<CommandGetTableTypes>,
+    options?: CallOptions,
+  ) => Effect.Effect<FlightInfo, FlightClientErrorLike>;
+  readonly executeQuery: (
+    request: RemoveTypeUrl<CommandStatementQuery>,
+    options?: CallOptions,
+  ) => Effect.Effect<FlightInfo, FlightClientErrorLike>;
+}
+
+export class ArrowFlightSqlClient extends Context.Service<
+  ArrowFlightSqlClient,
+  ArrowFlightSqlClientService
+>()("@useairfoil/flight/ArrowFlightSqlClient") {}
+
+export const make = Effect.fnUntraced(function* (
+  options: ArrowFlightSqlClientOptions,
+): Effect.fn.Return<ArrowFlightSqlClientService, never, Scope.Scope> {
+  const inner: ArrowFlightClientService = yield* makeArrowFlightClient(options);
+
+  return ArrowFlightSqlClient.of({
+    executeFlightInfo: (info, options) => inner.executeFlightInfo(info, options),
+    getCatalogs: (request, options) =>
+      inner.getFlightInfo(
+        createCommandDescriptor(
+          CommandGetCatalogs.$type,
+          CommandGetCatalogs.encode({
+            $type: CommandGetCatalogs.$type,
+            ...request,
+          }).finish(),
+        ),
+        options,
+      ),
+    getDbSchemas: (request, options) =>
+      inner.getFlightInfo(
+        createCommandDescriptor(
+          CommandGetDbSchemas.$type,
+          CommandGetDbSchemas.encode({
+            $type: CommandGetDbSchemas.$type,
+            ...request,
+          }).finish(),
+        ),
+        options,
+      ),
+    getTables: (request, options) =>
+      inner.getFlightInfo(
+        createCommandDescriptor(
+          CommandGetTables.$type,
+          CommandGetTables.encode({
+            $type: CommandGetTables.$type,
+            ...request,
+          }).finish(),
+        ),
+        options,
+      ),
+    getTableTypes: (request, options) =>
+      inner.getFlightInfo(
+        createCommandDescriptor(
+          CommandGetTableTypes.$type,
+          CommandGetTableTypes.encode({
+            $type: CommandGetTableTypes.$type,
+            ...request,
+          }).finish(),
+        ),
+        options,
+      ),
+    executeQuery: (request, options) =>
+      inner.getFlightInfo(
+        createCommandDescriptor(
+          CommandStatementQuery.$type,
+          CommandStatementQuery.encode({
+            $type: CommandStatementQuery.$type,
+            ...request,
+          }).finish(),
+        ),
+        options,
+      ),
+  });
+});
+
+export const layer = (options: ArrowFlightSqlClientOptions): Layer.Layer<ArrowFlightSqlClient> =>
+  Layer.effect(ArrowFlightSqlClient, make(options));
+
+export const layerConfig = (options: Config.Wrap<ArrowFlightSqlClientOptions>) =>
+  Layer.effect(
+    ArrowFlightSqlClient,
+    Effect.gen(function* () {
+      const resolved = yield* Config.unwrap(options);
+      return yield* make(resolved);
+    }),
+  );
+
+function createCommandDescriptor(typeUrl: string, value: Uint8Array): FlightDescriptor {
+  const cmd = Any.create({
+    typeUrl: `type.googleapis.com/${typeUrl}`,
+    value,
+  });
+
+  return FlightDescriptor.create({
+    type: FlightDescriptor_DescriptorType.CMD,
+    cmd: Any.encode(cmd).finish(),
+  });
+}
+
+export type { CallOptions, FlightInfo, RecordBatchWithMetadata };

package/src/arrow-utils.ts

@@ -1,10 +1,11 @@
 import { Message, type Schema } from "apache-arrow";
+
 import type { FlightData, FlightInfo } from "./proto/Flight";
+import type { RecordBatchWithMetadata } from "./record-batch-with-metadata";
+
 import { RecordBatchStreamReaderFromFlightData } from "./record-batch-decoder";
 
-export function decodeSchemaFromFlightInfo(
-  info: FlightInfo,
-): Schema | undefined {
+export function decodeSchemaFromFlightInfo(info: FlightInfo): Schema | undefined {
   // Notice that the `info.schema` field has the following format:
   // The schema of the dataset in its IPC form:
   //   4 bytes - an optional IPC_CONTINUATION_TOKEN prefix
@@ -24,7 +25,7 @@ export function getMessageSchema(message: Message): Schema | undefined {
 export function decodeFlightDataStream(
   stream: AsyncIterable<FlightData>,
   { expectedSchema: _expectedSchema }: { expectedSchema: Schema },
-) {
+): AsyncIterable<RecordBatchWithMetadata> {
   // TODO: we want to validate the schema of the stream?
   return new RecordBatchStreamReaderFromFlightData(stream);
 }

package/src/flight-client-error.ts

@@ -0,0 +1,27 @@
+import { Data } from "effect";
+
+export class FlightClientError extends Data.TaggedError("FlightClientError")<{
+  readonly message: string;
+  readonly cause?: unknown;
+}> {}
+
+export class FlightTransportError extends Data.TaggedError("FlightTransportError")<{
+  readonly message: string;
+  readonly cause?: unknown;
+}> {}
+
+export class FlightDecodeError extends Data.TaggedError("FlightDecodeError")<{
+  readonly message: string;
+  readonly cause?: unknown;
+}> {}
+
+export class FlightProtocolError extends Data.TaggedError("FlightProtocolError")<{
+  readonly message: string;
+  readonly cause?: unknown;
+}> {}
+
+export type FlightClientErrorLike =
+  | FlightClientError
+  | FlightTransportError
+  | FlightDecodeError
+  | FlightProtocolError;

package/src/flight-data-encoder.ts

@@ -1,6 +1,7 @@
 import { Message, type RecordBatch, type Schema } from "apache-arrow";
 import * as metadata from "apache-arrow/ipc/metadata/message";
 import { VectorAssembler } from "apache-arrow/visitor/vectorassembler";
+
 import { FlightData, type FlightDescriptor } from "./proto/Flight";
 
 export const FlightDataEncoder = {
@@ -23,21 +24,12 @@ export const FlightDataEncoder = {
     {
       appMetadata,
     }: {
-      appMetadata?: (args: {
-        index: number;
-        length: number;
-      }) => Uint8Array | undefined;
+      appMetadata?: (args: { index: number; length: number }) => Uint8Array | undefined;
     } = {},
   ): ReadonlyArray<FlightData> {
-    const { byteLength, nodes, bufferRegions, buffers } =
-      VectorAssembler.assemble(batch);
+    const { byteLength, nodes, bufferRegions, buffers } = VectorAssembler.assemble(batch);
 
-    const metadataRecordBatch = new metadata.RecordBatch(
-      batch.numRows,
-      nodes,
-      bufferRegions,
-      null,
-    );
+    const metadataRecordBatch = new metadata.RecordBatch(batch.numRows, nodes, bufferRegions, null);
 
     const message = Message.from(metadataRecordBatch, byteLength);
 
@@ -77,7 +69,7 @@ function encodeRecordBatchContent(
 ): Uint8Array {
   const out = new Uint8Array(byteLength);
   const bufGroupSize = compression != null ? 2 : 1;
-  const bufs = new Array(bufGroupSize);
+  const bufs = Array.from<ArrayBufferView<ArrayBufferLike>>({ length: bufGroupSize });
   let pos = 0;
 
   for (let i = 0; i < buffers.length; i += bufGroupSize) {

package/src/index.ts

@@ -1,16 +1,20 @@
-export { Metadata } from "nice-grpc";
-export { ArrowFlightClient } from "./arrow-flight";
-export { ArrowFlightSqlClient } from "./arrow-flight-sql";
+export * as ArrowFlightClient from "./arrow-flight-client";
+export * as ArrowFlightSqlClient from "./arrow-flight-sql-client";
+export * as FlightClientError from "./flight-client-error";
+export type { ArrowFlightClientOptions, ArrowFlightClientService } from "./arrow-flight-client";
+export type {
+  ArrowFlightSqlClientOptions,
+  ArrowFlightSqlClientService,
+} from "./arrow-flight-sql-client";
+export type { FlightClientErrorLike } from "./flight-client-error";
 export { FlightDataEncoder } from "./flight-data-encoder";
 export {
   FlightData,
   FlightDescriptor,
   FlightDescriptor_DescriptorType,
+  FlightInfo,
   PutResult,
   Ticket,
 } from "./proto/Flight";
-export {
-  type ClientOptions,
-  createChannelFromConfig,
-  type HostOrChannel,
-} from "./proto-utils";
+export { type ClientOptions, createChannelFromConfig, type HostOrChannel } from "./proto-utils";
+export type { RecordBatchWithMetadata } from "./record-batch-with-metadata";

package/src/proto/google/protobuf/descriptor.ts

@@ -34,6 +34,8 @@ PROTO3 = 999,
  */
 2023 = 1000,
 2024 = 1001,
+/** UNSTABLE - A placeholder edition for developing and testing unscheduled features. */
+UNSTABLE = 9999,
 /**
  * 1_TEST_ONLY - Placeholder editions for testing feature resolution.  These should not be
  * used or relied on outside of tests.
@@ -72,6 +74,9 @@ case 1000:
 case 1001:
       case "EDITION_2024":
         return Edition.2024;
+case 9999:
+      case "EDITION_UNSTABLE":
+        return Edition.UNSTABLE;
 case 1:
       case "EDITION_1_TEST_ONLY":
         return Edition.1_TEST_ONLY;
@@ -105,6 +110,7 @@ case Edition.PROTO2: return "EDITION_PROTO2";
 case Edition.PROTO3: return "EDITION_PROTO3";
 case Edition.2023: return "EDITION_2023";
 case Edition.2024: return "EDITION_2024";
+case Edition.UNSTABLE: return "EDITION_UNSTABLE";
 case Edition.1_TEST_ONLY: return "EDITION_1_TEST_ONLY";
 case Edition.2_TEST_ONLY: return "EDITION_2_TEST_ONLY";
 case Edition.99997_TEST_ONLY: return "EDITION_99997_TEST_ONLY";
@@ -1001,7 +1007,12 @@ readonly unverifiedLazy?: boolean | undefined,
  * is a formalization for deprecating fields.
  */
 readonly deprecated?: boolean | undefined,
-/** For Google-internal migration only. Do not use. */
+/**
+ * DEPRECATED. DO NOT USE!
+ * For Google-internal migration only. Do not use.
+ * 
+ * @deprecated
+ */
 readonly weak?: boolean | undefined,
 /**
  * Indicate that the field value should not be printed out when using debug

package/src/proto/google/protobuf/timestamp.ts

@@ -104,15 +104,16 @@ export const protobufPackage = "google.protobuf";
 export interface Timestamp {
   $type: "google.protobuf.Timestamp";
   /**
-   * Represents seconds of UTC time since Unix epoch
-   * 1970-01-01T00:00:00Z. Must be from 0001-01-01T00:00:00Z to
-   * 9999-12-31T23:59:59Z inclusive.
+   * Represents seconds of UTC time since Unix epoch 1970-01-01T00:00:00Z. Must
+   * be between -315576000000 and 315576000000 inclusive (which corresponds to
+   * 0001-01-01T00:00:00Z to 9999-12-31T23:59:59Z).
    */
   readonly seconds: bigint;
   /**
-   * Non-negative fractions of a second at nanosecond resolution. Negative
-   * second values with fractions must still have non-negative nanos values
-   * that count forward in time. Must be from 0 to 999,999,999
+   * Non-negative fractions of a second at nanosecond resolution. This field is
+   * the nanosecond portion of the duration, not an alternative to seconds.
+   * Negative second values with fractions must still have non-negative nanos
+   * values that count forward in time. Must be between 0 and 999,999,999
    * inclusive.
    */
   readonly nanos: number;

package/src/proto-utils.ts

@@ -28,11 +28,7 @@ export type HostOrChannel =
 
 export function createChannelFromConfig(config: HostOrChannel): Channel {
   if (config.host !== undefined) {
-    return createChannel(
-      config.host,
-      config.credentials,
-      config.channelOptions,
-    );
+    return createChannel(config.host, config.credentials, config.channelOptions);
   }
 
   return config.channel;