diff --git a/packages/amp/src/arrow-flight/service.ts b/packages/amp/src/arrow-flight/service.ts index f13e225..4d29bbb 100644 --- a/packages/amp/src/arrow-flight/service.ts +++ b/packages/amp/src/arrow-flight/service.ts @@ -169,6 +169,11 @@ const make = Effect.gen(function* () { return flightDataStream.pipe( Stream.mapEffect( Effect.fnUntraced(function* (flightData): Effect.fn.Return, ArrowFlightError> { + // Amp sends keep-alive frames with no IPC message on long-lived streams + if (flightData.dataHeader.length === 0) { + return Option.none() + } + const messageType = yield* Effect.orDie(getMessageType(flightData)) switch (messageType) { diff --git a/packages/amp/test/arrow-flight/service.test.ts b/packages/amp/test/arrow-flight/service.test.ts index 4437b38..2d1b640 100644 --- a/packages/amp/test/arrow-flight/service.test.ts +++ b/packages/amp/test/arrow-flight/service.test.ts @@ -168,4 +168,33 @@ describe("ArrowFlight", () => { expect(results[0]?.data).toEqual([{ bin: toHex(firstExpected) }]) }) ) + + it.effect("skips keep-alive frames", ({ expect }) => + Effect.gen(function* () { + const testSchema = SchemaBuilder.schema().int32("id", { nullable: false }).build() + + const generated = yield* FlightDataGenerator.generateFlightData(testSchema, { + numRows: 3, + seed: 144 + }) + + const keepAlive = create(FlightDataSchema, {}) + const transport = makeTransport([ + keepAlive, + toProtoFlightData(generated.schemaFlightData, new Uint8Array(0)), + keepAlive, + toProtoFlightData(generated.recordBatchFlightData, recordBatchMetadata), + keepAlive + ]) + const layer = ArrowFlight.layer.pipe(Layer.provide(Layer.succeed(ArrowFlight.Transport, transport))) + + const results = yield* Effect.gen(function* () { + const flight = yield* ArrowFlight.ArrowFlight + return yield* flight.query("SELECT id FROM test") + }).pipe(Effect.provide(layer)) + + expect(results).toHaveLength(1) + expect(results[0]?.data).toEqual(generated.expectedValues.id?.map((id) => ({ id }))) + }) + ) })