chore: fixup

Author: MarioCadenas

packages/app-kit-ui/src/react/hooks/__tests__/use-chart-data.test.ts

@@ -205,7 +205,7 @@ describe("useChartData", () => {
       );
     });
 
-    test("auto-selects ARROW by default when no heuristics match", () => {
+    test("auto-selects JSON by default when no heuristics match", () => {
       mockUseAnalyticsQuery.mockReturnValue({
         data: [],
         loading: false,
@@ -223,11 +223,11 @@ describe("useChartData", () => {
       expect(mockUseAnalyticsQuery).toHaveBeenCalledWith(
         "test",
         { limit: 100 },
-        expect.objectContaining({ format: "ARROW" }),
+        expect.objectContaining({ format: "JSON" }),
       );
     });
 
-    test("defaults to auto format (ARROW) when format is not specified", () => {
+    test("defaults to auto format (JSON) when format is not specified", () => {
       mockUseAnalyticsQuery.mockReturnValue({
         data: [],
         loading: false,
@@ -243,7 +243,7 @@ describe("useChartData", () => {
       expect(mockUseAnalyticsQuery).toHaveBeenCalledWith(
         "test",
         undefined,
-        expect.objectContaining({ format: "ARROW" }),
+        expect.objectContaining({ format: "JSON" }),
       );
     });
   });

packages/app-kit-ui/src/react/hooks/use-analytics-query.ts

@@ -126,14 +126,24 @@ export function useAnalyticsQuery<
 
           // success - Arrow format
           if (parsed.type === "arrow") {
-            const arrowData = await ArrowClient.fetchArrow(
-              getArrowStreamUrl(parsed.statement_id),
-            );
-            const table = await ArrowClient.processArrowBuffer(arrowData);
-            setLoading(false);
-            // Table is cast to TypedArrowTable with row type from QueryRegistry
-            setData(table as ResultType);
-            return;
+            try {
+              const arrowData = await ArrowClient.fetchArrow(
+                getArrowStreamUrl(parsed.statement_id),
+              );
+              const table = await ArrowClient.processArrowBuffer(arrowData);
+              setLoading(false);
+              // Table is cast to TypedArrowTable with row type from QueryRegistry
+              setData(table as ResultType);
+              return;
+            } catch (error) {
+              console.error(
+                "[useAnalyticsQuery] Failed to fetch Arrow data",
+                error,
+              );
+              setLoading(false);
+              setError("Unable to load data, please try again");
+              return;
+            }
           }
 
           // error

packages/app-kit-ui/src/react/hooks/use-chart-data.ts

@@ -72,11 +72,10 @@ function resolveFormat(
       return "ARROW";
     }
 
-    // Default to Arrow for better performance
-    return "ARROW";
+    return "JSON";
   }
 
-  return "ARROW";
+  return "JSON";
 }
 
 // ============================================================================

packages/app-kit/package.json

@@ -48,7 +48,6 @@
     "@opentelemetry/sdk-metrics": "^2.2.0",
     "@opentelemetry/sdk-node": "^0.208.0",
     "@opentelemetry/semantic-conventions": "^1.38.0",
-    "apache-arrow": "^21.1.0",
     "dotenv": "^16.6.1",
     "express": "^4.22.0",
     "pg": "^8.16.3",

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

@@ -8,7 +8,7 @@ import type {
 import { SQLWarehouseConnector } from "../connectors";
 import { Plugin, toPlugin } from "../plugin";
 import type { Request, Response } from "../utils";
-import { getRequestContext } from "../utils";
+import { getRequestContext, getWorkspaceClient } from "../utils";
 import { queryDefaults } from "./defaults";
 import { QueryProcessor } from "./query";
 import type {
@@ -41,8 +41,21 @@ export class AnalyticsPlugin extends Plugin {
   }
 
   injectRoutes(router: IAppRouter) {
-    // Inject core Arrow routes first (provides /arrow-result/:jobId endpoint)
-    this.injectCoreArrowRoutes(router);
+    this.route(router, {
+      method: "get",
+      path: "/arrow-result/:jobId",
+      handler: async (req: Request, res: Response) => {
+        await this._handleArrowRoute(req, res);
+      },
+    });
+
+    this.route(router, {
+      method: "get",
+      path: "/users/me/arrow-result/:jobId",
+      handler: async (req: Request, res: Response) => {
+        await this._handleArrowRoute(req, res, { asUser: true });
+      },
+    });
 
     this.route<AnalyticsQueryResponse>(router, {
       method: "post",
@@ -61,17 +74,58 @@ export class AnalyticsPlugin extends Plugin {
     });
   }
 
+  private async _handleArrowRoute(
+    req: Request,
+    res: Response,
+    { asUser = false }: { asUser?: boolean } = {},
+  ): Promise<void> {
+    try {
+      const { jobId } = req.params;
+
+      const workspaceClient = getWorkspaceClient(asUser);
+
+      console.log(
+        `Processing Arrow job request: ${jobId} for plugin: ${this.name}`,
+      );
+
+      const result = await this.getArrowData(workspaceClient, jobId);
+
+      res.setHeader("Content-Type", "application/octet-stream");
+      res.setHeader("Content-Length", result.data.length.toString());
+      res.setHeader("Cache-Control", "public, max-age=3600");
+
+      console.log(
+        `Sending Arrow buffer: ${result.data.length} bytes for job ${jobId}`,
+      );
+      res.send(Buffer.from(result.data));
+    } catch (error) {
+      console.error(`Arrow job error for ${this.name}:`, error);
+      res.status(404).json({
+        error: error instanceof Error ? error.message : "Arrow job not found",
+        plugin: this.name,
+      });
+    }
+  }
+
   private async _handleQueryRoute(
     req: Request,
     res: Response,
     { asUser = false }: { asUser?: boolean } = {},
   ): Promise<void> {
     const { query_key } = req.params;
     const { parameters, format = "JSON" } = req.body as IAnalyticsQueryRequest;
-    const formatParameters =
+    const queryParameters =
       format === "ARROW"
-        ? { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }
-        : {};
+        ? {
+            formatParameters: {
+              disposition: "EXTERNAL_LINKS",
+              format: "ARROW_STREAM",
+            },
+            type: "arrow",
+          }
+        : {
+            type: "result",
+          };
 
     const requestContext = getRequestContext();
     const userKey = asUser
@@ -126,17 +180,14 @@ export class AnalyticsPlugin extends Plugin {
         const result = await this.query(
           query,
           processedParams,
-          formatParameters,
+          queryParameters.formatParameters,
           signal,
           {
             asUser,
           },
         );
 
-        const type =
-          formatParameters.format === "ARROW_STREAM" ? "arrow" : "result";
-
-        return { type, ...result };
+        return { type: queryParameters.type, ...result };
       },
       streamExecutionSettings,
       userKey,
@@ -151,21 +202,11 @@ export class AnalyticsPlugin extends Plugin {
     { asUser = false }: { asUser?: boolean } = {},
   ): Promise<any> {
     const requestContext = getRequestContext();
+    const workspaceClient = getWorkspaceClient(asUser);
+
     const { statement, parameters: sqlParameters } =
       this.queryProcessor.convertToSQLParameters(query, parameters);
 
-    let workspaceClient: WorkspaceClient;
-    if (asUser) {
-      if (!requestContext.userDatabricksClient) {
-        throw new Error(
-          `User token passthrough feature is not enabled for workspace ${requestContext.workspaceId}.`,
-        );
-      }
-      workspaceClient = requestContext.userDatabricksClient;
-    } else {
-      workspaceClient = requestContext.serviceDatabricksClient;
-    }
-
     const response = await this.SQLClient.executeStatement(
       workspaceClient,
       {
@@ -180,11 +221,14 @@ export class AnalyticsPlugin extends Plugin {
     return response.result;
   }
 
+  // If we need arrow stream in more plugins we can define this as a base method in the core plugin class
+  // and have a generic endpoint for each plugin that consumes this arrow data.
   protected async getArrowData(
     workspaceClient: WorkspaceClient,
     jobId: string,
-  ): Promise<any> {
-    return await this.SQLClient.getArrowData(workspaceClient, jobId);
+    signal?: AbortSignal,
+  ): Promise<ReturnType<typeof this.SQLClient.getArrowData>> {
+    return await this.SQLClient.getArrowData(workspaceClient, jobId, signal);
   }
 
   async shutdown(): Promise<void> {

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

@@ -4,9 +4,10 @@ export interface IAnalyticsConfig extends BasePluginConfig {
   timeout?: number;
 }
 
+export type AnalyticsFormat = "JSON" | "ARROW";
 export interface IAnalyticsQueryRequest {
   parameters?: Record<string, any>;
-  format?: "JSON" | "ARROW";
+  format?: AnalyticsFormat;
 }
 
 export interface AnalyticsQueryResponse {

packages/app-kit/src/connectors/sql-warehouse/client.ts

@@ -63,9 +63,10 @@ export class SQLWarehouseConnector {
   private get arrowProcessor(): ArrowStreamProcessor {
     if (!this._arrowProcessor) {
       this._arrowProcessor = new ArrowStreamProcessor({
-        maxConcurrentDownloads: 5,
         timeout: this.config.timeout || executeStatementDefaults.timeout,
-        retries: 3,
+        maxConcurrentDownloads:
+          ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS,
+        retries: ArrowStreamProcessor.DEFAULT_RETRIES,
       });
     }
     return this._arrowProcessor;
@@ -403,25 +404,68 @@ export class SQLWarehouseConnector {
     jobId: string,
     signal?: AbortSignal,
   ): Promise<ReturnType<typeof this.arrowProcessor.processChunks>> {
-    try {
-      const response = await workspaceClient.statementExecution.getStatement(
-        {
-          statement_id: jobId,
+    const startTime = Date.now();
+
+    return this.telemetry.startActiveSpan(
+      "arrow.getData",
+      {
+        kind: SpanKind.CLIENT,
+        attributes: {
+          "db.system": "databricks",
+          "arrow.job_id": jobId,
         },
-        this._createContext(signal),
-      );
-      const chunks = response.result?.external_links;
-      const schema = response.manifest?.schema;
-
-      if (!chunks || !schema) {
-        throw new Error("No chunks or schema found in response");
-      }
-
-      return await this.arrowProcessor.processChunks(chunks, schema);
-    } catch (error) {
-      console.error(`Failed Arrow job: ${jobId}`, error);
-      throw error;
-    }
+      },
+      async (span: Span) => {
+        try {
+          const response =
+            await workspaceClient.statementExecution.getStatement(
+              { statement_id: jobId },
+              this._createContext(signal),
+            );
+
+          const chunks = response.result?.external_links;
+          const schema = response.manifest?.schema;
+
+          if (!chunks || !schema) {
+            throw new Error("No chunks or schema found in response");
+          }
+
+          span.setAttribute("arrow.chunk_count", chunks.length);
+
+          const result = await this.arrowProcessor.processChunks(
+            chunks,
+            schema,
+            signal,
+          );
+
+          span.setAttribute("arrow.data_size_bytes", result.data.length);
+          span.setStatus({ code: SpanStatusCode.OK });
+
+          const duration = Date.now() - startTime;
+          this.telemetryMetrics.queryDuration.record(duration, {
+            operation: "arrow.getData",
+            status: "success",
+          });
+
+          return result;
+        } catch (error) {
+          span.setStatus({
+            code: SpanStatusCode.ERROR,
+            message: error instanceof Error ? error.message : "Unknown error",
+          });
+          span.recordException(error as Error);
+
+          const duration = Date.now() - startTime;
+          this.telemetryMetrics.queryDuration.record(duration, {
+            operation: "arrow.getData",
+            status: "error",
+          });
+
+          console.error(`Failed Arrow job: ${jobId}`, error);
+          throw error;
+        }
+      },
+    );
   }
 
   // create context for cancellation token