@useairfoil/connector-kit 0.6.0 → 0.6.1

package/src/demo/definition.ts

@@ -0,0 +1,68 @@
+import type { Schema } from "effect";
+
+import type { ResourceRequiredField, ResourceSchema } from "../core/types";
+import type { DemoContext } from "./generate";
+
+/** A table a demo writes: the same name, row schema and version as its resource. */
+export type DemoTable<
+  Name extends string = string,
+  S extends ResourceSchema = ResourceSchema,
+  Version extends string = string,
+> = {
+  readonly name: Name;
+  readonly rowSchema: S;
+  readonly version: Version;
+};
+
+export const table = <
+  const Name extends string,
+  S extends ResourceSchema,
+  const Version extends ResourceRequiredField<S>,
+>(
+  definition: DemoTable<Name, S, Version>,
+): DemoTable<Name, S, Version> => definition;
+
+/** Rows by table. Required fields a row leaves out get placeholder values. */
+export type DemoRows<Tables extends ReadonlyArray<DemoTable>> = {
+  readonly [Table in Tables[number] as Table["name"]]?: ReadonlyArray<
+    Partial<Schema.Schema.Type<Table["rowSchema"]>>
+  >;
+};
+
+export type DemoDefinition<
+  Setup extends ReadonlyArray<DemoTable> = ReadonlyArray<DemoTable>,
+  Group extends ReadonlyArray<DemoTable> = ReadonlyArray<DemoTable>,
+> = {
+  readonly name: string;
+  /** Tables written once, before any group, such as a product catalog. */
+  readonly setup: {
+    readonly tables: Setup;
+    readonly rows: (context: DemoContext) => DemoRows<Setup>;
+  };
+  /**
+   * Tables for one related unit, such as a customer and their orders. The run
+   * adds groups until it reaches the size. Build keys from the index, so they
+   * are unique and other groups can point to them.
+   */
+  readonly group: {
+    readonly tables: Group;
+    readonly rows: (index: number, context: DemoContext) => DemoRows<Group>;
+    /**
+     * Groups written per chunk, 100 by default. A run can go over the size by
+     * one chunk, so lower it when a group is big.
+     */
+    readonly groupsPerChunk?: number;
+    /**
+     * The number of groups there are, when the provider only keeps so much
+     * history. The run stops there, even below the size.
+     */
+    readonly count?: number;
+  };
+};
+
+export const define = <
+  const Setup extends ReadonlyArray<DemoTable>,
+  const Group extends ReadonlyArray<DemoTable>,
+>(
+  definition: DemoDefinition<Setup, Group>,
+): DemoDefinition<Setup, Group> => definition;

package/src/demo/generate.ts

@@ -0,0 +1,167 @@
+import type { Faker } from "@faker-js/faker";
+
+import { Duration, Effect, Schema, type SchemaRepresentation } from "effect";
+
+import { ConnectorError } from "../errors";
+import { checkHasRepresentation } from "../iceberg/table";
+
+/** What one run shares: the seeded Faker and the time the run started. */
+export type DemoContext = {
+  readonly faker: Faker;
+  readonly now: Date;
+};
+
+/** Completes a row: keeps what it has, and fills the required fields it lacks. */
+export type RowFiller = (context: DemoContext, row: object) => Record<string, unknown>;
+
+type Representation = SchemaRepresentation.Representation;
+type Value = (context: DemoContext) => unknown;
+
+const twoYears = Duration.toMillis(Duration.days(730));
+
+class Unsupported extends Error {}
+
+const unsupported = (path: string, what: string): never => {
+  throw new Unsupported(`${path}: ${what} is not supported in demo data`);
+};
+
+/**
+ * Compiles a row schema into a filler for demo rows. It reads the same
+ * representation as `compileTable`, so values fit the Iceberg types. It fills
+ * only required fields the row lacks, so optional ones such as `_deleted` stay
+ * out.
+ */
+export const compileRow = (schema: Schema.Top): Effect.Effect<RowFiller, ConnectorError> =>
+  Effect.try({
+    try: () => {
+      const { representation, references } = Schema.toRepresentation(Schema.toType(schema));
+
+      const resolve = (node: Representation, path: string): Representation => {
+        if (node._tag !== "Reference") return node;
+        const target = references[node.$ref];
+
+        return target === undefined ? unsupported(path, `reference ${node.$ref}`) : target;
+      };
+
+      const requiredFields = (node: Representation, path: string) => {
+        const resolved = resolve(node, path);
+
+        if (resolved._tag !== "Objects" || resolved.indexSignatures.length > 0) {
+          return unsupported(path, resolved._tag);
+        }
+
+        return resolved.propertySignatures
+          .filter((property) => !property.isOptional)
+          .map((property) => {
+            const name = String(property.name);
+            return [name, value(property.type, `${path}.${name}`)] as const;
+          });
+      };
+
+      const object = (node: Representation, path: string): Value => {
+        const fields = requiredFields(node, path);
+        return (context) => Object.fromEntries(fields.map(([name, make]) => [name, make(context)]));
+      };
+
+      const value = (node: Representation, path: string): Value => {
+        const resolved = resolve(node, path);
+
+        switch (resolved._tag) {
+          case "String":
+            return ({ faker }) => faker.word.words({ count: { min: 1, max: 3 } });
+          case "Boolean":
+            return ({ faker }) => faker.datatype.boolean();
+          case "BigInt":
+            return ({ faker }) => BigInt(faker.number.int({ max: 100_000 }));
+          case "Number":
+            return checkHasRepresentation(resolved.checks, "effect/schema/isInt")
+              ? ({ faker }) => faker.number.int({ max: 1000 })
+              : ({ faker }) => faker.number.float({ max: 1000, fractionDigits: 2 });
+          case "Literal": {
+            const literal = resolved.literal;
+            return () => literal;
+          }
+          case "Union": {
+            // Fill every field, so nullable unions get their value type.
+            const members = resolved.types.filter(
+              (member) => member._tag !== "Null" && member._tag !== "Undefined",
+            );
+
+            const literals = members.flatMap((member) =>
+              member._tag === "Literal" ? [member.literal] : [],
+            );
+
+            if (literals.length > 0 && literals.length === members.length) {
+              return ({ faker }) => faker.helpers.arrayElement(literals);
+            }
+
+            const [only] = members;
+
+            return only !== undefined && members.length === 1
+              ? value(only, path)
+              : unsupported(path, "a union of several types");
+          }
+          case "Objects":
+            return object(resolved, path);
+          case "Arrays": {
+            const [element] = resolved.rest;
+
+            if (resolved.elements.length > 0 || resolved.rest.length !== 1 || !element) {
+              return unsupported(path, "a tuple");
+            }
+
+            const item = value(element, `${path}[]`);
+
+            return (context) =>
+              Array.from({ length: context.faker.number.int({ max: 3 }) }, () => item(context));
+          }
+          case "Declaration":
+            return declaration(resolved, path);
+          case "Suspend":
+            return value(resolved.thunk, path);
+          default:
+            return unsupported(path, resolved._tag);
+        }
+      };
+
+      const declaration = (node: SchemaRepresentation.Declaration, path: string): Value => {
+        switch (node.representation?.id) {
+          case "effect/schema/Date":
+            return ({ faker, now }) =>
+              faker.date.between({ from: now.getTime() - twoYears, to: now });
+          case "effect/schema/Uint8Array":
+            return ({ faker }) => Uint8Array.from({ length: 8 }, () => faker.number.int(255));
+          case "effect/schema/ReadonlyMap": {
+            const [keyNode, valueNode] = node.typeParameters;
+            if (!keyNode || !valueNode) return unsupported(path, "a map without key or value");
+            const key = value(keyNode, `${path}.key`);
+            const item = value(valueNode, `${path}.value`);
+
+            return (context) =>
+              new Map(
+                Array.from({ length: context.faker.number.int({ max: 3 }) }, () => [
+                  key(context),
+                  item(context),
+                ]),
+              );
+          }
+          default:
+            return unsupported(path, node.representation?.id ?? "an unnamed declaration");
+        }
+      };
+
+      const fields = requiredFields(representation, "$");
+
+      return (context, row) => ({
+        ...Object.fromEntries(
+          fields.filter(([name]) => !(name in row)).map(([name, make]) => [name, make(context)]),
+        ),
+        ...row,
+      });
+    },
+    catch: (cause) =>
+      new ConnectorError({
+        message: cause instanceof Unsupported ? cause.message : "Failed to compile demo rows",
+        cause,
+      }),
+  });

package/src/demo/index.ts

@@ -0,0 +1,3 @@
+export { type DemoContext } from "./generate";
+export { define, table, type DemoDefinition, type DemoRows, type DemoTable } from "./definition";
+export { run, type DemoTotals, type RunOptions } from "./run";

package/src/demo/run.ts

@@ -0,0 +1,142 @@
+import { Faker, base, en } from "@faker-js/faker";
+import { Array as Arr, Effect, Option, Schema, Stream } from "effect";
+
+import type { DemoDefinition, DemoTable } from "./definition";
+
+import { ConnectorError } from "../errors";
+import { ingestBatch } from "../ingestor/instrumented";
+import { Ingestor, type IngestResult } from "../ingestor/service";
+import { compileRow, type DemoContext } from "./generate";
+
+export type RunOptions = {
+  /** Stop once this many Arrow bytes are written. */
+  readonly bytes: number;
+  readonly seed: number;
+  /** The time data is made around. Same seed, time and tables give the same rows. */
+  readonly now: Date;
+  /** Groups per chunk. Overrides the demo's. */
+  readonly groupsPerChunk?: number;
+};
+
+export type DemoTotals = IngestResult & {
+  readonly tables: Readonly<Record<string, IngestResult>>;
+};
+
+const empty: DemoTotals = { rows: 0, bytes: 0, tables: {} };
+
+// Wings rejects Flight messages over 4 MiB. 500 rows keeps a batch of the
+// widest demo rows well under that.
+const batchRows = 500;
+
+const sum = (results: ReadonlyArray<IngestResult>): IngestResult => ({
+  rows: results.reduce((total, result) => total + result.rows, 0),
+  bytes: results.reduce((total, result) => total + result.bytes, 0),
+});
+
+const add = (totals: DemoTotals, written: ReadonlyArray<readonly [string, IngestResult]>) =>
+  written.reduce<DemoTotals>(
+    (current, [name, result]) => ({
+      rows: current.rows + result.rows,
+      bytes: current.bytes + result.bytes,
+      tables: {
+        ...current.tables,
+        [name]: {
+          rows: (current.tables[name]?.rows ?? 0) + result.rows,
+          bytes: (current.tables[name]?.bytes ?? 0) + result.bytes,
+        },
+      },
+    }),
+    totals,
+  );
+
+/**
+ * Writes the setup rows, then groups, until it reaches the size or the demo's
+ * count. It can go over the size by one chunk. Only bound tables are
+ * written; keys still point to the others.
+ */
+export const run = Effect.fnUntraced(function* (demo: DemoDefinition, options: RunOptions) {
+  const ingestor = yield* Ingestor;
+
+  const bound = (table: DemoTable) =>
+    ingestor.resourceNames === undefined || ingestor.resourceNames.includes(table.name);
+
+  // A fresh Faker per run, so its sequence depends only on the seed.
+  const faker = new Faker({ locale: [en, base] });
+  faker.seed(options.seed);
+  const context: DemoContext = { faker, now: options.now };
+
+  const prepare = (tables: ReadonlyArray<DemoTable>) =>
+    Effect.forEach(tables.filter(bound), (table) =>
+      compileRow(table.rowSchema).pipe(
+        Effect.map((fill) => ({
+          table,
+          fill,
+          decode: Schema.decodeUnknownEffect(Schema.toType(table.rowSchema)),
+        })),
+      ),
+    );
+
+  const write = (
+    tables: Effect.Success<ReturnType<typeof prepare>>,
+    rowsOf: (name: string) => ReadonlyArray<object>,
+  ) =>
+    Effect.forEach(tables, ({ table, fill, decode }) =>
+      Effect.forEach(rowsOf(table.name), (row) =>
+        decode(fill(context, row)).pipe(
+          Effect.mapError(
+            (cause) => new ConnectorError({ message: `Invalid demo row for ${table.name}`, cause }),
+          ),
+        ),
+      ).pipe(
+        Effect.flatMap((rows) =>
+          Effect.forEach(Arr.chunksOf(rows, batchRows), (batch) =>
+            ingestBatch({
+              connector: demo.name,
+              resource: table.name,
+              source: "demo",
+              batch: { rows: batch },
+            }),
+          ),
+        ),
+        Effect.map((results) => [table.name, sum(results)] as const),
+      ),
+    );
+
+  const setupTables = yield* prepare(demo.setup.tables);
+  const groupTables = yield* prepare(demo.group.tables);
+
+  const setup = demo.setup.rows(context);
+  const afterSetup = add(empty, yield* write(setupTables, (name) => setup[name] ?? []));
+
+  const totals =
+    afterSetup.bytes >= options.bytes || groupTables.length === 0
+      ? afterSetup
+      : yield* Stream.range(0, (demo.group.count ?? Infinity) - 1).pipe(
+          Stream.grouped(options.groupsPerChunk ?? demo.group.groupsPerChunk ?? 100),
+          Stream.mapEffect((indexes) => {
+            const groups = indexes.map((index) => demo.group.rows(index, context));
+            return write(groupTables, (name) => groups.flatMap((group) => group[name] ?? []));
+          }),
+          Stream.scan(() => afterSetup, add),
+          Stream.takeUntil((current) => current.bytes >= options.bytes),
+          Stream.tap((current) =>
+            Effect.logInfo("demo progress").pipe(
+              Effect.annotateLogs({ rows: current.rows, bytes: current.bytes }),
+            ),
+          ),
+          Stream.runLast,
+          Effect.map(Option.getOrElse(() => afterSetup)),
+        );
+
+  if (totals.bytes < options.bytes) {
+    yield* Effect.logWarning(
+      "The selected tables ran out of rows, so the run wrote less than the size asked for",
+    ).pipe(Effect.annotateLogs({ bytes: totals.bytes, target: options.bytes }));
+  }
+
+  yield* Effect.logInfo("demo done").pipe(
+    Effect.annotateLogs({ rows: totals.rows, bytes: totals.bytes, tables: totals.tables }),
+  );
+
+  return totals;
+});

package/src/iceberg/annotations.ts

@@ -0,0 +1,62 @@
+import { Schema } from "effect";
+
+declare module "effect/Schema" {
+  namespace Annotations {
+    interface Annotations {
+      readonly fieldId?: number | undefined;
+      readonly icebergIdentifierFields?: ReadonlyArray<string> | undefined;
+    }
+  }
+}
+
+/** Sets a field's Iceberg ID, and any other annotations. */
+export const field =
+  <A = unknown>(fieldId: number, annotations: Schema.Annotations.Key<A> = {}) =>
+  <S extends Schema.Schema<A>>(schema: S): S["Rebuild"] =>
+    schema.pipe(Schema.annotateKey({ ...annotations, fieldId }));
+
+/** Adds an offset to a base ID for reusable schemas. */
+export const ids =
+  (baseId: number) =>
+  (offset: number): number =>
+    baseId + offset;
+
+type KnownField<S extends Schema.Top> = keyof S["Type"] & string;
+
+/** A row field that cannot be null or undefined. */
+export type RequiredField<S extends Schema.Top> = [KnownField<S>] extends [never]
+  ? string
+  : {
+      readonly [Field in KnownField<S>]-?: undefined extends S["Type"][Field]
+        ? never
+        : null extends S["Type"][Field]
+          ? never
+          : Field;
+    }[KnownField<S>];
+
+/** A row schema whose identifier fields are `Key`. Only the type has it. */
+export interface Identified<Key extends string> {
+  readonly "~iceberg.identifier": Key;
+}
+
+/**
+ * Marks the fields that identify a row. Wings upserts rows by them.
+ *
+ * Put it last. A later `.annotate()` keeps the runtime annotation, but the type
+ * loses the key, so `Resource.entity` stops compiling.
+ *
+ * @example
+ * ```ts
+ * const Post = Schema.Struct({
+ *   id: Schema.String.pipe(Iceberg.field(1)),
+ *   updatedAt: Schema.String.pipe(Iceberg.field(2)),
+ * }).pipe(Iceberg.identifier("id"));
+ * ```
+ */
+export const identifier =
+  <S extends Schema.Top, const Key extends RequiredField<S>>(
+    ...keys: readonly [Key, ...ReadonlyArray<Key>]
+  ) =>
+  (schema: S): S["Rebuild"] & Identified<Key> =>
+    // Effect has no slot for the key in the schema type, so we add it.
+    schema.annotate({ icebergIdentifierFields: keys }) as S["Rebuild"] & Identified<Key>;

package/src/iceberg/error.ts

@@ -0,0 +1,9 @@
+import { Data } from "effect";
+
+/** A schema could not be turned into an Iceberg table. */
+export class IcebergSchemaError extends Data.TaggedError("IcebergSchemaError")<{
+  /** Schema path, such as `$.orders.element`. */
+  readonly path: string;
+  readonly message: string;
+  readonly cause?: unknown;
+}> {}

package/src/iceberg/index.ts

@@ -0,0 +1,3 @@
+export { field, identifier, ids } from "./annotations";
+export { IcebergSchemaError } from "./error";
+export { makeCreateTableCommitRequest, type CreateTableCommitOptions } from "./table";