feat: appkit exposed apis (#69)

Author: MarioCadenas

apps/dev-playground/server/reconnect-plugin.ts

@@ -15,7 +15,7 @@ interface ReconnectStreamResponse {
 
 export class ReconnectPlugin extends Plugin {
   public name = "reconnect";
-  public envVars = [];
+  protected envVars: string[] = [];
 
   injectRoutes(router: IAppRouter): void {
     this.route<ReconnectResponse>(router, {

apps/dev-playground/server/telemetry-example-plugin.ts

@@ -17,7 +17,7 @@ import type { Request, Response, Router } from "express";
 
 class TelemetryExamples extends Plugin {
   public name = "telemetry-examples" as const;
-  public envVars: string[] = [];
+  protected envVars: string[] = [];
 
   private requestCounter: Counter;
   private durationHistogram: Histogram;
@@ -516,5 +516,5 @@ class TelemetryExamples extends Plugin {
 export const telemetryExamples = toPlugin<
   typeof TelemetryExamples,
   BasePluginConfig,
-  "telemetry-examples"
->(TelemetryExamples, "telemetry-examples");
+  "telemetryExamples"
+>(TelemetryExamples, "telemetryExamples");

docs/docs/api/appkit/Class.AppKitError.md

@@ -76,6 +76,12 @@ readonly optional cause: Error;
 
 Optional cause of the error
 
+#### Overrides
+
+```ts
+Error.cause
+```
+
 ***
 
 ### code

docs/docs/api/appkit/Class.Plugin.md

@@ -143,10 +143,8 @@ asUser(req: Request): this;
 ```
 
 Execute operations using the user's identity from the request.
-
-Returns a scoped instance of this plugin where all method calls
-will execute with the user's Databricks credentials instead of
-the service principal.
+Returns a proxy of this plugin where all method calls execute
+with the user's Databricks credentials instead of the service principal.
 
 #### Parameters
 
@@ -158,31 +156,12 @@ the service principal.
 
 `this`
 
-A scoped plugin instance that executes as the user
+A proxied plugin instance that executes as the user
 
 #### Throws
 
 Error if user token is not available in request headers
 
-#### Example
-
-```typescript
-// In route handler - execute query as the requesting user
-router.post('/users/me/query/:key', async (req, res) => {
-  const result = await this.asUser(req).query(req.params.key)
-  res.json(result)
-})
-
-// Mixed execution in same handler
-router.post('/dashboard', async (req, res) => {
-  const [systemData, userData] = await Promise.all([
-    this.getSystemStats(),                  // Service principal
-    this.asUser(req).getUserPreferences(), // User context
-  ])
-  res.json({ systemData, userData })
-})
-```
-
 ***
 
 ### execute()
@@ -245,6 +224,28 @@ userKey?: string): Promise<void>;
 
 ***
 
+### exports()
+
+```ts
+exports(): unknown;
+```
+
+Returns the public exports for this plugin.
+Override this to define a custom public API.
+By default, returns an empty object.
+
+#### Returns
+
+`unknown`
+
+#### Implementation of
+
+```ts
+BasePlugin.exports
+```
+
+***
+
 ### getEndpoints()
 
 ```ts

docs/docs/api/appkit/Function.getExecutionContext.md

@@ -0,0 +1,18 @@
+# Function: getExecutionContext()
+
+```ts
+function getExecutionContext(): ExecutionContext;
+```
+
+Get the current execution context.
+
+- If running inside a user context (via asUser), returns the user context
+- Otherwise, returns the service context
+
+## Returns
+
+`ExecutionContext`
+
+## Throws
+
+Error if ServiceContext is not initialized

docs/docs/api/appkit/index.md

@@ -46,4 +46,5 @@ plugin architecture, and React integration.
 | ------ | ------ |
 | [appKitTypesPlugin](Function.appKitTypesPlugin.md) | Vite plugin to generate types for AppKit queries. Calls generateFromEntryPoint under the hood. |
 | [createApp](Function.createApp.md) | Bootstraps AppKit with the provided configuration. |
+| [getExecutionContext](Function.getExecutionContext.md) | Get the current execution context. |
 | [isSQLTypeMarker](Function.isSQLTypeMarker.md) | Type guard to check if a value is a SQL type marker |

docs/docs/api/appkit/typedoc-sidebar.ts

@@ -124,6 +124,11 @@ const typedocSidebar: SidebarsConfig = {
           id: "api/appkit/Function.createApp",
           label: "createApp"
         },
+        {
+          type: "doc",
+          id: "api/appkit/Function.getExecutionContext",
+          label: "getExecutionContext"
+        },
         {
           type: "doc",
           id: "api/appkit/Function.isSQLTypeMarker",

docs/docs/plugins.md

@@ -219,17 +219,25 @@ import type express from "express";
 
 class MyPlugin extends Plugin {
   name = "myPlugin";
-  envVars = [];                 // list required env vars here
-
-  injectRoutes(router: express.Router) {
-    this.route(router, {
-      name: "hello",
-      method: "get",
-      path: "/hello",
-      handler: async (_req, res) => {
-        res.json({ ok: true });
-      },
-    });
+  envVars = ["MY_API_KEY"];
+
+  async setup() {
+    // Initialize your plugin
+  }
+
+  myCustomMethod() {
+    // Some implementation
+  }
+
+  async shutdown() {
+    // Clean up resources
+  }
+
+  exports() {
+    // an object with the methods from this plugin to expose
+    return {
+      myCustomMethod: this.myCustomMethod
+    }
   }
 }
 
@@ -248,6 +256,23 @@ export const myPlugin = toPlugin<typeof MyPlugin, Record<string, never>, "myPlug
   - **Telemetry**: Instrument your plugin with traces and metrics via `this.telemetry`. See [`ITelemetry`](api/appkit/Interface.ITelemetry.md).
 - **Execution interceptors**: Use `execute()` and `executeStream()` with [`StreamExecutionSettings`](api/appkit/Interface.StreamExecutionSettings.md)
 
+**Consuming your plugin programmatically**
+
+Optionally, you may want to provide a way to consume your plugin programmatically using the AppKit object.
+To do that, your plugin needs to implement the `exports` method, returning an object with the methods you want to expose. From the previous example, the plugin could be consumed as follows:
+
+```ts
+const AppKit = await createApp({
+  plugins: [
+    server({ port: 8000 }),
+    analytics(),
+    myPlugin(),
+  ],
+});
+
+AppKit.myPlugin.myCustomMethod();
+```
+
 See the [`Plugin`](api/appkit/Class.Plugin.md) API reference for complete documentation.
 
 ## Caching

packages/appkit/src/analytics/analytics.ts

@@ -26,7 +26,7 @@ const logger = createLogger("analytics");
 
 export class AnalyticsPlugin extends Plugin {
   name = "analytics";
-  envVars = [];
+  protected envVars: string[] = [];
 
   protected static description = "Analytics plugin for data analysis";
   protected declare config: IAnalyticsConfig;
@@ -264,6 +264,19 @@ export class AnalyticsPlugin extends Plugin {
   async shutdown(): Promise<void> {
     this.streamManager.abortAll();
   }
+
+  /**
+   * Returns the public exports for the analytics plugin.
+   * Note: `asUser()` is automatically added by AppKit.
+   */
+  exports() {
+    return {
+      /**
+       * Execute a SQL query using service principal credentials.
+       */
+      query: this.query,
+    };
+  }
 }
 
 /**

packages/appkit/src/core/appkit.ts

@@ -13,9 +13,8 @@ import type { TelemetryConfig } from "../telemetry";
 import { TelemetryManager } from "../telemetry";
 
 export class AppKit<TPlugins extends InputPluginMap> {
-  private static _instance: AppKit<InputPluginMap> | null = null;
-  private pluginInstances: Record<string, BasePlugin> = {};
-  private setupPromises: Promise<void>[] = [];
+  #pluginInstances: Record<string, BasePlugin> = {};
+  #setupPromises: Promise<void>[] = [];
 
   private constructor(config: { plugins: TPlugins }) {
     const { plugins, ...globalConfig } = config;
@@ -47,7 +46,7 @@ export class AppKit<TPlugins extends InputPluginMap> {
     for (const [name, pluginData] of deferredPlugins) {
       if (pluginData) {
         this.createAndRegisterPlugin(globalConfig, name, pluginData, {
-          plugins: this.pluginInstances,
+          plugins: this.#pluginInstances,
         });
       }
     }
@@ -69,20 +68,72 @@ export class AppKit<TPlugins extends InputPluginMap> {
     };
     const pluginInstance = new Plugin(baseConfig);
 
-    this.pluginInstances[name] = pluginInstance;
+    this.#pluginInstances[name] = pluginInstance;
 
     pluginInstance.validateEnv();
 
-    this.setupPromises.push(pluginInstance.setup());
+    this.#setupPromises.push(pluginInstance.setup());
+
+    const self = this;
 
     Object.defineProperty(this, name, {
       get() {
-        return this.pluginInstances[name];
+        const plugin = self.#pluginInstances[name];
+        return self.wrapWithAsUser(plugin);
       },
       enumerable: true,
     });
   }
 
+  /**
+   * Binds all function properties in an exports object to the given context.
+   */
+  private bindExportMethods(
+    exports: Record<string, unknown>,
+    context: BasePlugin,
+  ) {
+    for (const key in exports) {
+      if (Object.hasOwn(exports, key) && typeof exports[key] === "function") {
+        exports[key] = (exports[key] as (...args: unknown[]) => unknown).bind(
+          context,
+        );
+      }
+    }
+  }
+
+  /**
+   * Wraps a plugin's exports with an `asUser` method that returns
+   * a user-scoped version of the exports.
+   */
+  private wrapWithAsUser<T extends BasePlugin>(plugin: T) {
+    // If plugin doesn't implement exports(), return empty object
+    const pluginExports = (plugin.exports?.() ?? {}) as Record<string, unknown>;
+    this.bindExportMethods(pluginExports, plugin);
+
+    // If plugin doesn't support asUser (no asUser method), return exports as-is
+    if (typeof (plugin as any).asUser !== "function") {
+      return pluginExports;
+    }
+
+    return {
+      ...pluginExports,
+      /**
+       * Execute operations using the user's identity from the request.
+       * Returns user-scoped exports where all methods execute with the
+       * user's Databricks credentials instead of the service principal.
+       */
+      asUser: (req: import("express").Request) => {
+        const userPlugin = (plugin as any).asUser(req);
+        const userExports = (userPlugin.exports?.() ?? {}) as Record<
+          string,
+          unknown
+        >;
+        this.bindExportMethods(userExports, userPlugin);
+        return userExports;
+      },
+    };
+  }
+
   static async _createApp<
     T extends PluginData<PluginConstructor, unknown, string>[],
   >(
@@ -106,11 +157,11 @@ export class AppKit<TPlugins extends InputPluginMap> {
       plugins: preparedPlugins,
     };
 
-    AppKit._instance = new AppKit(mergedConfig);
+    const instance = new AppKit(mergedConfig);
 
-    await Promise.all(AppKit._instance.setupPromises);
+    await Promise.all(instance.#setupPromises);
 
-    return AppKit._instance as unknown as PluginMap<T>;
+    return instance as unknown as PluginMap<T>;
   }
 
   private static preparePlugins(