feat: implement type-safe getResource

invakid404 · Nov 16, 2024 · dead321 · dead321
1 parent deadf31
commit dead321
Show file tree

Hide file tree

Showing 7 changed files with 120 additions and 91 deletions.
diff --git a/src/generator/common.ts b/src/generator/common.ts
@@ -1,80 +1,44 @@
 import toValidIdentifier from "to-valid-identifier";
 import { jsonSchemaToZod } from "json-schema-to-zod";
-import type { ResourceTypes } from "../windmill/resourceTypes.js";
 import { getContext, run } from "./context.js";
 import { JSONSchema } from "./types.js";
 import { InMemoryDuplex } from "../utils/inMemoryDuplex.js";
+import { resourceReferencesSchemaName } from "./resources.js";
 
 export const runWithBuffer = async <T,>(cb: () => T) => {
+  const { allResourceTypes } = getContext()!;
+
   const buffer = new InMemoryDuplex();
-  const result = await run(buffer, cb);
+  const result = await run(buffer, allResourceTypes, cb);
 
   return { buffer, result };
 };
 
-export type Runnable = {
+export type ResourceWithSchema = {
   path: string;
   schema?: JSONSchema;
 };
 
 export type GenerateSchemasOptions = {
-  allResourceTypes: ResourceTypes;
-  generator: AsyncGenerator<Runnable>;
+  generator: AsyncGenerator<ResourceWithSchema>;
   mapName: string;
 };
 
 export const generateSchemas = async ({
-  allResourceTypes,
   generator,
   mapName,
 }: GenerateSchemasOptions) => {
-  const { write } = getContext()!;
+  const { write, allResourceTypes } = getContext()!;
 
   const pathToSchemaMap = new Map<string, string>();
-  const referencedResourceTypes = new Set<string>();
 
   for await (const { path, schema } of generator) {
     if (schema == null) {
       continue;
     }
 
     const schemaName = toValidIdentifier(`${mapName}_${path}`);
-    const zodSchema = jsonSchemaToZod(schema, {
-      parserOverride: (schema, _refs) => {
-        // NOTE: Windmill sometimes has `default: null` on required fields,
-        //       which is incorrect for obvious reasons, so as a rule of thumb,
-        //       we remove null default values as a whole
-        if ("default" in schema && schema.default == null) {
-          delete schema.default;
-        }
-
-        // NOTE: Windmill sometimes has `enum: null` on string fields, and the
-        //       library doesn't like that, so we need to delete it
-        if (schema.type === "string" && schema.enum == null) {
-          delete schema.enum;
-
-          return;
-        }
-
-        const resourceTypeOrFalse = extractResourceTypeFromSchema(
-          schema as never,
-        );
-        if (resourceTypeOrFalse) {
-          const { resourceType } = resourceTypeOrFalse;
-          // NOTE: we could do a best-effort attempt to resolve non-resource
-          //       argument types by parsing the script sources (for TS only),
-          //       but handling things like types imported from elsewhere would
-          //       not be easy
-          if (!(resourceType in allResourceTypes)) {
-            return `z.any()`;
-          }
-
-          referencedResourceTypes.add(resourceType);
-
-          return toValidIdentifier(resourceType);
-        }
-      },
-    });
+    const zodSchema = schemaToZod(schema);
 
     write(`const ${schemaName} = lazyObject(() => ${zodSchema});`);
     pathToSchemaMap.set(path, schemaName);
@@ -85,8 +49,45 @@ export const generateSchemas = async ({
     write(`${JSON.stringify(scriptPath)}: ${schemaName},`);
   }
   write("} as const))");
+};
+
+export const schemaToZod = (schema: JSONSchema) => {
+  const { allResourceTypes } = getContext()!;
+
+  return jsonSchemaToZod(schema, {
+    parserOverride: (schema, _refs) => {
+      // NOTE: Windmill sometimes has `default: null` on required fields,
+      //       which is incorrect for obvious reasons, so as a rule of thumb,
+      //       we remove null default values as a whole
+      if ("default" in schema && schema.default == null) {
+        delete schema.default;
+      }
+
+      // NOTE: Windmill sometimes has `enum: null` on string fields, and the
+      //       library doesn't like that, so we need to delete it
+      if (schema.type === "string" && schema.enum == null) {
+        delete schema.enum;
+
+        return;
+      }
+
+      const resourceTypeOrFalse = extractResourceTypeFromSchema(
+        schema as never,
+      );
+      if (resourceTypeOrFalse) {
+        const { resourceType } = resourceTypeOrFalse;
+        // NOTE: we could do a best-effort attempt to resolve non-resource
+        //       argument types by parsing the script sources (for TS only),
+        //       but handling things like types imported from elsewhere would
+        //       not be easy
+        if (!(resourceType in allResourceTypes)) {
+          return `z.any()`;
+        }
 
-  return referencedResourceTypes;
+        return resourceReferencesSchemaName(resourceType);
+      }
+    },
+  });
 };
 
 const RESOURCE_TYPE_PREFIX = "resource-";

diff --git a/src/generator/context.ts b/src/generator/context.ts
@@ -1,13 +1,19 @@
 import { AsyncLocalStorage } from "node:async_hooks";
 import { Writable } from "node:stream";
+import { ResourceTypes } from "../windmill/resourceTypes.js";
 
 type GenerateContext = {
   write: (content: string) => Promise<void>;
+  allResourceTypes: ResourceTypes;
 };
 
 const generateStore = new AsyncLocalStorage<GenerateContext>();
 
-export const run = <T,>(output: Writable, cb: () => T) => {
+export const run = <T,>(
+  output: Writable,
+  allResourceTypes: ResourceTypes,
+  cb: () => T,
+) => {
   const write = (content: string) =>
     new Promise<void>((resolve, reject) =>
       output.write(content + "\n", (err) => {
@@ -19,7 +25,7 @@ export const run = <T,>(output: Writable, cb: () => T) => {
       }),
     );
 
-  return generateStore.run({ write }, cb);
+  return generateStore.run({ write, allResourceTypes }, cb);
 };
 
 export const getContext = () => generateStore.getStore();
diff --git a/src/generator/flows.ts b/src/generator/flows.ts
@@ -26,13 +26,12 @@ const preamble = dedent`
   };
 `;
 
-export const generateFlows = async (allResourceTypes: ResourceTypes) => {
+export const generateFlows = async () => {
   const { write } = getContext()!;
 
   await write(preamble);
 
   return generateSchemas({
-    allResourceTypes,
     generator: listFlows(),
     mapName,
   });

diff --git a/src/generator/index.ts b/src/generator/index.ts
@@ -7,25 +7,19 @@ import { generateResources } from "./resources.js";
 import { generateFlows } from "./flows.js";
 import { runWithBuffer } from "./common.js";
 
-export const generate = async (output: Writable) =>
-  run(output, async () => {
+export const generate = async (output: Writable) => {
+  const allResourceTypes = await listResourceTypes();
+  return run(output, allResourceTypes, async () => {
     writePreamble();
 
-    const allResourceTypes = await listResourceTypes();
     const results = await Promise.all(
-      [generateScripts, generateFlows].map((fn) =>
-        runWithBuffer(() => fn(allResourceTypes)),
+      [generateResources, generateScripts, generateFlows].map((fn) =>
+        runWithBuffer(fn),
       ),
     );
 
-    const referencedResourceTypes = results.reduce(
-      (acc, { buffer, result }) => {
-        buffer.pipe(output);
-
-        return acc.union(result);
-      },
-      new Set<string>(),
-    );
-
-    await generateResources([...referencedResourceTypes]);
+    results.forEach(({ buffer }) => {
+      buffer.pipe(output);
+    });
   });
+};
diff --git a/src/generator/resources.ts b/src/generator/resources.ts
@@ -1,39 +1,68 @@
 import toValidIdentifier from "to-valid-identifier";
-import { jsonSchemaToZod } from "json-schema-to-zod";
 import PQueue from "p-queue";
-import { listResourcesByType } from "../windmill/resources.js";
+import { listResources } from "../windmill/resources.js";
 import type { JSONSchema } from "./types.js";
 import { getContext } from "./context.js";
+import { schemaToZod } from "./common.js";
+import dedent from "dedent";
 
-export const generateResources = async (resourceTypes: string[]) => {
-  const { write } = getContext()!;
+const resourceToTypeMap = "resourceToType";
 
-  const resourceQueue = new PQueue({ concurrency: 5 });
+const preamble = dedent`
+  export const getResource = <Path extends keyof typeof resourceToType>(
+    path: Path,
+  ): Promise<z.infer<(typeof resourceToType)[Path]>> => wmill.getResource(path);
+`;
 
-  const resources = resourceTypes.map((resourceType) =>
-    resourceQueue.add(
-      async () => ({
-        resourceType,
-        paths: await Array.fromAsync(
-          listResourcesByType(resourceType),
-          ({ path }) => path,
-        ),
-      }),
-      { throwOnTimeout: true },
-    ),
-  );
+export const generateResources = async () => {
+  const { write, allResourceTypes } = getContext()!;
 
-  for await (const { resourceType, paths } of resources) {
-    const resourceSchema = makeResourceSchema(paths);
+  write(preamble);
 
-    const schemaName = toValidIdentifier(resourceType);
-    const zodSchema = jsonSchemaToZod(resourceSchema);
+  const resourcesByType = new Map<string, string[]>();
+  for await (const {
+    resource_type: resourceTypeName,
+    path,
+  } of listResources()) {
+    const paths = resourcesByType.get(resourceTypeName) ?? [];
+    resourcesByType.set(resourceTypeName, [...paths, path]);
+  }
+
+  for (const [resourceTypeName, paths] of resourcesByType) {
+    const resourceType = allResourceTypes[resourceTypeName]!;
+
+    const typeSchemaName = resourceTypeSchemaName(resourceType.name);
+    const resourceTypeSchema = schemaToZod(resourceType.schema as never);
+
+    write(`const ${typeSchemaName} = lazyObject(() => ${resourceTypeSchema});`);
+
+    const referencesSchemaName = resourceReferencesSchemaName(
+      resourceType.name,
+    );
+    const referencesSchema = schemaToZod(makeReferencesSchema(paths));
 
-    write(`const ${schemaName} = lazyObject(() => ${zodSchema});`);
+    write(
+      `const ${referencesSchemaName} = lazyObject(() => ${referencesSchema});`,
+    );
   }
+
+  write(`const ${resourceToTypeMap} = lazyObject(() => ({`);
+  for (const [resourceTypeName, paths] of resourcesByType) {
+    const typeSchemaName = resourceTypeSchemaName(resourceTypeName);
+    for (const path of paths) {
+      write(`${JSON.stringify(path)}: ${typeSchemaName},`);
+    }
+  }
+  write(`} as const));`);
 };
 
-const makeResourceSchema = (paths: string[]) => {
+export const resourceReferencesSchemaName = (resourceType: string) =>
+  toValidIdentifier(`${resourceType}_references`);
+
+export const resourceTypeSchemaName = (resourceType: string) =>
+  toValidIdentifier(`${resourceType}_type`);
+
+const makeReferencesSchema = (paths: string[]) => {
   const refs = paths.map((path) => `$res:${path}`);
 
   return {

diff --git a/src/generator/scripts.ts b/src/generator/scripts.ts
@@ -26,13 +26,12 @@ const preamble = dedent`
   };
 `;
 
-export const generateScripts = async (allResourceTypes: ResourceTypes) => {
+export const generateScripts = async () => {
   const { write } = getContext()!;
 
   await write(preamble);
 
   return generateSchemas({
-    allResourceTypes,
     generator: listScripts(),
     mapName,
   });

diff --git a/src/windmill/resources.ts b/src/windmill/resources.ts
@@ -2,15 +2,16 @@ import * as wmill from "windmill-client";
 
 const PER_PAGE = 20;
 
-export async function* listResourcesByType(resourceType: string) {
+export async function* listResources(resourceType?: string) {
   const workspace = process.env["WM_WORKSPACE"]!;
 
   for (let page = 1; ; ++page) {
     const pageData = await wmill.ResourceService.listResource({
       workspace,
       page,
       perPage: PER_PAGE,
-      resourceType,
+      resourceTypeExclude: "state,app_theme",
+      ...(resourceType != null && { resourceType }),
     });
 
     if (pageData.length === 0) {