Skip to content

Commit 0a07b6b

Browse files
ditadiditadi
committed
feat: add type-safe SQL parameters markets with @param annotations
1 parent bd52e76 commit 0a07b6b

File tree

14 files changed

+314
-67
lines changed

14 files changed

+314
-67
lines changed

‎apps/dev-playground/client/src/appKitTypes.d.ts‎

Lines changed: 53 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Auto-generated by AppKit - DO NOT EDIT
22
// Generated by 'npx appkit-generate-types' or Vite plugin during build
33
import "@databricks/app-kit-ui/react";
4+
import type { SQLTypeMarker, SQLStringMarker, SQLNumberMarker, SQLBooleanMarker, SQLBinaryMarker, SQLDateMarker, SQLTimestampMarker } from "@databricks/app-kit-ui/js";
45

56
declare module "@databricks/app-kit-ui/react" {
67
interface QueryRegistry {
@@ -40,7 +41,20 @@ declare module "@databricks/app-kit-ui/react" {
4041
};
4142
spend_data: {
4243
name: "spend_data";
43-
parameters: { groupBy: string; aggregationLevel: string; startDate: string; endDate: string; appId: string; creator: string };
44+
parameters: {
45+
/** STRING - use sql.string() */
46+
groupBy: SQLStringMarker;
47+
/** STRING - use sql.string() */
48+
aggregationLevel: SQLStringMarker;
49+
/** DATE - use sql.date() */
50+
startDate: SQLDateMarker;
51+
/** DATE - use sql.date() */
52+
endDate: SQLDateMarker;
53+
/** STRING - use sql.string() */
54+
appId: SQLStringMarker;
55+
/** STRING - use sql.string() */
56+
creator: SQLStringMarker;
57+
};
4458
result: Array<{
4559
/** @sqlType STRING */
4660
group_key: string;
@@ -52,7 +66,14 @@ declare module "@databricks/app-kit-ui/react" {
5266
};
5367
spend_summary: {
5468
name: "spend_summary";
55-
parameters: { aggregationLevel: string; endDate: string; startDate: string };
69+
parameters: {
70+
/** STRING - use sql.string() */
71+
aggregationLevel: SQLStringMarker;
72+
/** DATE - use sql.date() */
73+
endDate: SQLDateMarker;
74+
/** DATE - use sql.date() */
75+
startDate: SQLDateMarker;
76+
};
5677
result: Array<{
5778
/** @sqlType DECIMAL */
5879
total: number;
@@ -64,7 +85,20 @@ declare module "@databricks/app-kit-ui/react" {
6485
};
6586
sql_helpers_test: {
6687
name: "sql_helpers_test";
67-
parameters: { stringParam: string; numberParam: string; booleanParam: string; dateParam: string; timestampParam: string; binaryParam: string };
88+
parameters: {
89+
/** STRING - use sql.string() */
90+
stringParam: SQLStringMarker;
91+
/** NUMERIC - use sql.number() */
92+
numberParam: SQLNumberMarker;
93+
/** BOOLEAN - use sql.boolean() */
94+
booleanParam: SQLBooleanMarker;
95+
/** DATE - use sql.date() */
96+
dateParam: SQLDateMarker;
97+
/** TIMESTAMP - use sql.timestamp() */
98+
timestampParam: SQLTimestampMarker;
99+
/** STRING - use sql.string() */
100+
binaryParam: SQLStringMarker;
101+
};
68102
result: Array<{
69103
/** @sqlType STRING */
70104
string_value: string;
@@ -86,7 +120,14 @@ declare module "@databricks/app-kit-ui/react" {
86120
};
87121
top_contributors: {
88122
name: "top_contributors";
89-
parameters: { aggregationLevel: string; startDate: string; endDate: string };
123+
parameters: {
124+
/** STRING - use sql.string() */
125+
aggregationLevel: SQLStringMarker;
126+
/** DATE - use sql.date() */
127+
startDate: SQLDateMarker;
128+
/** DATE - use sql.date() */
129+
endDate: SQLDateMarker;
130+
};
90131
result: Array<{
91132
/** @sqlType STRING */
92133
app_name: string;
@@ -96,7 +137,14 @@ declare module "@databricks/app-kit-ui/react" {
96137
};
97138
untagged_apps: {
98139
name: "untagged_apps";
99-
parameters: { aggregationLevel: string; startDate: string; endDate: string };
140+
parameters: {
141+
/** STRING - use sql.string() */
142+
aggregationLevel: SQLStringMarker;
143+
/** DATE - use sql.date() */
144+
startDate: SQLDateMarker;
145+
/** DATE - use sql.date() */
146+
endDate: SQLDateMarker;
147+
};
100148
result: Array<{
101149
/** @sqlType STRING */
102150
app_name: string;

‎apps/dev-playground/config/queries/spend_data.sql‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
-- @param groupBy STRING
2+
-- @param aggregationLevel STRING
3+
-- @param startDate DATE
4+
-- @param endDate DATE
5+
-- @param appId STRING
6+
-- @param creator STRING
17
SELECT
28
COALESCE(:groupBy, 'default') as group_key,
39
date_trunc(COALESCE(:aggregationLevel, 'day'), u.usage_date) AS aggregation_period,

‎apps/dev-playground/config/queries/spend_summary.sql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
-- @param aggregationLevel STRING
2+
-- @param startDate DATE
3+
-- @param endDate DATE
14
SELECT
25
ROUND(SUM(u.usage_quantity * lp.pricing.effective_list.default)) AS total,
36
ROUND(SUM(u.usage_quantity * lp.pricing.effective_list.default) /

‎apps/dev-playground/config/queries/sql_helpers_test.sql‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
-- @param stringParam STRING
2+
-- @param numberParam NUMERIC
3+
-- @param booleanParam BOOLEAN
4+
-- @param dateParam DATE
5+
-- @param timestampParam TIMESTAMP
6+
-- @param binaryParam STRING
17
SELECT
28
:stringParam as string_value,
39
:numberParam as number_value,

‎apps/dev-playground/config/queries/top_contributors.sql‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
-- Top contributors by app with aggregation support
2-
-- Parameters: workspaceId, startDate, endDate, aggregationLevel (daily, weekly, monthly)
2+
-- @param aggregationLevel STRING
3+
-- @param startDate DATE
4+
-- @param endDate DATE
35
WITH aggregated_costs AS (
46
SELECT
57
u.usage_metadata.app_name AS app_name,

‎apps/dev-playground/config/queries/untagged_apps.sql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
-- @param aggregationLevel STRING
2+
-- @param startDate DATE
3+
-- @param endDate DATE
14
WITH app_periods AS (
25
SELECT
36
u.usage_metadata.app_name AS app_name,

‎packages/app-kit-ui/src/js/index.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,12 @@
11
export {
22
isSQLTypeMarker,
3+
type SQLBinaryMarker,
4+
type SQLBooleanMarker,
5+
type SQLDateMarker,
6+
type SQLNumberMarker,
7+
type SQLStringMarker,
8+
type SQLTimestampMarker,
9+
type SQLTypeMarker,
310
sql,
411
} from "shared";
512
export * from "./sse";

‎packages/app-kit/src/type-generator/index.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import fs from "node:fs";
22
import dotenv from "dotenv";
3-
import { generateQueriesFromExplain, type QuerySchema } from "./query-registry";
3+
import { generateQueriesFromExplain } from "./query-registry";
4+
import type { QuerySchema } from "./types";
45

56
dotenv.config();
67

@@ -26,6 +27,7 @@ function generateTypeDeclarations(querySchemas: QuerySchema[] = []): string {
2627
return `// Auto-generated by AppKit - DO NOT EDIT
2728
// Generated by 'npx appkit-generate-types' or Vite plugin during build
2829
import "@databricks/app-kit-ui/react";
30+
import type { SQLTypeMarker, SQLStringMarker, SQLNumberMarker, SQLBooleanMarker, SQLBinaryMarker, SQLDateMarker, SQLTimestampMarker } from "@databricks/app-kit-ui/js";
2931
3032
declare module "@databricks/app-kit-ui/react" {
3133
interface QueryRegistry {${querySection}}

‎packages/app-kit/src/type-generator/query-registry.ts‎

Lines changed: 33 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3,40 +3,12 @@ import path from "node:path";
33
import { WorkspaceClient } from "@databricks/sdk-experimental";
44
import { CACHE_VERSION, hashSQL, loadCache, saveCache } from "./cache";
55
import { Spinner } from "./spinner";
6-
7-
/**
8-
* Query schema interface
9-
* @property name - the name of the query
10-
* @property type - the type of the query (string, number, boolean, object, array, etc.)
11-
*/
12-
export interface QuerySchema {
13-
name: string;
14-
type: string;
15-
}
16-
17-
/**
18-
* Databricks statement execution response interface
19-
* @property statement_id - the id of the statement
20-
* @property status - the status of the statement
21-
* @property manifest - the manifest of the statement
22-
* @property schema - the schema of the statement
23-
*/
24-
export interface DatabricksStatementExecutionResponse {
25-
statement_id: string;
26-
status: { state: string };
27-
manifest: {
28-
schema: {
29-
column_count: number;
30-
columns: {
31-
name: string;
32-
type_text: string;
33-
type_name: string;
34-
position: number;
35-
comment?: string;
36-
}[];
37-
};
38-
};
39-
}
6+
import {
7+
type DatabricksStatementExecutionResponse,
8+
type QuerySchema,
9+
sqlTypeToHelper,
10+
sqlTypeToMarker,
11+
} from "./types";
4012

4113
/**
4214
* Extract parameters from a SQL query
@@ -65,10 +37,22 @@ export function convertToQueryType(
6537
(p) => !SERVER_INJECTED_PARAMS.includes(p),
6638
);
6739

68-
// generate parameters types
40+
const paramTypes = extractParameterTypes(sql);
41+
42+
// generate parameters types with JSDoc hints
6943
const paramsType =
7044
params.length > 0
71-
? `{ ${params.map((p) => `${p}: string`).join("; ")} }`
45+
? `{\n ${params
46+
.map((p) => {
47+
const sqlType = paramTypes[p];
48+
// if no type annotation, use SQLTypeMarker (union type)
49+
const markerType = sqlType
50+
? sqlTypeToMarker[sqlType]
51+
: "SQLTypeMarker";
52+
const helper = sqlType ? sqlTypeToHelper[sqlType] : "sql.*()";
53+
return `/** ${sqlType || "any"} - use ${helper} */\n ${p}: ${markerType}`;
54+
})
55+
.join(";\n ")};\n }`
7256
: "Record<string, never>";
7357

7458
// generate result fields with JSDoc
@@ -96,6 +80,19 @@ export function convertToQueryType(
9680
}`;
9781
}
9882

83+
export function extractParameterTypes(sql: string): Record<string, string> {
84+
const paramTypes: Record<string, string> = {};
85+
const regex =
86+
/--\s*@param\s+(\w+)\s+(STRING|NUMERIC|BOOLEAN|DATE|TIMESTAMP|BINARY)/gi;
87+
const matches = sql.matchAll(regex);
88+
for (const match of matches) {
89+
const [, paramName, paramType] = match;
90+
paramTypes[paramName] = paramType.toUpperCase();
91+
}
92+
93+
return paramTypes;
94+
}
95+
9996
/**
10097
* Generate query schemas from a folder of SQL files
10198
* It uses the EXPLAIN command to generate the query schemas

‎packages/app-kit/src/type-generator/tests/query-registry.test.ts‎

Lines changed: 80 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
import { describe, expect, test } from "vitest";
22
import {
33
convertToQueryType,
4-
type DatabricksStatementExecutionResponse,
54
extractParameters,
5+
extractParameterTypes,
66
SERVER_INJECTED_PARAMS,
77
} from "../query-registry";
8+
import type { DatabricksStatementExecutionResponse } from "../types";
89

910
describe("extractParameters", () => {
1011
test("extracts parameters from SQL query", () => {
@@ -47,6 +48,69 @@ describe("SERVER_INJECTED_PARAMS", () => {
4748
});
4849
});
4950

51+
describe("extractParameterTypes", () => {
52+
test("extracts parameter types from SQL comments", () => {
53+
const sql = `-- @param startDate DATE
54+
-- @param endDate DATE
55+
-- @param groupBy STRING
56+
SELECT * FROM users WHERE date BETWEEN :startDate AND :endDate`;
57+
const types = extractParameterTypes(sql);
58+
59+
expect(types.startDate).toBe("DATE");
60+
expect(types.endDate).toBe("DATE");
61+
expect(types.groupBy).toBe("STRING");
62+
});
63+
64+
test("returns empty object for SQL without @param comments", () => {
65+
const sql = "SELECT * FROM users WHERE date = :startDate";
66+
const types = extractParameterTypes(sql);
67+
68+
expect(Object.keys(types).length).toBe(0);
69+
});
70+
71+
test("handles all supported types", () => {
72+
const sql = `-- @param str STRING
73+
-- @param num NUMERIC
74+
-- @param bool BOOLEAN
75+
-- @param dt DATE
76+
-- @param ts TIMESTAMP
77+
-- @param bin BINARY
78+
SELECT 1`;
79+
const types = extractParameterTypes(sql);
80+
81+
expect(types.str).toBe("STRING");
82+
expect(types.num).toBe("NUMERIC");
83+
expect(types.bool).toBe("BOOLEAN");
84+
expect(types.dt).toBe("DATE");
85+
expect(types.ts).toBe("TIMESTAMP");
86+
expect(types.bin).toBe("BINARY");
87+
});
88+
89+
test("ignores malformed @param comments", () => {
90+
const sql = `-- @param startDate
91+
-- @param INVALID
92+
-- @param noType
93+
-- this is not a param comment
94+
SELECT 1`;
95+
const types = extractParameterTypes(sql);
96+
97+
expect(Object.keys(types).length).toBe(0);
98+
});
99+
100+
test("handles mixed valid and invalid annotations", () => {
101+
const sql = `-- @param validDate DATE
102+
-- @param invalidParam
103+
-- @param validString STRING
104+
SELECT 1`;
105+
const types = extractParameterTypes(sql);
106+
107+
expect(types.validDate).toBe("DATE");
108+
expect(types.validString).toBe("STRING");
109+
expect(types.invalidParam).toBeUndefined();
110+
expect(Object.keys(types).length).toBe(2);
111+
});
112+
});
113+
50114
describe("convertToQueryType", () => {
51115
const mockResponse: DatabricksStatementExecutionResponse = {
52116
statement_id: "test-123",
@@ -74,7 +138,7 @@ describe("convertToQueryType", () => {
74138

75139
expect(result).toContain('name: "users"');
76140
expect(result).toContain("parameters:");
77-
expect(result).toContain("startDate: string");
141+
expect(result).toContain("startDate: SQLTypeMarker");
78142
expect(result).toContain("result: Array<{");
79143
});
80144

@@ -83,8 +147,20 @@ describe("convertToQueryType", () => {
83147
"SELECT * FROM users WHERE workspace_id = :workspaceId AND date = :startDate";
84148
const result = convertToQueryType(mockResponse, sql, "users");
85149

86-
expect(result).toContain("startDate: string");
87-
expect(result).not.toContain("workspaceId: string");
150+
expect(result).toContain("startDate: SQLTypeMarker");
151+
expect(result).not.toContain("workspaceId:");
152+
});
153+
154+
test("uses specific marker types when @param annotation is provided", () => {
155+
const sql = `-- @param startDate DATE
156+
-- @param count NUMERIC
157+
-- @param name STRING
158+
SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name`;
159+
const result = convertToQueryType(mockResponse, sql, "users");
160+
161+
expect(result).toContain("startDate: SQLDateMarker");
162+
expect(result).toContain("count: SQLNumberMarker");
163+
expect(result).toContain("name: SQLStringMarker");
88164
});
89165

90166
test("generates Record<string, never> for queries without params", () => {

0 commit comments

Comments
 (0)