Skip to content

Commit 99567b9

Browse files
BraydenBrayden
authored
Merge pull request #54 from Ehesp/ehesp/rework
Added type safety, feature toggles and shared RPC
2 parents 7c5a783 + 2ccce1b commit 99567b9

File tree

22 files changed

+2228
-1822
lines changed

22 files changed

+2228
-1822
lines changed

‎package.json‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,10 @@
99
"cf-typegen": "wrangler types"
1010
},
1111
"devDependencies": {
12-
"@cloudflare/workers-types": "^4.20240925.0",
12+
"@cloudflare/workers-types": "^4.20241216.0",
1313
"@types/pg": "^8.11.10",
14-
"typescript": "^5.5.2",
15-
"wrangler": "^3.60.3"
14+
"typescript": "^5.7.2",
15+
"wrangler": "^3.96.0"
1616
},
1717
"dependencies": {
1818
"@libsql/client": "^0.14.0",

‎pnpm-lock.yaml‎

Lines changed: 137 additions & 238 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/allowlist/index.ts‎

Lines changed: 78 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -1,86 +1,95 @@
1-
import { HandlerConfig } from "../handler";
2-
import { DataSource } from "../types";
1+
import { StarbaseDBConfiguration } from "../handler";
2+
import { DataSource, QueryResult } from "../types";
33

4-
const parser = new (require('node-sql-parser').Parser)();
4+
const parser = new (require("node-sql-parser").Parser)();
55

66
let allowlist: string[] | null = null;
77
let normalizedAllowlist: any[] | null = null;
88

99
function normalizeSQL(sql: string) {
10-
// Remove trailing semicolon. This allows a user to send a SQL statement that has
11-
// a semicolon where the allow list might not include it but both statements can
12-
// equate to being the same. AST seems to have an issue with matching the difference
13-
// when included in one query vs another.
14-
return sql.trim().replace(/;\s*$/, '');
10+
// Remove trailing semicolon. This allows a user to send a SQL statement that has
11+
// a semicolon where the allow list might not include it but both statements can
12+
// equate to being the same. AST seems to have an issue with matching the difference
13+
// when included in one query vs another.
14+
return sql.trim().replace(/;\s*$/, "");
1515
}
1616

17-
async function loadAllowlist(dataSource?: DataSource): Promise<string[]> {
18-
try {
19-
const statement = 'SELECT sql_statement FROM tmp_allowlist_queries'
20-
const result = await dataSource?.internalConnection?.durableObject.executeQuery(statement, [], false) as any[];
21-
return result.map((row: any) => row.sql_statement);
22-
} catch (error) {
23-
console.error('Error loading allowlist:', error);
24-
return [];
25-
}
17+
async function loadAllowlist(dataSource: DataSource): Promise<string[]> {
18+
try {
19+
const statement = "SELECT sql_statement FROM tmp_allowlist_queries";
20+
const result = await dataSource.rpc.executeQuery({ sql: statement }) as QueryResult[]
21+
return result.map((row) => String(row.sql_statement));
22+
} catch (error) {
23+
console.error("Error loading allowlist:", error);
24+
return [];
25+
}
2626
}
2727

28-
export async function isQueryAllowed(sql: string, isEnabled: boolean, dataSource?: DataSource, config?: HandlerConfig): Promise<boolean | Error> {
29-
// If the feature is not turned on then by default the query is allowed
30-
if (!isEnabled) return true;
28+
export async function isQueryAllowed(opts: {
29+
sql: string;
30+
isEnabled: boolean;
31+
dataSource: DataSource;
32+
config: StarbaseDBConfiguration;
33+
}): Promise<boolean | Error> {
34+
const { sql, isEnabled, dataSource, config } = opts;
35+
36+
// If the feature is not turned on then by default the query is allowed
37+
if (!isEnabled) return true;
38+
39+
// If we are using the administrative AUTHORIZATION token value, this request is allowed.
40+
// We want database UI's to be able to have more free reign to run queries so we can load
41+
// tables, run queries, and more. If you want to block queries with the allowlist then we
42+
// advise you to do so by implementing user authentication with JWT.
43+
if (config.role === "admin") {
44+
return true;
45+
}
3146

32-
// If we are using the administrative AUTHORIZATION token value, this request is allowed.
33-
// We want database UI's to be able to have more free reign to run queries so we can load
34-
// tables, run queries, and more. If you want to block queries with the allowlist then we
35-
// advise you to do so by implementing user authentication with JWT.
36-
if (dataSource?.request.headers.get('Authorization') === `Bearer ${config?.adminAuthorizationToken}`) {
37-
return true;
47+
allowlist = await loadAllowlist(dataSource);
48+
normalizedAllowlist = allowlist.map((query) =>
49+
parser.astify(normalizeSQL(query))
50+
);
51+
52+
try {
53+
if (!sql) {
54+
return Error("No SQL provided for allowlist check");
3855
}
3956

40-
allowlist = await loadAllowlist(dataSource);
41-
normalizedAllowlist = allowlist.map(query => parser.astify(normalizeSQL(query)));
42-
43-
try {
44-
if (!sql) {
45-
return Error('No SQL provided for allowlist check')
46-
}
57+
const normalizedQuery = parser.astify(normalizeSQL(sql));
58+
59+
// Compare ASTs while ignoring specific values
60+
const isCurrentAllowed = normalizedAllowlist?.some((allowedQuery) => {
61+
// Create deep copies to avoid modifying original ASTs
62+
const allowedAst = JSON.parse(JSON.stringify(allowedQuery));
63+
const queryAst = JSON.parse(JSON.stringify(normalizedQuery));
4764

48-
const normalizedQuery = parser.astify(normalizeSQL(sql));
49-
50-
// Compare ASTs while ignoring specific values
51-
const isCurrentAllowed = normalizedAllowlist?.some(allowedQuery => {
52-
// Create deep copies to avoid modifying original ASTs
53-
const allowedAst = JSON.parse(JSON.stringify(allowedQuery));
54-
const queryAst = JSON.parse(JSON.stringify(normalizedQuery));
55-
56-
// Remove or normalize value fields from both ASTs
57-
const normalizeAst = (ast: any) => {
58-
if (Array.isArray(ast)) {
59-
ast.forEach(normalizeAst);
60-
} else if (ast && typeof ast === 'object') {
61-
// Remove or normalize fields that contain specific values
62-
if ('value' in ast) {
63-
ast.value = '?';
64-
}
65-
66-
Object.values(ast).forEach(normalizeAst);
67-
}
68-
69-
return ast;
70-
};
71-
72-
normalizeAst(allowedAst);
73-
normalizeAst(queryAst);
74-
75-
return JSON.stringify(allowedAst) === JSON.stringify(queryAst);
76-
});
77-
78-
if (!isCurrentAllowed) {
79-
throw new Error("Query not allowed");
65+
// Remove or normalize value fields from both ASTs
66+
const normalizeAst = (ast: any) => {
67+
if (Array.isArray(ast)) {
68+
ast.forEach(normalizeAst);
69+
} else if (ast && typeof ast === "object") {
70+
// Remove or normalize fields that contain specific values
71+
if ("value" in ast) {
72+
ast.value = "?";
73+
}
74+
75+
Object.values(ast).forEach(normalizeAst);
8076
}
8177

82-
return true;
83-
} catch (error: any) {
84-
throw new Error(error?.message ?? 'Error');
78+
return ast;
79+
};
80+
81+
normalizeAst(allowedAst);
82+
normalizeAst(queryAst);
83+
84+
return JSON.stringify(allowedAst) === JSON.stringify(queryAst);
85+
});
86+
87+
if (!isCurrentAllowed) {
88+
throw new Error("Query not allowed");
8589
}
86-
}
90+
91+
return true;
92+
} catch (error: any) {
93+
throw new Error(error?.message ?? "Error");
94+
}
95+
}

‎src/cache/index.ts‎

Lines changed: 118 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -1,52 +1,79 @@
1-
import { DataSource, Source } from "../types";
2-
const parser = new (require('node-sql-parser').Parser)();
1+
import { StarbaseDBConfiguration } from "../handler";
2+
import { DataSource, QueryResult } from "../types";
3+
import sqlparser from "node-sql-parser";
4+
const parser = new sqlparser.Parser();
35

46
function hasModifyingStatement(ast: any): boolean {
5-
// Check if current node is a modifying statement
6-
if (ast.type && ['insert', 'update', 'delete'].includes(ast.type.toLowerCase())) {
7-
return true;
8-
}
7+
// Check if current node is a modifying statement
8+
if (
9+
ast.type &&
10+
["insert", "update", "delete"].includes(ast.type.toLowerCase())
11+
) {
12+
return true;
13+
}
914

10-
// Recursively check all properties of the AST
11-
for (const key in ast) {
12-
if (typeof ast[key] === 'object' && ast[key] !== null) {
13-
if (Array.isArray(ast[key])) {
14-
if (ast[key].some(item => hasModifyingStatement(item))) {
15-
return true;
16-
}
17-
} else if (hasModifyingStatement(ast[key])) {
18-
return true;
19-
}
15+
// Recursively check all properties of the AST
16+
for (const key in ast) {
17+
if (typeof ast[key] === "object" && ast[key] !== null) {
18+
if (Array.isArray(ast[key])) {
19+
if (ast[key].some((item) => hasModifyingStatement(item))) {
20+
return true;
2021
}
22+
} else if (hasModifyingStatement(ast[key])) {
23+
return true;
24+
}
2125
}
26+
}
2227

23-
return false;
28+
return false;
2429
}
2530

26-
export async function beforeQueryCache(sql: string, params?: any[], dataSource?: DataSource, dialect?: string): Promise<any | null> {
27-
// Currently we do not support caching queries that have dynamic parameters
28-
if (params?.length) return null
29-
if (dataSource?.source === Source.internal || !dataSource?.request.headers.has('X-Starbase-Cache')) return null
31+
export async function beforeQueryCache(opts: {
32+
sql: string;
33+
params?: unknown[];
34+
dataSource: DataSource;
35+
}): Promise<unknown | null> {
36+
const { sql, params = [], dataSource } = opts;
3037

31-
if (!dialect) dialect = 'sqlite'
32-
if (dialect.toLowerCase() === 'postgres') dialect = 'postgresql'
38+
// Currently we do not support caching queries that have dynamic parameters
39+
if (params.length) {
40+
return null;
41+
}
3342

34-
let ast = parser.astify(sql, { database: dialect });
35-
if (hasModifyingStatement(ast)) return null
36-
37-
const fetchCacheStatement = `SELECT timestamp, ttl, query, results FROM tmp_cache WHERE query = ?`
38-
const result = await dataSource.internalConnection?.durableObject.executeQuery(fetchCacheStatement, [sql], false) as any[];
39-
40-
if (result?.length) {
41-
const { timestamp, ttl, results } = result[0];
42-
const expirationTime = new Date(timestamp).getTime() + (ttl * 1000);
43-
44-
if (Date.now() < expirationTime) {
45-
return JSON.parse(results)
46-
}
43+
// If it's an internal request, or cache is not enabled, return null.
44+
if (dataSource.source === "internal" || !dataSource.cache) {
45+
return null;
46+
}
47+
48+
const dialect =
49+
dataSource.source === "external" ? dataSource.external!.dialect : "sqlite";
50+
51+
let ast = parser.astify(sql, { database: dialect });
52+
if (hasModifyingStatement(ast)) return null;
53+
54+
const fetchCacheStatement = `SELECT timestamp, ttl, query, results FROM tmp_cache WHERE query = ?`;
55+
56+
type QueryResult = {
57+
timestamp: string;
58+
ttl: number;
59+
results: string;
60+
};
61+
62+
const result = await dataSource.rpc.executeQuery({
63+
sql: fetchCacheStatement,
64+
params: [sql],
65+
}) as any[];
66+
67+
if (result?.length) {
68+
const { timestamp, ttl, results } = result[0] as QueryResult;
69+
const expirationTime = new Date(timestamp).getTime() + ttl * 1000;
70+
71+
if (Date.now() < expirationTime) {
72+
return JSON.parse(results);
4773
}
74+
}
4875

49-
return null
76+
return null;
5077
}
5178

5279
// Serialized RPC arguemnts are limited to 1MiB in size at the moment for Cloudflare
@@ -55,36 +82,57 @@ export async function beforeQueryCache(sql: string, params?: any[], dataSource?:
5582
// to look into include using Cloudflare Cache but need to find a good way to cache the
5683
// response in a safe way for our use case. Another option is another service for queues
5784
// or another way to ingest it directly to the Durable Object.
58-
export async function afterQueryCache(sql: string, params: any[] | undefined, result: any, dataSource?: DataSource, dialect?: string) {
59-
// Currently we do not support caching queries that have dynamic parameters
60-
if (params?.length) return;
61-
if (dataSource?.source === Source.internal || !dataSource?.request.headers.has('X-Starbase-Cache')) return null
62-
63-
try {
64-
if (!dialect) dialect = 'sqlite'
65-
if (dialect.toLowerCase() === 'postgres') dialect = 'postgresql'
66-
67-
let ast = parser.astify(sql, { database: dialect });
68-
69-
// If any modifying query exists within our SQL statement then we shouldn't proceed
70-
if (hasModifyingStatement(ast)) return;
71-
72-
const timestamp = Date.now();
73-
const results = JSON.stringify(result);
74-
75-
const exists = await dataSource.internalConnection?.durableObject.executeQuery(
76-
'SELECT 1 FROM tmp_cache WHERE query = ? LIMIT 1',
77-
[sql],
78-
false
79-
) as any[];
80-
81-
const query = exists?.length
82-
? { sql: 'UPDATE tmp_cache SET timestamp = ?, results = ? WHERE query = ?', params: [timestamp, results, sql] }
83-
: { sql: 'INSERT INTO tmp_cache (timestamp, ttl, query, results) VALUES (?, ?, ?, ?)', params: [timestamp, 60, sql, results] };
84-
85-
await dataSource.internalConnection?.durableObject.executeQuery(query.sql, query.params, false);
86-
} catch (error) {
87-
console.error('Error in cache operation:', error);
88-
return;
89-
}
90-
}
85+
export async function afterQueryCache(opts: {
86+
sql: string;
87+
params: unknown[] | undefined;
88+
result: unknown;
89+
dataSource: DataSource;
90+
}) {
91+
const { sql, params, result, dataSource } = opts;
92+
93+
// Currently we do not support caching queries that have dynamic parameters
94+
if (params?.length) return;
95+
96+
// If it's an internal request, or cache is not enabled, return null.
97+
if (dataSource.source === "internal" || !dataSource.cache) {
98+
return null;
99+
}
100+
101+
try {
102+
const dialect =
103+
dataSource.source === "external"
104+
? dataSource.external!.dialect
105+
: "sqlite";
106+
107+
let ast = parser.astify(sql, { database: dialect });
108+
109+
// If any modifying query exists within our SQL statement then we shouldn't proceed
110+
if (hasModifyingStatement(ast)) return;
111+
112+
const timestamp = Date.now();
113+
const results = JSON.stringify(result);
114+
115+
const exists = await dataSource.rpc.executeQuery({
116+
sql: "SELECT 1 FROM tmp_cache WHERE query = ? LIMIT 1",
117+
params: [sql],
118+
}) as QueryResult[];
119+
120+
const query = exists?.length
121+
? {
122+
sql: "UPDATE tmp_cache SET timestamp = ?, results = ? WHERE query = ?",
123+
params: [timestamp, results, sql],
124+
}
125+
: {
126+
sql: "INSERT INTO tmp_cache (timestamp, ttl, query, results) VALUES (?, ?, ?, ?)",
127+
params: [timestamp, dataSource.cacheTTL ?? 60, sql, results],
128+
};
129+
130+
await dataSource.rpc.executeQuery({
131+
sql: query.sql,
132+
params: query.params,
133+
});
134+
} catch (error) {
135+
console.error("Error in cache operation:", error);
136+
return;
137+
}
138+
}

0 commit comments

Comments
 (0)