From 506f1914fc537bb4fdc861676c084e335ba627ec Mon Sep 17 00:00:00 2001 From: Jason Stirnaman Date: Wed, 29 Jul 2026 18:03:03 -0500 Subject: [PATCH] fix: forward query routing options --- src/services/query.service.ts | 70 ++++++++++++++++-- tests/query-routing.test.ts | 132 ++++++++++++++++++++++++++++++++++ 2 files changed, 196 insertions(+), 6 deletions(-) create mode 100644 tests/query-routing.test.ts diff --git a/src/services/query.service.ts b/src/services/query.service.ts index 986d6e6..6d14604 100644 --- a/src/services/query.service.ts +++ b/src/services/query.service.ts @@ -5,6 +5,7 @@ */ import { BaseConnectionService } from "./base-connection.service.js"; +import type { QParamType } from "@influxdata/influxdb3-client"; import { InfluxProductType } from "../helpers/enums/influx-product-types.enum.js"; import { QueryLanguage, QuerySafetyService } from "./query-safety.service.js"; import { createRequestId } from "./telemetry.service.js"; @@ -118,11 +119,11 @@ export class QueryService { const connectionInfo = this.baseService.getConnectionInfo(); switch (connectionInfo.type) { case InfluxProductType.CloudDedicated: - return this.executeCloudDedicatedQuery(query, database); + return this.executeCloudDedicatedQuery(query, database, options); case InfluxProductType.Clustered: - return this.executeClusteredQuery(query, database); + return this.executeClusteredQuery(query, database, options); case InfluxProductType.CloudServerless: - return this.executeCloudServerlessQuery(query, database); + return this.executeCloudServerlessQuery(query, database, options); case InfluxProductType.Core: case InfluxProductType.Enterprise: return this.executeCoreEnterpriseQuery(query, database, { @@ -388,7 +389,7 @@ export class QueryService { ); case InfluxProductType.CloudDedicated: case InfluxProductType.Clustered: - return this.executeClusteredQuery(query, database); + return this.executeClusteredQuery(query, database, options); default: throw new Error( `InfluxQL queries are not supported for ${connectionInfo.type}`, @@ -508,11 +509,20 @@ export class QueryService { private async executeCloudDedicatedQuery( query: string, database: string, + options: { + params?: Record | unknown[]; + timeoutMs?: number; + }, ): Promise { + const params = this.flightParams(options); + try { const client = this.baseService.getClient(); if (!client) throw new Error("InfluxDB client not initialized"); - const result = client.queryPoints(query, database, { type: "sql" }); + const result = client.queryPoints(query, database, { + type: "sql", + ...(params !== undefined && { params }), + }); const rows: any[] = []; for await (const row of result) { rows.push(row); @@ -529,6 +539,10 @@ export class QueryService { private async executeClusteredQuery( query: string, database: string, + options: { + params?: Record | unknown[]; + timeoutMs?: number; + } = {}, ): Promise { try { const httpClient = this.baseService.getInfluxHttpClient(); @@ -536,7 +550,9 @@ export class QueryService { params: { db: database, q: query, + ...(options.params !== undefined && { params: options.params }), }, + ...(options.timeoutMs !== undefined && { timeout: options.timeoutMs }), }); return response; } catch (error: any) { @@ -550,11 +566,20 @@ export class QueryService { private async executeCloudServerlessQuery( query: string, database: string, + options: { + params?: Record | unknown[]; + timeoutMs?: number; + }, ): Promise { + const params = this.flightParams(options); + try { const client = this.baseService.getClient(); if (!client) throw new Error("InfluxDB client not initialized"); - const result = client.queryPoints(query, database, { type: "sql" }); + const result = client.queryPoints(query, database, { + type: "sql", + ...(params !== undefined && { params }), + }); const rows: any[] = []; for await (const row of result) { rows.push(row); @@ -565,6 +590,39 @@ export class QueryService { } } + private flightParams(options: { + params?: Record | unknown[]; + timeoutMs?: number; + }): Record | undefined { + if (options.timeoutMs !== undefined) { + throw new Error( + "timeoutMs is not supported for Flight queries; configure the client query timeout instead", + ); + } + + if (Array.isArray(options.params)) { + throw new Error( + "Array query params are not supported for Flight queries; use a named parameter object instead", + ); + } + + if ( + options.params !== undefined && + Object.values(options.params).some( + (value) => + typeof value !== "string" && + typeof value !== "number" && + typeof value !== "boolean", + ) + ) { + throw new Error( + "Flight query params must be strings, numbers, or booleans", + ); + } + + return options.params as Record | undefined; + } + /** * Centralized error handler for query methods */ diff --git a/tests/query-routing.test.ts b/tests/query-routing.test.ts new file mode 100644 index 0000000..1f9d63b --- /dev/null +++ b/tests/query-routing.test.ts @@ -0,0 +1,132 @@ +import { describe, expect, it, vi } from "vitest"; +import { InfluxProductType } from "../src/helpers/enums/influx-product-types.enum.js"; +import { QueryService } from "../src/services/query.service.js"; +import { BaseConnectionService } from "../src/services/base-connection.service.js"; + +const QUERY = "SELECT * FROM cpu WHERE host = $host"; +const DATABASE = "metrics"; +const PARAMS = { host: "edge-01" }; +const TIMEOUT_MS = 1_500; + +function stubBaseService(type: InfluxProductType): BaseConnectionService { + return { + validateDataCapabilities: vi.fn(), + getConnectionInfo: vi.fn().mockReturnValue({ type }), + getInfluxHttpClient: vi.fn(), + getClient: vi.fn(), + } as unknown as BaseConnectionService; +} + +function emptyAsyncResult() { + return (async function* () {})(); +} + +describe("query routing options", () => { + it.each([ + ["core", InfluxProductType.Core], + ["enterprise", InfluxProductType.Enterprise], + ])( + "%s forwards params and timeout to the v3 SQL API", + async (_name, type) => { + const base = stubBaseService(type); + const httpClient = { post: vi.fn().mockResolvedValue([]) }; + vi.mocked(base.getInfluxHttpClient).mockReturnValue(httpClient as any); + + await new QueryService(base).executeQuery(QUERY, DATABASE, { + params: PARAMS, + timeoutMs: TIMEOUT_MS, + }); + + expect(httpClient.post).toHaveBeenCalledWith( + "/api/v3/query_sql", + { db: DATABASE, q: QUERY, format: "json", params: PARAMS }, + expect.objectContaining({ timeout: TIMEOUT_MS }), + ); + }, + ); + + it.each([ + ["cloud-dedicated", InfluxProductType.CloudDedicated], + ["cloud-serverless", InfluxProductType.CloudServerless], + ])("%s forwards params to the Flight client", async (_name, type) => { + const base = stubBaseService(type); + const client = { queryPoints: vi.fn().mockReturnValue(emptyAsyncResult()) }; + vi.mocked(base.getClient).mockReturnValue(client as any); + + await new QueryService(base).executeQuery(QUERY, DATABASE, { + params: PARAMS, + }); + + expect(client.queryPoints).toHaveBeenCalledWith(QUERY, DATABASE, { + type: "sql", + params: PARAMS, + }); + }); + + it.each([ + ["cloud-dedicated", InfluxProductType.CloudDedicated], + ["cloud-serverless", InfluxProductType.CloudServerless], + ])("%s rejects per-query timeout explicitly", async (_name, type) => { + const base = stubBaseService(type); + vi.mocked(base.getClient).mockReturnValue({ + queryPoints: vi.fn().mockReturnValue(emptyAsyncResult()), + } as any); + + await expect( + new QueryService(base).executeQuery(QUERY, DATABASE, { + timeoutMs: TIMEOUT_MS, + }), + ).rejects.toThrow(/timeoutMs is not supported.*Flight/i); + }); + + it("clustered forwards params and timeout to its query API", async () => { + const base = stubBaseService(InfluxProductType.Clustered); + const httpClient = { get: vi.fn().mockResolvedValue({}) }; + vi.mocked(base.getInfluxHttpClient).mockReturnValue(httpClient as any); + const service = new QueryService(base); + + await service.executeQuery(QUERY, DATABASE, { + params: PARAMS, + timeoutMs: TIMEOUT_MS, + }); + + expect(httpClient.get).toHaveBeenCalledWith("/query", { + params: { db: DATABASE, q: QUERY, params: PARAMS }, + timeout: TIMEOUT_MS, + }); + }); + + it("preserves InfluxQL params and timeout on the v3 API", async () => { + const base = stubBaseService(InfluxProductType.Core); + const httpClient = { post: vi.fn().mockResolvedValue([]) }; + vi.mocked(base.getInfluxHttpClient).mockReturnValue(httpClient as any); + + await new QueryService(base).executeInfluxqlQuery(QUERY, DATABASE, { + params: PARAMS, + timeoutMs: TIMEOUT_MS, + }); + + expect(httpClient.post).toHaveBeenCalledWith( + "/api/v3/query_influxql", + { db: DATABASE, q: QUERY, format: "json", params: PARAMS }, + expect.objectContaining({ timeout: TIMEOUT_MS }), + ); + }); + + it("preserves the Clustered InfluxQL timeout", async () => { + const base = stubBaseService(InfluxProductType.Clustered); + const httpClient = { get: vi.fn().mockResolvedValue({}) }; + vi.mocked(base.getInfluxHttpClient).mockReturnValue(httpClient as any); + + await new QueryService(base).executeInfluxqlQuery( + "SELECT * FROM cpu", + DATABASE, + { timeoutMs: TIMEOUT_MS }, + ); + + expect(httpClient.get).toHaveBeenCalledWith( + "/query", + expect.objectContaining({ timeout: TIMEOUT_MS }), + ); + }); +});