refactor: move type generation responsabilit to individual plugins (#45)

Author: ditadi

packages/app-kit/src/analytics/analytics.ts

@@ -1,20 +1,24 @@
-import { SQLWarehouseConnector } from "../connectors";
-import { Plugin, toPlugin } from "../plugin";
+import { existsSync, type FSWatcher, watch } from "node:fs";
+import path from "node:path";
+import type { WorkspaceClient } from "@databricks/sdk-experimental";
 import type {
   IAppRouter,
   PluginExecuteConfig,
+  QuerySchemas,
   StreamExecutionSettings,
 } from "shared";
+import { generateQueryRegistryTypes } from "@/utils/type-generator";
+import { SQLWarehouseConnector } from "../connectors";
+import { Plugin, toPlugin } from "../plugin";
+import type { Request, Response } from "../utils";
+import { getRequestContext } from "../utils";
 import { queryDefaults } from "./defaults";
 import { QueryProcessor } from "./query";
 import {
   analyticsQueryResponseSchema,
   type IAnalyticsConfig,
   type IAnalyticsQueryRequest,
 } from "./types";
-import type { Request, Response } from "../utils";
-import { getRequestContext } from "../utils";
-import type { WorkspaceClient } from "@databricks/sdk-experimental";
 
 export class AnalyticsPlugin extends Plugin {
   name = "analytics";
@@ -27,6 +31,8 @@ export class AnalyticsPlugin extends Plugin {
   private SQLClient: SQLWarehouseConnector;
   private queryProcessor: QueryProcessor;
 
+  private schemaWatcher: FSWatcher | null = null;
+
   constructor(config: IAnalyticsConfig) {
     super(config);
     this.config = config;
@@ -36,6 +42,10 @@ export class AnalyticsPlugin extends Plugin {
       timeout: config.timeout,
       telemetry: this.telemetry,
     });
+
+    if (process.env.NODE_ENV === "development") {
+      this._generateQueryTypes();
+    }
   }
 
   injectRoutes(router: IAppRouter) {
@@ -162,6 +172,51 @@ export class AnalyticsPlugin extends Plugin {
 
   async shutdown(): Promise<void> {
     this.streamManager.abortAll();
+
+    if (this.schemaWatcher) {
+      this.schemaWatcher.close();
+      this.schemaWatcher = null;
+    }
+  }
+
+  // generate query types for development
+  private _generateQueryTypes() {
+    const schemaDir = path.join(process.cwd(), "config/queries");
+    const schemaPath = path.join(schemaDir, "schema.ts");
+
+    const typePath =
+      this.config.typePath || path.join(process.cwd(), "client", "src");
+
+    const generate = () => {
+      let querySchemas: QuerySchemas = {};
+      try {
+        delete require.cache[require.resolve(schemaPath)];
+        querySchemas = require(schemaPath).querySchemas;
+      } catch (error) {
+        if (existsSync(schemaPath)) {
+          console.warn(
+            `[AppKit] Failed to load query schemas from ${schemaPath}:`,
+            error instanceof Error ? error.message : error,
+          );
+        }
+      }
+      generateQueryRegistryTypes(querySchemas, typePath);
+    };
+
+    generate();
+
+    if (existsSync(schemaPath)) {
+      this.schemaWatcher = watch(
+        schemaDir,
+        { recursive: true },
+        (_event, filename) => {
+          if (filename === "schema.ts") {
+            console.log(`[AppKit] Query schema changed, regenerating types...`);
+            generate();
+          }
+        },
+      );
+    }
   }
 }
 

packages/app-kit/src/analytics/types.ts

@@ -3,6 +3,7 @@ import { z } from "zod";
 
 export interface IAnalyticsConfig extends BasePluginConfig {
   timeout?: number;
+  typePath?: string;
 }
 
 export interface IAnalyticsQueryRequest {

packages/app-kit/src/core/app-kit.ts

@@ -1,23 +1,18 @@
-import { existsSync, type FSWatcher, watch } from "node:fs";
-import path from "node:path";
 import type {
   BasePlugin,
   InputPluginMap,
   OptionalConfigPluginDef,
   PluginConstructor,
   PluginData,
   PluginMap,
-  QuerySchemas,
 } from "shared";
 import type { TelemetryConfig } from "../telemetry";
 import { TelemetryManager } from "../telemetry";
-import { generatePluginTypes } from "./type-generator";
 
 export class AppKit<TPlugins extends InputPluginMap> {
   private static _instance: AppKit<InputPluginMap> | null = null;
   private pluginInstances: Record<string, BasePlugin> = {};
   private setupPromises: Promise<void>[] = [];
-  private schemaWatcher: FSWatcher | null = null;
 
   private constructor(config: { plugins: TPlugins }) {
     const { plugins, ...globalConfig } = config;
@@ -102,20 +97,9 @@ export class AppKit<TPlugins extends InputPluginMap> {
 
     await Promise.all(AppKit._instance.setupPromises);
 
-    if (process.env.NODE_ENV === "development") {
-      AppKit._instance._generatePluginTypes(rawPlugins);
-    }
-
     return AppKit._instance as unknown as PluginMap<T>;
   }
 
-  cleanup() {
-    if (this.schemaWatcher) {
-      this.schemaWatcher.close();
-      this.schemaWatcher = null;
-    }
-  }
-
   private static preparePlugins(
     plugins: PluginData<PluginConstructor, unknown, string>[],
   ) {
@@ -128,60 +112,6 @@ export class AppKit<TPlugins extends InputPluginMap> {
     }
     return result;
   }
-
-  private _generatePluginTypes(
-    rawPlugins: PluginData<PluginConstructor, unknown, string>[],
-  ) {
-    const serverPlugin = rawPlugins.find((p) => p.name === "server");
-    const staticPath = (serverPlugin?.config as { staticPath?: string })
-      ?.staticPath;
-
-    const schemaDir = path.join(process.cwd(), "config/queries");
-    const querySchemaPath = path.join(schemaDir, "schema.ts");
-
-    const generate = () => {
-      let querySchemas: QuerySchemas = {};
-
-      try {
-        delete require.cache[require.resolve(querySchemaPath)];
-        querySchemas = require(querySchemaPath).querySchemas;
-      } catch (error) {
-        if (existsSync(querySchemaPath)) {
-          console.warn(
-            `[AppKit] Failed to load query schemas from ${querySchemaPath}:`,
-            error instanceof Error ? error.message : error,
-          );
-        }
-      }
-
-      generatePluginTypes(
-        rawPlugins.map((p) => ({ name: p.name })),
-        querySchemas,
-        staticPath,
-      );
-    };
-    generate();
-
-    if (this.schemaWatcher) {
-      this.schemaWatcher.close();
-      this.schemaWatcher = null;
-    }
-
-    if (existsSync(querySchemaPath)) {
-      this.schemaWatcher = watch(
-        schemaDir,
-        { recursive: true },
-        (_event, filename) => {
-          if (filename?.endsWith(".ts")) {
-            console.log(
-              `[AppKit] Query schema changed, regenerating types for ${filename}`,
-            );
-            generate();
-          }
-        },
-      );
-    }
-  }
 }
 
 export async function createApp<

packages/app-kit/src/core/type-generator.ts

@@ -4,6 +4,7 @@ import path from "node:path";
 import dotenv from "dotenv";
 import express from "express";
 import type { PluginPhase } from "shared";
+import { generatePluginRegistryTypes } from "@/utils/type-generator";
 import { Plugin, toPlugin } from "../plugin";
 import { instrumentations } from "../telemetry";
 import { databricksClientMiddleware, isRemoteServerEnabled } from "../utils";
@@ -68,6 +69,11 @@ export class ServerPlugin extends Plugin {
     this.serverApplication.use(await databricksClientMiddleware());
 
     this.extendRoutes();
+
+    if (process.env.NODE_ENV === "development") {
+      generatePluginRegistryTypes(this.config.plugins, this.config.staticPath);
+    }
+
     for (const extension of this.serverExtensions) {
       extension(this.serverApplication);
     }