Merge pull request #42 from Brayden/bwilmoth/before-query

Brayden · web-flow · commit 0b7cc3cb3e20 · 2024-12-01T13:07:13.000-05:00
beforeQuery functional hook
diff --git a/src/handler.ts b/src/handler.ts
@@ -1,4 +1,4 @@
-import { DataSource, Source } from ".";
+import { DataSource, Source } from "./types";
 import { LiteREST } from "./literest";
 import { executeQuery, executeTransaction } from "./operation";
 import { createResponse, QueryRequest, QueryTransactionRequest } from "./utils";
diff --git a/src/index.ts b/src/index.ts
@@ -1,7 +1,7 @@
 import { createResponse } from './utils';
 import handleStudioRequest from "./studio";
 import { Handler } from "./handler";
-import { QueryResponse } from "./operation";
+import { DatabaseStub, DataSource, RegionLocationHint, Source } from './types';
 export { DatabaseDurableObject } from './do'; 
 
 const DURABLE_OBJECT_ID = 'sql-durable-object';
@@ -37,43 +37,6 @@ export interface Env {
     // ## DO NOT REMOVE: TEMPLATE INTERFACE ##
 }
 
-export enum Source {
-    internal = 'internal',  // Durable Object's SQLite instance
-    external = 'external'   // External data source (e.g. Outerbase)
-}
-
-export type DataSource = {
-    source: Source;
-    request: Request;
-    internalConnection?: InternalConnection;
-    externalConnection?: {
-        outerbaseApiKey: string;
-    };
-}
-
-interface InternalConnection {
-    durableObject: DatabaseStub;
-}
-
-type DatabaseStub = DurableObjectStub & {
-    fetch: (init?: RequestInit | Request) => Promise<Response>;
-    executeQuery(sql: string, params: any[] | undefined, isRaw: boolean): QueryResponse;
-    executeTransaction(queries: { sql: string; params?: any[] }[], isRaw: boolean): any[];
-};
-
-enum RegionLocationHint {
-    AUTO = 'auto',
-    WNAM = 'wnam', // Western North America
-    ENAM = 'enam', // Eastern North America
-    SAM = 'sam', // South America
-    WEUR = 'weur', // Western Europe
-    EEUR = 'eeur', // Eastern Europe
-    APAC = 'apac', // Asia Pacific
-    OC = 'oc', // Oceania
-    AFR = 'afr', // Africa
-    ME = 'me', // Middle East
-}
-
 export default {
 	/**
 	 * This is the standard fetch handler for a Cloudflare Worker
@@ -84,67 +47,73 @@ export default {
 	 * @returns The response to be sent back to the client
 	 */
 	async fetch(request, env, ctx): Promise<Response> {
-        const url = new URL(request.url);
-        const pathname = url.pathname;
-        const isWebSocket = request.headers.get("Upgrade") === "websocket";
-
-        /**
-         * If the request is a GET request to the /studio endpoint, we can handle the request
-         * directly in the Worker to avoid the need to deploy a separate Worker for the Studio.
-         * Studio provides a user interface to interact with the SQLite database in the Durable
-         * Object.
-         */
-        if (env.STUDIO_USER && env.STUDIO_PASS && request.method === 'GET' && pathname === '/studio') {
-            return handleStudioRequest(request, {
-                username: env.STUDIO_USER,
-                password: env.STUDIO_PASS, 
-                apiToken: env.AUTHORIZATION_TOKEN
-            });
-        }
+        try {
+            const url = new URL(request.url);
+            const pathname = url.pathname;
+            const isWebSocket = request.headers.get("Upgrade") === "websocket";
 
-        /**
-         * Prior to proceeding to the Durable Object, we can perform any necessary validation or
-         * authorization checks here to ensure the request signature is valid and authorized to
-         * interact with the Durable Object.
-         */
-        if (request.headers.get('Authorization') !== `Bearer ${env.AUTHORIZATION_TOKEN}` && !isWebSocket) {
-            return createResponse(undefined, 'Unauthorized request', 401)
-        } else if (isWebSocket) {
             /**
-             * Web socket connections cannot pass in an Authorization header into their requests,
-             * so we can use a query parameter to validate the connection.
+             * If the request is a GET request to the /studio endpoint, we can handle the request
+             * directly in the Worker to avoid the need to deploy a separate Worker for the Studio.
+             * Studio provides a user interface to interact with the SQLite database in the Durable
+             * Object.
              */
-            const token = url.searchParams.get('token');
-
-            if (token !== env.AUTHORIZATION_TOKEN) {
-                return new Response('WebSocket connections are not supported at this endpoint.', { status: 440 });
+            if (env.STUDIO_USER && env.STUDIO_PASS && request.method === 'GET' && pathname === '/studio') {
+                return handleStudioRequest(request, {
+                    username: env.STUDIO_USER,
+                    password: env.STUDIO_PASS, 
+                    apiToken: env.AUTHORIZATION_TOKEN
+                });
             }
-        }
-
-        /**
-         * Retrieve the Durable Object identifier from the environment bindings and instantiate a
-         * Durable Object stub to interact with the Durable Object.
-         */
-        const region = env.REGION ?? RegionLocationHint.AUTO;
-        const id: DurableObjectId = env.DATABASE_DURABLE_OBJECT.idFromName(DURABLE_OBJECT_ID);
-        const stub = region !== RegionLocationHint.AUTO ? env.DATABASE_DURABLE_OBJECT.get(id, { locationHint: region as DurableObjectLocationHint }) : env.DATABASE_DURABLE_OBJECT.get(id);
 
-        const source: Source = request.headers.get('X-Starbase-Source') as Source ?? url.searchParams.get('source') as Source ?? 'internal';
-        const dataSource: DataSource = {
-            source: source,
-            request: request.clone(),
-            internalConnection: {
-                durableObject: stub as unknown as DatabaseStub,
-            },
-            externalConnection: {
-                outerbaseApiKey: env.OUTERBASE_API_KEY ?? ''
+            /**
+             * Prior to proceeding to the Durable Object, we can perform any necessary validation or
+             * authorization checks here to ensure the request signature is valid and authorized to
+             * interact with the Durable Object.
+             */
+            if (request.headers.get('Authorization') !== `Bearer ${env.AUTHORIZATION_TOKEN}` && !isWebSocket) {
+                return createResponse(undefined, 'Unauthorized request', 401)
+            } else if (isWebSocket) {
+                /**
+                 * Web socket connections cannot pass in an Authorization header into their requests,
+                 * so we can use a query parameter to validate the connection.
+                 */
+                const token = url.searchParams.get('token');
+
+                if (token !== env.AUTHORIZATION_TOKEN) {
+                    return new Response('WebSocket connections are not supported at this endpoint.', { status: 440 });
+                }
             }
-        };
 
-        const response = await new Handler().handle(request, dataSource, env);
-
-        // ## DO NOT REMOVE: TEMPLATE ROUTING ##
-
-        return response;
+            /**
+             * Retrieve the Durable Object identifier from the environment bindings and instantiate a
+             * Durable Object stub to interact with the Durable Object.
+             */
+            const region = env.REGION ?? RegionLocationHint.AUTO;
+            const id: DurableObjectId = env.DATABASE_DURABLE_OBJECT.idFromName(DURABLE_OBJECT_ID);
+            const stub = region !== RegionLocationHint.AUTO ? env.DATABASE_DURABLE_OBJECT.get(id, { locationHint: region as DurableObjectLocationHint }) : env.DATABASE_DURABLE_OBJECT.get(id);
+
+            const source: Source = request.headers.get('X-Starbase-Source') as Source ?? url.searchParams.get('source') as Source ?? 'internal';
+            const dataSource: DataSource = {
+                source: source,
+                request: request.clone(),
+                internalConnection: {
+                    durableObject: stub as unknown as DatabaseStub,
+                },
+                externalConnection: {
+                    outerbaseApiKey: env.OUTERBASE_API_KEY ?? ''
+                }
+            };
+
+            const response = await new Handler().handle(request, dataSource, env);
+            return response;
+        } catch (error) {
+            // Return error response to client
+            return createResponse(
+                undefined, 
+                error instanceof Error ? error.message : 'An unexpected error occurred',
+                400
+            );
+        }
 	},
 } satisfies ExportedHandler<Env>;
diff --git a/src/literest/index.ts b/src/literest/index.ts
@@ -1,5 +1,5 @@
 import { createResponse } from '../utils';
-import { DataSource, Source } from "..";
+import { DataSource, Source } from "../types";
 import { executeQuery, executeTransaction } from "../operation";
 import { Env } from "../index"
 
diff --git a/src/operation.ts b/src/operation.ts
@@ -5,7 +5,7 @@ import { createClient as createTursoConnection } from '@libsql/client/web';
 
 // Import how we interact with the databases through the Outerbase SDK
 import { CloudflareD1Connection, MongoDBConnection, MySQLConnection, PostgreSQLConnection, StarbaseConnection, TursoConnection } from '@outerbase/sdk';
-import { DataSource } from '.';
+import { DataSource } from './types';
 import { Env } from './'
 import { MongoClient } from 'mongodb';
 
@@ -33,14 +33,53 @@ export type ConnectionDetails = {
     defaultSchema: string,
 }
 
+async function beforeQuery(sql: string, params?: any[], dataSource?: DataSource, env?: Env): Promise<{ sql: string, params?: any[] }> {
+    // ## DO NOT REMOVE: PRE QUERY HOOK ##
+    
+    return {
+        sql,
+        params
+    }
+}
+
 async function afterQuery(sql: string, result: any, isRaw: boolean, dataSource?: DataSource, env?: Env): Promise<any> {
-    // ## DO NOT REMOVE: TEMPLATE AFTER QUERY HOOK ##
+    result = isRaw ? transformRawResults(result, 'from') : result;
+
+    // ## DO NOT REMOVE: POST QUERY HOOK ##
     
-    return result;
+    return isRaw ? transformRawResults(result, 'to') : result;
 }
 
-function cleanseQuery(sql: string): string {
-    return sql.replaceAll('\n', ' ')
+function transformRawResults(result: any, direction: 'to' | 'from'): Record<string, any> {
+    if (direction === 'from') {
+        // Convert our result from the `raw` output to a traditional object
+        result = {
+            ...result,
+            rows: result.rows.map((row: any) => 
+                result.columns.reduce((obj: any, column: string, index: number) => {
+                    obj[column] = row[index];
+                    return obj;
+                }, {})
+            )
+        };
+
+        return result.rows
+    } else if (direction === 'to') {
+        // Convert our traditional object to the `raw` output format
+        const columns = Object.keys(result[0] || {});
+        const rows = result.map((row: any) => columns.map(col => row[col]));
+        
+        return {
+            columns,
+            rows,
+            meta: {
+                rows_read: rows.length,
+                rows_written: 0
+            }
+        };
+    }
+    
+    return result
 }
 
 // Outerbase API supports more data sources than can be supported via Cloudflare Workers. For those data
@@ -68,6 +107,7 @@ async function executeExternalQuery(sql: string, params: any, isRaw: boolean, da
         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',
@@ -76,14 +116,13 @@ async function executeExternalQuery(sql: string, params: any, isRaw: boolean, da
             'X-Source-Token': dataSource.externalConnection.outerbaseApiKey,
         },
         body: JSON.stringify({
-            query: cleanseQuery(sql),
+            query: sql.replaceAll('\n', ' '),
             params: convertedParams,
         }),
     });
 
     const results: any = await response.json();
-    const items = results.response.results?.items;
-    return await afterQuery(sql, items, isRaw, dataSource, env);
+    return results.response.results?.items;
 }
 
 export async function executeQuery(sql: string, params: any | undefined, isRaw: boolean, dataSource?: DataSource, env?: Env): Promise<QueryResponse> {
@@ -92,12 +131,16 @@ export async function executeQuery(sql: string, params: any | undefined, isRaw:
         return []
     }
 
+    const { sql: updatedSQL, params: updatedParams } = await beforeQuery(sql, params, dataSource, env)
+    let response;
+
     if (dataSource.source === 'internal') {
-        const response = await dataSource.internalConnection?.durableObject.executeQuery(sql, params, isRaw);
-        return await afterQuery(sql, response, isRaw, dataSource, env);
+        response = await dataSource.internalConnection?.durableObject.executeQuery(updatedSQL, updatedParams, isRaw);
     } else {
-        return executeExternalQuery(sql, params, isRaw, dataSource, env);
+        response = await executeExternalQuery(updatedSQL, updatedParams, isRaw, dataSource, env);
     }
+
+    return await afterQuery(updatedSQL, response, isRaw, dataSource, env);
 }
 
 export async function executeTransaction(queries: { sql: string; params?: any[] }[], isRaw: boolean, dataSource?: DataSource, env?: Env): Promise<QueryResponse> {
@@ -106,30 +149,14 @@ export async function executeTransaction(queries: { sql: string; params?: any[]
         return []
     }
     
-    if (dataSource.source === 'internal') {
-        const results: any[] = [];
-
-        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 = [];
+    const results = [];
 
-        for (const query of queries) {
-            const result = await executeExternalQuery(query.sql, query.params, isRaw, dataSource, env);
-            results.push(result);
-        }
-        
-        return results;
+    for (const query of queries) {
+        const result = await executeQuery(query.sql, query.params, isRaw, dataSource, env);
+        results.push(result);
     }
+    
+    return results;
 }
 
 async function createSDKPostgresConnection(env: Env): Promise<ConnectionDetails> {
@@ -246,5 +273,5 @@ export async function executeSDKQuery(sql: string, params: any | undefined, isRa
     await db.connect();
     const { data } = await db.raw(sql, params);
     
-    return await afterQuery(sql, data, isRaw, dataSource, env);
+    return data;
 }
diff --git a/src/types.ts b/src/types.ts
@@ -0,0 +1,38 @@
+import { QueryResponse } from "./operation";
+
+export enum Source {
+    internal = 'internal',  // Durable Object's SQLite instance
+    external = 'external'   // External data source (e.g. Outerbase)
+}
+
+export type DataSource = {
+    source: Source;
+    request: Request;
+    internalConnection?: InternalConnection;
+    externalConnection?: {
+        outerbaseApiKey: string;
+    };
+}
+
+export interface InternalConnection {
+    durableObject: DatabaseStub;
+}
+
+export type DatabaseStub = DurableObjectStub & {
+    fetch: (init?: RequestInit | Request) => Promise<Response>;
+    executeQuery(sql: string, params: any[] | undefined, isRaw: boolean): QueryResponse;
+    executeTransaction(queries: { sql: string; params?: any[] }[], isRaw: boolean): any[];
+};
+
+export enum RegionLocationHint {
+    AUTO = 'auto',
+    WNAM = 'wnam', // Western North America
+    ENAM = 'enam', // Eastern North America
+    SAM = 'sam', // South America
+    WEUR = 'weur', // Western Europe
+    EEUR = 'eeur', // Eastern Europe
+    APAC = 'apac', // Asia Pacific
+    OC = 'oc', // Oceania
+    AFR = 'afr', // Africa
+    ME = 'me', // Middle East
+}

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-import { DataSource, Source } from ".";`
	`1`	`+import { DataSource, Source } from "./types";`
`2`	`2`	`import { LiteREST } from "./literest";`
`3`	`3`	`import { executeQuery, executeTransaction } from "./operation";`
`4`	`4`	`import { createResponse, QueryRequest, QueryTransactionRequest } from "./utils";`