address CR and add tests

Author: bogdantruta-firebolt

src/connection/base.ts

@@ -179,10 +179,10 @@ export abstract class Connection {
 
   abstract cancelAsyncQuery(token: string): Promise<void>;
 
-  protected async prepareAndExecuteQuery(
+  protected prepareQuery(
     query: string,
     executeQueryOptions: ExecuteQueryOptions
-  ): Promise<{ formattedQuery: string; response: Response }> {
+  ): { formattedQuery: string; setKey: string } {
     executeQueryOptions.response = {
       ...defaultResponseSettings,
       ...(executeQueryOptions.response ?? {})
@@ -205,13 +205,14 @@ export abstract class Connection {
       );
     }
 
-    return await this.executeQuery(formattedQuery, executeQueryOptions, setKey);
+    return { formattedQuery, setKey };
   }
 
   protected async executeQuery(
     formattedQuery: string,
     executeQueryOptions: ExecuteQueryOptions,
-    setKey = ""
+    setKey = "",
+    returnResponse = false
   ) {
     const { httpClient } = this.context;
 
@@ -229,7 +230,14 @@ export abstract class Connection {
     try {
       const response = await request.ready();
       await this.processHeaders(response.headers);
-      return { formattedQuery, response };
+      if (returnResponse) {
+        return { response, text: "" };
+      }
+
+      const text = await response.text();
+      await this.throwErrorIfErrorBody(text, response);
+
+      return { response, text };
     } catch (error) {
       // In case it was a set query, remove set parameter if query fails
       if (setKey.length > 0) {
@@ -241,17 +249,42 @@ export abstract class Connection {
     }
   }
 
+  protected async prepareAndExecuteQuery(
+    query: string,
+    executeQueryOptions: ExecuteQueryOptions,
+    returnResponse = false
+  ): Promise<{
+    formattedQuery: string;
+    response: Response;
+    text: string;
+  }> {
+    const { formattedQuery, setKey } = this.prepareQuery(
+      query,
+      executeQueryOptions
+    );
+    const { response, text } = await this.executeQuery(
+      formattedQuery,
+      executeQueryOptions,
+      setKey,
+      returnResponse
+    );
+
+    return {
+      formattedQuery,
+      response,
+      text
+    };
+  }
+
   async execute(
     query: string,
     executeQueryOptions: ExecuteQueryOptions = {}
   ): Promise<Statement> {
-    const { formattedQuery, response } = await this.prepareAndExecuteQuery(
+    const { formattedQuery, text } = await this.prepareAndExecuteQuery(
       query,
       executeQueryOptions
     );
 
-    const text = await response.text();
-    await this.throwErrorIfErrorBody(text, response);
     return new Statement(this.context, {
       query: formattedQuery,
       text,

src/connection/connection_v2.ts

@@ -91,13 +91,11 @@ export class ConnectionV2 extends BaseConnection {
       // can't have an async set query
       throw new Error("SET statements cannot be executed asynchronously.");
     }
-    const { formattedQuery, response } = await this.prepareAndExecuteQuery(
+    const { formattedQuery, text } = await this.prepareAndExecuteQuery(
       query,
       asyncExecuteQueryOptions
     );
 
-    const text = await response.text();
-    await this.throwErrorIfErrorBody(text, response);
     return new AsyncStatement(this.context, {
       query: formattedQuery,
       text,
@@ -109,13 +107,17 @@ export class ConnectionV2 extends BaseConnection {
     query: string,
     executeQueryOptions: ExecuteQueryOptions = {}
   ): Promise<StreamStatement> {
-    const { response } = await this.prepareAndExecuteQuery(query, {
-      ...executeQueryOptions,
-      settings: {
-        ...executeQueryOptions?.settings,
-        output_format: OutputFormat.JSON_LINES
-      }
-    });
+    const { response } = await this.prepareAndExecuteQuery(
+      query,
+      {
+        ...executeQueryOptions,
+        settings: {
+          ...executeQueryOptions?.settings,
+          output_format: OutputFormat.JSON_LINES
+        }
+      },
+      true
+    );
 
     return new StreamStatement({
       response,
@@ -137,48 +139,36 @@ export class ConnectionV2 extends BaseConnection {
     query: string,
     executeQueryOptions: ExecuteQueryOptions
   ): Promise<Statement> {
-    this.checkParameters(query, executeQueryOptions.namedParameters);
-    const { response } = await this.executeQuery(
+    const { text } = await this.executeQuery(
       query,
       this.getExecuteQueryOptionsForPreparedStatement(executeQueryOptions)
     );
 
-    const text = await response.text();
-    await this.throwErrorIfErrorBody(text, response);
     return new Statement(this.context, {
       query: query,
       text,
       executeQueryOptions
     });
   }
 
-  private checkParameters(
-    query: string,
-    namedParameters: Record<string, unknown> | undefined
-  ) {
-    const fireboltNumericParameters =
-      this.queryFormatter.getFireboltNumericParameters(query);
-    const providedParameters = Object.keys(namedParameters || {});
-    if (fireboltNumericParameters.length != providedParameters.length) {
-      throw new Error(
-        `Number of parameters in the query (${fireboltNumericParameters.length}) does not match the number of named parameters provided (${providedParameters.length}).`
-      );
-    }
-    fireboltNumericParameters.forEach(param => {
-      if (!providedParameters.includes(param)) {
-        throw new Error(
-          `Parameter "${param}" is missing from the named parameters.`
-        );
-      }
-    });
-  }
-
   private getExecuteQueryOptionsForPreparedStatement(
     executeQueryOptions: ExecuteQueryOptions
   ): ExecuteQueryOptions {
-    const queryParameters = Object.entries(
-      executeQueryOptions.namedParameters || {}
-    ).map(([key, value]) => ({ name: key, value }));
+    let queryParameters;
+    if (!executeQueryOptions.parameters) {
+      queryParameters = Object.entries(
+        executeQueryOptions.namedParameters || {}
+      ).map(([key, value]) => ({ name: key, value }));
+    } else if (!executeQueryOptions.namedParameters) {
+      queryParameters = executeQueryOptions.parameters.map((value, index) => ({
+        name: `$${index + 1}`,
+        value
+      }));
+    } else {
+      throw new Error(
+        "Server-side prepared statement can only use either parameters or namedParameters"
+      );
+    }
 
     return {
       settings: {

src/formatter/base.ts

@@ -220,20 +220,6 @@ export abstract class QueryFormatter {
     return query;
   }
 
-  getFireboltNumericParameters(query: string): string[] {
-    const regex = /'[^']*'|"[^"]*"|\/\*[\s\S]*?\*\/|--.*|(\$[0-9]+)/g;
-
-    let match;
-    const params = [];
-
-    while ((match = regex.exec(query)) !== null) {
-      if (match[1]) {
-        params.push(match[1]);
-      }
-    }
-    return params;
-  }
-
   isSetStatement(query: string): boolean {
     return query.trim().toLowerCase().startsWith(SET_PREFIX);
   }