-
Notifications
You must be signed in to change notification settings - Fork 9
feat: add INLINE + ARROW_STREAM format support for analytics plugin #256
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
9a385b5
dd1e1d5
ec677c6
b30f9eb
f6017a9
c3b7318
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,6 +19,7 @@ import { queryDefaults } from "./defaults"; | |
| import manifest from "./manifest.json"; | ||
| import { QueryProcessor } from "./query"; | ||
| import type { | ||
| AnalyticsFormat, | ||
| AnalyticsQueryResponse, | ||
| IAnalyticsConfig, | ||
| IAnalyticsQueryRequest, | ||
|
|
@@ -119,7 +120,8 @@ export class AnalyticsPlugin extends Plugin { | |
| res: express.Response, | ||
| ): Promise<void> { | ||
| const { query_key } = req.params; | ||
| const { parameters, format = "JSON" } = req.body as IAnalyticsQueryRequest; | ||
| const { parameters, format = "ARROW_STREAM" } = | ||
| req.body as IAnalyticsQueryRequest; | ||
|
|
||
| // Request-scoped logging with WideEvent tracking | ||
| logger.debug(req, "Executing query: %s (format=%s)", query_key, format); | ||
|
|
@@ -155,19 +157,6 @@ export class AnalyticsPlugin extends Plugin { | |
| const userKey = getCurrentUserId(); | ||
| const executorKey = isAsUser ? userKey : "global"; | ||
|
|
||
| const queryParameters = | ||
| format === "ARROW" | ||
| ? { | ||
| formatParameters: { | ||
| disposition: "EXTERNAL_LINKS", | ||
| format: "ARROW_STREAM", | ||
| }, | ||
| type: "arrow", | ||
| } | ||
| : { | ||
| type: "result", | ||
| }; | ||
|
|
||
| const hashedQuery = this.queryProcessor.hashQuery(query); | ||
|
|
||
| const defaultConfig: PluginExecuteConfig = { | ||
|
|
@@ -197,20 +186,115 @@ export class AnalyticsPlugin extends Plugin { | |
| parameters, | ||
| ); | ||
|
|
||
| const result = await executor.query( | ||
| return this._executeWithFormatFallback( | ||
| executor, | ||
| query, | ||
| processedParams, | ||
| queryParameters.formatParameters, | ||
| format, | ||
| signal, | ||
| ); | ||
|
|
||
| return { type: queryParameters.type, ...result }; | ||
| }, | ||
| streamExecutionSettings, | ||
| executorKey, | ||
| ); | ||
| } | ||
|
|
||
| /** Format configurations in fallback order. */ | ||
| private static readonly FORMAT_CONFIGS = { | ||
| ARROW_STREAM: { | ||
| formatParameters: { disposition: "INLINE", format: "ARROW_STREAM" }, | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. from this URL
so before changing anything this was already supporting arrow, can I know what's the case where this was failing? I would like to see it |
||
| type: "result" as const, | ||
| }, | ||
| JSON: { | ||
| formatParameters: undefined, | ||
| type: "result" as const, | ||
| }, | ||
| ARROW: { | ||
| formatParameters: { | ||
| disposition: "EXTERNAL_LINKS", | ||
| format: "ARROW_STREAM", | ||
| }, | ||
| type: "arrow" as const, | ||
| }, | ||
| }; | ||
|
|
||
| /** | ||
| * Execute a query with automatic format fallback. | ||
| * | ||
| * For the default ARROW_STREAM format, tries formats in order until one | ||
| * succeeds: ARROW_STREAM → JSON → ARROW. This handles warehouses that | ||
| * only support a subset of format/disposition combinations. | ||
| * | ||
| * Explicit format requests (JSON, ARROW) are not retried. | ||
| */ | ||
| private async _executeWithFormatFallback( | ||
| executor: AnalyticsPlugin, | ||
| query: string, | ||
| processedParams: | ||
| | Record<string, SQLTypeMarker | null | undefined> | ||
| | undefined, | ||
| requestedFormat: AnalyticsFormat, | ||
| signal?: AbortSignal, | ||
| ): Promise<{ type: string; [key: string]: any }> { | ||
| // Explicit format — no fallback. | ||
| if (requestedFormat === "JSON" || requestedFormat === "ARROW") { | ||
| const config = AnalyticsPlugin.FORMAT_CONFIGS[requestedFormat]; | ||
| const result = await executor.query( | ||
| query, | ||
| processedParams, | ||
| config.formatParameters, | ||
| signal, | ||
| ); | ||
| return { type: config.type, ...result }; | ||
| } | ||
|
|
||
| // Default (ARROW_STREAM) — try each format in order. | ||
| const fallbackOrder: AnalyticsFormat[] = ["ARROW_STREAM", "JSON", "ARROW"]; | ||
|
|
||
| for (let i = 0; i < fallbackOrder.length; i++) { | ||
| const fmt = fallbackOrder[i]; | ||
| const config = AnalyticsPlugin.FORMAT_CONFIGS[fmt]; | ||
| try { | ||
| const result = await executor.query( | ||
| query, | ||
| processedParams, | ||
| config.formatParameters, | ||
| signal, | ||
| ); | ||
| if (i > 0) { | ||
| logger.info( | ||
| "Query succeeded with fallback format %s (preferred %s was rejected)", | ||
| fmt, | ||
| fallbackOrder[0], | ||
| ); | ||
| } | ||
| return { type: config.type, ...result }; | ||
| } catch (err: unknown) { | ||
| const msg = err instanceof Error ? err.message : String(err); | ||
| const isFormatError = | ||
| msg.includes("ARROW_STREAM") || | ||
| msg.includes("JSON_ARRAY") || | ||
| msg.includes("EXTERNAL_LINKS") || | ||
| msg.includes("INVALID_PARAMETER_VALUE") || | ||
| msg.includes("NOT_IMPLEMENTED"); | ||
|
|
||
| if (!isFormatError || i === fallbackOrder.length - 1) { | ||
| throw err; | ||
| } | ||
|
|
||
| logger.warn( | ||
| "Format %s rejected by warehouse, falling back to %s: %s", | ||
| fmt, | ||
| fallbackOrder[i + 1], | ||
| msg, | ||
| ); | ||
| } | ||
| } | ||
|
|
||
| // Unreachable — last format in fallbackOrder throws on failure. | ||
| throw new Error("All format fallbacks exhausted"); | ||
| } | ||
|
|
||
| /** | ||
| * Execute a SQL query using the current execution context. | ||
| * | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
in theory arrow is the same as arrow_stream, so I'm not following what's the problem?