Support schema in LiteREST for external connections

Brayden · Brayden · commit 400b9afb202b · 2024-11-22T11:51:24.000-05:00
diff --git a/src/handler.ts b/src/handler.ts
@@ -14,12 +14,14 @@ import { importTableFromCsvRoute } from "./import/csv";
 export class Handler {
     private liteREST?: LiteREST;
     private dataSource?: DataSource;
+    private env?: Env;
 
     constructor() { }
 
     public async handle(request: Request, dataSource: DataSource, env: Env): Promise<Response> {
         this.dataSource = dataSource;
         this.liteREST = new LiteREST(dataSource, env);
+        this.env = env;
         const url = new URL(request.url);
 
         if (request.method === 'POST' && url.pathname === '/query/raw') {
@@ -76,7 +78,7 @@ export class Handler {
             if (this.dataSource.source === Source.external) {
                 return createResponse(undefined, 'Function is only available for internal data source.', 400);
             }
-            
+
             const tableName = url.pathname.split('/').pop();
             if (!tableName) {
                 return createResponse(undefined, 'Table name is required', 400);
@@ -111,15 +113,15 @@ export class Handler {
                     return { sql, params };
                 });
 
-                const response = await executeTransaction(queries, isRaw, this.dataSource);
+                const response = await executeTransaction(queries, isRaw, this.dataSource, this.env);
                 return createResponse(response, undefined, 200);
             } else if (typeof sql !== 'string' || !sql.trim()) {
                 return createResponse(undefined, 'Invalid or empty "sql" field.', 400);
             } else if (params !== undefined && !Array.isArray(params) && typeof params !== 'object') {
                 return createResponse(undefined, 'Invalid "params" field. Must be an array or object.', 400);
             }
 
-            const response = await executeQuery(sql, params, isRaw, this.dataSource);
+            const response = await executeQuery(sql, params, isRaw, this.dataSource, this.env);
             return createResponse(response, undefined, 200);
         } catch (error: any) {
             console.error('Query Route Error:', error);
diff --git a/src/index.ts b/src/index.ts
@@ -123,7 +123,7 @@ export default {
             },
             externalConnection: {
                 outerbaseApiKey: env.OUTERBASE_API_KEY ?? ''
-            },
+            }
         };
 
         const response = await new Handler().handle(request, dataSource, env);
diff --git a/src/literest/index.ts b/src/literest/index.ts
@@ -29,23 +29,34 @@ export class LiteREST {
      * @param tableName - The name of the table.
      * @returns An array of primary key column names.
      */
-    private async getPrimaryKeyColumns(tableName: string): Promise<string[]> {
+    private async getPrimaryKeyColumns(tableName: string, schemaName?: string): Promise<string[]> {
         let query = `PRAGMA table_info(${tableName});`;
-
+    
         if (this.dataSource.source === Source.external) {
             if (this.env.EXTERNAL_DB_TYPE?.toLowerCase() === "postgres") {
-                query = `SELECT * FROM information_schema.table_constraints WHERE table_name = '${tableName}' AND constraint_type = 'PRIMARY KEY';`;
+                query = `
+                    SELECT kcu.column_name AS name 
+                    FROM information_schema.table_constraints tc 
+                    JOIN information_schema.key_column_usage kcu 
+                    ON tc.constraint_name = kcu.constraint_name
+                    AND tc.table_schema = kcu.table_schema 
+                    AND tc.table_name = kcu.table_name 
+                    WHERE tc.constraint_type = 'PRIMARY KEY'
+                    AND tc.table_name = '${tableName}' 
+                    AND tc.table_schema = '${schemaName ?? 'public'}';`;
             } else if (this.env.EXTERNAL_DB_TYPE?.toLowerCase() === "mysql") {
-                query = `SELECT COLUMN_NAME FROM information_schema.key_column_usage 
-                        WHERE table_name = '${tableName}' 
-                        AND constraint_name = 'PRIMARY'
-                        AND table_schema = DATABASE();`;
+                query = `
+                    SELECT COLUMN_NAME AS name FROM information_schema.key_column_usage 
+                    WHERE table_name = '${tableName}' 
+                    AND constraint_name = 'PRIMARY'
+                    AND table_schema = ${schemaName ?? 'DATABASE()'};
+                `;
             }
         }
 
-        const schemaInfo = (await executeQuery(query, [], false, this.dataSource)) as any[];
+        const schemaInfo = (await executeQuery(query, this.dataSource.source === Source.external ? {} : [], false, this.dataSource, this.env)) as any[];
         const pkColumns = schemaInfo
-            .filter(col => typeof col.pk === 'number' && col.pk > 0 && col.name !== null)
+            // .filter(col => typeof col.pk === 'number' && col.pk > 0 && col.name !== null)
             .map(col => col.name as string);
         return pkColumns;
     }
@@ -90,7 +101,7 @@ export class LiteREST {
         const conditions: string[] = [];
         const params: any[] = [];
 
-        if (pkColumns.length === 1) {
+        if (pkColumns?.length === 1) {
             const pk = pkColumns[0];
             const pkValue = id || data[pk] || searchParams.get(pk);
             if (!pkValue) {
@@ -118,7 +129,7 @@ export class LiteREST {
      * @param queries - The operations to execute.
      */
     private async executeOperation(queries: { sql: string, params: any[] }[]): Promise<{ result?: any, error?: string | undefined, status: number }> {
-        const results: any[] = (await executeTransaction(queries, false, this.dataSource)) as any[];
+        const results: any[] = (await executeTransaction(queries, false, this.dataSource, this.env)) as any[];
         return { result: results?.length > 0 ? results[0] : undefined, status: 200 };
     }
 
@@ -128,20 +139,20 @@ export class LiteREST {
      * @returns The response to the request.
      */
     async handleRequest(request: Request): Promise<Response> {
-        const { method, tableName, id, searchParams, body } = await this.parseRequest(request);
+        const { method, tableName, schemaName, id, searchParams, body } = await this.parseRequest(request);
 
         try {
             switch (method) {
                 case 'GET':
-                    return await this.handleGet(tableName, id, searchParams);
+                    return await this.handleGet(tableName, schemaName, id, searchParams);
                 case 'POST':
-                    return await this.handlePost(tableName, body);
+                    return await this.handlePost(tableName, schemaName, body);
                 case 'PATCH':
-                    return await this.handlePatch(tableName, id, body);
+                    return await this.handlePatch(tableName, schemaName, id, body);
                 case 'PUT':
-                    return await this.handlePut(tableName, id, body);
+                    return await this.handlePut(tableName, schemaName, id, body);
                 case 'DELETE':
-                    return await this.handleDelete(tableName, id);
+                    return await this.handleDelete(tableName, schemaName, id);
                 default:
                     return createResponse(undefined, 'Method not allowed', 405);
             }
@@ -156,7 +167,7 @@ export class LiteREST {
      * @param request - The incoming request.
      * @returns An object containing the method, table name, id, search parameters, and body.
      */
-    private async parseRequest(request: Request): Promise<{ method: string, tableName: string, id?: string, searchParams: URLSearchParams, body?: any }> {
+    private async parseRequest(request: Request): Promise<{ method: string, tableName: string, schemaName: string | undefined, id?: string, searchParams: URLSearchParams, body?: any }> {
         const liteRequest = new Request(request.url.replace('/rest', ''), request);
         const url = new URL(liteRequest.url);
         const pathParts = url.pathname.split('/').filter(Boolean);
@@ -165,24 +176,26 @@ export class LiteREST {
             throw new Error('Invalid route');
         }
 
-        const tableName = this.sanitizeIdentifier(pathParts[0]);
-        const id = pathParts[1];
+        const tableName = this.sanitizeIdentifier(pathParts.length === 1 ? pathParts[0] : pathParts[1]);
+        const schemaName = pathParts.length === 1 ? undefined : this.sanitizeIdentifier(pathParts[0])
+        const id = pathParts.length === 1 ? pathParts[1] : pathParts[2];
         const body = ['POST', 'PUT', 'PATCH'].includes(liteRequest.method) ? await liteRequest.json() : undefined;
 
         return {
             method: liteRequest.method,
             tableName,
+            schemaName,
             id,
             searchParams: url.searchParams,
             body
         };
     }
 
-    private async buildSelectQuery(tableName: string, id: string | undefined, searchParams: URLSearchParams): Promise<{ query: string, params: any[] }> {
-        let query = `SELECT * FROM ${tableName}`;
+    private async buildSelectQuery(tableName: string, schemaName: string | undefined, id: string | undefined, searchParams: URLSearchParams): Promise<{ query: string, params: any[] }> {
+        let query = `SELECT * FROM ${schemaName ? `${schemaName}.` : ''}${tableName}`;
         const params: any[] = [];
         const conditions: string[] = [];
-        const pkColumns = await this.getPrimaryKeyColumns(tableName);
+        const pkColumns = await this.getPrimaryKeyColumns(tableName, schemaName);
         const { conditions: pkConditions, params: pkParams, error } = this.getPrimaryKeyConditions(pkColumns, id, {}, searchParams);
     
         if (!error) {
@@ -249,8 +262,8 @@ export class LiteREST {
         return { query, params };
     }
 
-    private async handleGet(tableName: string, id: string | undefined, searchParams: URLSearchParams): Promise<Response> {
-        const { query, params } = await this.buildSelectQuery(tableName, id, searchParams);
+    private async handleGet(tableName: string, schemaName: string | undefined, id: string | undefined, searchParams: URLSearchParams): Promise<Response> {
+        const { query, params } = await this.buildSelectQuery(tableName, schemaName, id, searchParams);
 
         try {
             const response = await this.executeOperation([{ sql: query, params }]);
@@ -262,7 +275,7 @@ export class LiteREST {
         }
     } 
 
-    private async handlePost(tableName: string, data: any): Promise<Response> {
+    private async handlePost(tableName: string, schemaName: string | undefined, data: any): Promise<Response> {
         if (!this.isDataValid(data)) {
             console.error('Invalid data format for POST:', data);
             return createResponse(undefined, 'Invalid data format', 400);
@@ -293,8 +306,8 @@ export class LiteREST {
         }
     }
 
-    private async handlePatch(tableName: string, id: string | undefined, data: any): Promise<Response> {
-        const pkColumns = await this.getPrimaryKeyColumns(tableName);
+    private async handlePatch(tableName: string, schemaName: string | undefined, id: string | undefined, data: any): Promise<Response> {
+        const pkColumns = await this.getPrimaryKeyColumns(tableName, schemaName);
 
         const { conditions: pkConditions, params: pkParams, error } = this.getPrimaryKeyConditions(pkColumns, id, data, new URLSearchParams());
 
@@ -342,8 +355,8 @@ export class LiteREST {
         }
     }
     
-    private async handlePut(tableName: string, id: string | undefined, data: any): Promise<Response> {
-        const pkColumns = await this.getPrimaryKeyColumns(tableName);
+    private async handlePut(tableName: string, schemaName: string | undefined, id: string | undefined, data: any): Promise<Response> {
+        const pkColumns = await this.getPrimaryKeyColumns(tableName, schemaName);
 
         const { conditions: pkConditions, params: pkParams, error } = this.getPrimaryKeyConditions(pkColumns, id, data, new URLSearchParams());
 
@@ -383,8 +396,8 @@ export class LiteREST {
         }
     }
 
-    private async handleDelete(tableName: string, id: string | undefined): Promise<Response> {
-        const pkColumns = await this.getPrimaryKeyColumns(tableName);
+    private async handleDelete(tableName: string, schemaName: string | undefined, id: string | undefined): Promise<Response> {
+        const pkColumns = await this.getPrimaryKeyColumns(tableName, schemaName);
 
         const { conditions: pkConditions, params: pkParams, error } = this.getPrimaryKeyConditions(pkColumns, id, {}, new URLSearchParams());
 
diff --git a/src/operation.ts b/src/operation.ts
@@ -1,4 +1,5 @@
 import { DataSource } from '.';
+import { Env } from './'
 
 export type OperationQueueItem = {
     queries: { sql: string; params?: any[] }[];
@@ -19,82 +20,97 @@ export type RawQueryResponse = {
 
 export type QueryResponse = any[] | RawQueryResponse;
 
-async function afterQuery(sql: string, result: any, isRaw: boolean, dataSource?: DataSource): Promise<any> {
+async function afterQuery(sql: string, result: any, isRaw: boolean, dataSource?: DataSource, env?: Env): Promise<any> {
     // ## DO NOT REMOVE: TEMPLATE AFTER QUERY HOOK ##
 
     return result;
 }
 
-export async function executeQuery(sql: string, params: any | undefined, isRaw: boolean, dataSource?: DataSource): Promise<QueryResponse> {
+function cleanseQuery(sql: string): string {
+    return sql.replaceAll('\n', ' ')
+}
+
+ // NOTE: This is a temporary stop-gap solution to connect to external data sources. Outerbase offers
+ // an API to handle connecting to a large number of database types in a secure manner. However, the
+ // goal here is to optimize on query latency from your data sources by connecting to them directly.
+ // An upcoming update will move the Outerbase SDK to be used in StarbaseDB so this service can connect
+// to those database types without being required to funnel requests through the Outerbase API.
+async function executeExternalQuery(sql: string, params: any, isRaw: boolean, dataSource: DataSource, env?: Env): Promise<any> {
+    if (!dataSource.externalConnection) {
+        throw new Error('External connection not found.');
+    }
+
+    // Convert params from array to object if needed
+    let convertedParams = params;
+    if (Array.isArray(params)) {
+        let paramIndex = 0;
+        convertedParams = params.reduce((acc, value, index) => ({
+            ...acc,
+            [`param${index}`]: value
+        }), {});
+        sql = sql.replace(/\?/g, () => `:param${paramIndex++}`);
+    }
+
+    const API_URL = 'https://app.outerbase.com';
+    const response = await fetch(`${API_URL}/api/v1/ezql/raw`, {
+        method: 'POST',
+        headers: {
+            'Content-Type': 'application/json',
+            'X-Source-Token': dataSource.externalConnection.outerbaseApiKey,
+        },
+        body: JSON.stringify({
+            query: cleanseQuery(sql),
+            params: convertedParams,
+        }),
+    });
+
+    const results: any = await response.json();
+    const items = results.response.results?.items;
+    return await afterQuery(sql, items, isRaw, dataSource, env);
+}
+
+export async function executeQuery(sql: string, params: any | undefined, isRaw: boolean, dataSource?: DataSource, env?: Env): Promise<QueryResponse> {
     if (!dataSource) {
         console.error('Data source not found.')
         return []
     }
 
     if (dataSource.source === 'internal') {
         const response = await dataSource.internalConnection?.durableObject.executeQuery(sql, params, isRaw);
-        return response ?? [];
+        return await afterQuery(sql, response, isRaw, dataSource, env);
     } else {
-        if (!dataSource.externalConnection) {
-            throw new Error('External connection not found.');
-        }
-
-        const API_URL = 'https://app.outerbase.com'
-        const response = await fetch(`${API_URL}/api/v1/ezql/raw`, {
-            method: 'POST',
-            headers: {
-                'Content-Type': 'application/json',
-                'X-Source-Token': dataSource.externalConnection.outerbaseApiKey,
-            },
-            body: JSON.stringify({
-                query: sql,
-                // Params does not support arrays, so we ensure we only pass them an object.
-                params: Array.isArray(params) ? {} : params,
-            }),
-        })
-
-        let results: any = await response.json();
-        let items = results.response.results?.items;
-        return afterQuery(sql, items, isRaw, dataSource);
-    } 
+        return executeExternalQuery(sql, params, isRaw, dataSource, env);
+    }
 }
 
-export async function executeTransaction(queries: { sql: string; params?: any[] }[], isRaw: boolean, dataSource?: DataSource): Promise<QueryResponse> {
+export async function executeTransaction(queries: { sql: string; params?: any[] }[], isRaw: boolean, dataSource?: DataSource, env?: Env): Promise<QueryResponse> {
     if (!dataSource) {
         console.error('Data source not found.')
         return []
     }
     
     if (dataSource.source === 'internal') {
-        const response = await dataSource.internalConnection?.durableObject.executeTransaction(queries, isRaw);
-        return response ?? [];
-    } else {
-        if (!dataSource.externalConnection) {
-            throw new Error('External connection not found.');
-        }
+        const results: any[] = [];
 
-        const API_URL = 'https://app.outerbase.com';
+        for (const query of queries) {
+            const result = await dataSource.internalConnection?.durableObject.executeTransaction(queries, isRaw);
+            if (result) {
+                const processedResults = await Promise.all(
+                    result.map(r => afterQuery(query.sql, r, isRaw, dataSource, env))
+                );
+                results.push(...processedResults);
+            }
+        }
+        
+        return results;
+    } else {
         const results = [];
 
         for (const query of queries) {
-            const response = await fetch(`${API_URL}/api/v1/ezql/raw`, {
-                method: 'POST',
-                headers: {
-                    'Content-Type': 'application/json',
-                    'X-Source-Token': dataSource.externalConnection.outerbaseApiKey,
-                },
-                body: JSON.stringify({
-                    query: query.sql,
-                    // Params does not support arrays, so we ensure we only pass them an object
-                    params: Array.isArray(query.params) ? {} : query.params,
-                }),
-            });
-
-            const result: any = await response.json();
-            const items = result.response.results?.items;
-            results.push(afterQuery(query.sql, items, isRaw, dataSource));
+            const result = await executeExternalQuery(query.sql, query.params, isRaw, dataSource, env);
+            results.push(result);
         }
-
+        
         return results;
     }
 }
diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts
@@ -3,7 +3,5 @@
 interface Env {
 	AUTHORIZATION_TOKEN: "ABC123";
 	REGION: "auto";
-	EXTERNAL_DB_TYPE: "postgres";
-	OUTERBASE_API_KEY: "";
 	DATABASE_DURABLE_OBJECT: DurableObjectNamespace<import("./src/index").DatabaseDurableObject>;
 }
diff --git a/wrangler.toml b/wrangler.toml

Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,5 @@`
`3`	`3`	`interface Env {`
`4`	`4`	`AUTHORIZATION_TOKEN: "ABC123";`
`5`	`5`	`REGION: "auto";`
`6`		`- EXTERNAL_DB_TYPE: "postgres";`
`7`		`- OUTERBASE_API_KEY: "";`
`8`	`6`	`DATABASE_DURABLE_OBJECT: DurableObjectNamespace<import("./src/index").DatabaseDurableObject>;`
`9`	`7`	`}`