@useairfoil/flight 0.4.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,6 +1,6 @@
 {
   "name": "@useairfoil/flight",
-  "version": "0.4.0",
+  "version": "0.4.1",
   "repository": "useairfoil/typescript-sdk",
   "files": [
     "dist",
@@ -20,9 +20,7 @@
   "scripts": {
     "build": "tsdown",
     "build:proto": "buf generate proto",
-    "typecheck": "tsc --noEmit",
-    "test": "vitest",
-    "test:ci": "vitest run"
+    "typecheck": "tsc --noEmit"
   },
   "dependencies": {
     "@bufbuild/protobuf": "^2.9.0",
@@ -35,18 +33,13 @@
   },
   "devDependencies": {
     "@bufbuild/buf": "^1.5.7",
-    "@effect/vitest": "4.0.0-beta.54",
+    "@effect/vitest": "4.0.0-beta.83",
     "@types/node": "^24.2.0",
-    "@useairfoil/wings-testing": "0.1.0",
-    "effect": "4.0.0-beta.54",
-    "testcontainers": "^11.14.0",
+    "effect": "4.0.0-beta.83",
     "tsdown": "^0.15.6",
     "vitest": "^3.2.4"
   },
   "peerDependencies": {
-    "effect": "4.0.0-beta.54"
-  },
-  "optionalDependencies": {
-    "testcontainers": "^11.14.0"
+    "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,6 +1,7 @@
 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";
 
@@ -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/index.ts

@@ -1,12 +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 { RecordBatchWithMetadata } from "./record-batch-with-metadata";