feat: Fir 29768 add support for set statements in node sdk (#83)

Author: stepansergeevitch

src/connection/base.ts

@@ -26,6 +26,7 @@ const allowedUpdateParameters = ["database"];
 const updateEndpointHeader = "Firebolt-Update-Endpoint";
 const resetSessionHeader = "Firebolt-Reset-Session";
 const immutableParameters = ["database", "account_id", "output_format"];
+const testConnectionQuery = "SELECT 1";
 
 export abstract class Connection {
   protected context: Context;
@@ -171,11 +172,21 @@ export abstract class Connection {
     };
 
     const { parameters, namedParameters } = executeQueryOptions;
-    const formattedQuery = queryFormatter.formatQuery(
-      query,
-      parameters,
-      namedParameters
-    );
+
+    let setKey = "",
+      setValue = "",
+      formattedQuery: string;
+    if (queryFormatter.isSetStatement(query)) {
+      [setKey, setValue] = queryFormatter.splitSetStatement(query);
+      this.parameters[setKey] = setValue;
+      formattedQuery = testConnectionQuery;
+    } else {
+      formattedQuery = queryFormatter.formatQuery(
+        query,
+        parameters,
+        namedParameters
+      );
+    }
 
     const body = formattedQuery;
     const url = this.getRequestUrl(executeQueryOptions);
@@ -197,6 +208,12 @@ export abstract class Connection {
         executeQueryOptions
       });
       return statement;
+    } catch (error) {
+      // In case it was a set query, remove set parameter if query fails
+      if (setKey.length > 0) {
+        delete this.parameters[setKey];
+      }
+      throw error;
     } finally {
       this.activeRequests.delete(request);
     }

src/formatter/index.ts

@@ -15,6 +15,7 @@ const CHARS_ESCAPE_MAP: Record<string, string> = {
   "'": "\\'",
   "\\": "\\\\"
 };
+const SET_PREFIX = "set ";
 
 export class Tuple {
   value: unknown[];
@@ -238,4 +239,43 @@ export class QueryFormatter {
     }
     return query;
   }
+
+  isSetStatement(query: string): boolean {
+    return query.trim().toLowerCase().startsWith(SET_PREFIX);
+  }
+
+  trimTrailingSemicolon(text: string): string {
+    let end = text.length;
+    while (end > 0 && text[end - 1] == ";") {
+      end--;
+    }
+    return text.substring(0, end);
+  }
+
+  trimStringQuotes(text: string): string {
+    if (text.length < 2) {
+      return text;
+    }
+
+    const firstChar = text[0];
+    const lastChar = text[text.length - 1];
+    if (firstChar === lastChar && (firstChar === "'" || firstChar === '"')) {
+      return text.substring(1, text.length - 1);
+    }
+
+    return text;
+  }
+
+  splitSetStatement(query: string): [string, string] {
+    query = this.trimTrailingSemicolon(query);
+
+    const strippedEquation = query.trim().substring(SET_PREFIX.length).trim();
+    const index = strippedEquation.indexOf("=");
+    const key = strippedEquation.slice(0, index),
+      value = strippedEquation.slice(index + 1);
+    if (key.length == 0 || value.length == 0) {
+      throw new Error("Invalid SET statement format");
+    }
+    return [key.trim(), this.trimStringQuotes(value.trim())];
+  }
 }

test/integration/v2/setStatement.test.ts

@@ -0,0 +1,35 @@
+import { Firebolt } from "../../../src/index";
+
+const connectionParams = {
+  auth: {
+    client_id: process.env.FIREBOLT_CLIENT_ID as string,
+    client_secret: process.env.FIREBOLT_CLIENT_SECRET as string
+  },
+  account: process.env.FIREBOLT_ACCOUNT_V1 as string,
+  database: process.env.FIREBOLT_DATABASE as string,
+  engineName: process.env.FIREBOLT_ENGINE_NAME as string
+};
+
+jest.setTimeout(100000);
+
+describe("set statements execution", () => {
+  it("set statements are executed correctly", async () => {
+    const setQuery = "SET time_zone=America/New_York";
+    const query =
+      "SELECT '2023-01-05 17:04:42.123456 Europe/Berlin'::TIMESTAMPTZ;";
+
+    const firebolt = Firebolt({
+      apiEndpoint: process.env.FIREBOLT_API_ENDPOINT as string
+    });
+
+    const connection = await firebolt.connect(connectionParams);
+
+    await connection.execute(setQuery);
+    const statement = await connection.execute(query);
+
+    const { data } = await statement.fetchResult();
+    const expected = new Date("2023-01-05T11:04:42.123456-05:00");
+
+    expect(data[0][0]).toEqual(expected);
+  });
+});

test/unit/connection.test.ts

@@ -605,4 +605,72 @@ describe("Connection", () => {
     await connection.execute("SELECT 1");
     expect(databaseUsed).toEqual("dummy");
   });
+
+  it("handles set statements correctly", async () => {
+    let searchParamsUsed = new URLSearchParams();
+    server.use(
+      rest.post(`https://some_engine.com`, async (req, res, ctx) => {
+        const body = await req.text();
+        if (body.startsWith("SELECT 1")) {
+          searchParamsUsed = req.url.searchParams;
+          return res(ctx.json(selectOneResponse));
+        }
+      })
+    );
+
+    const connectionParams: ConnectionOptions = {
+      auth: {
+        client_id: "dummy",
+        client_secret: "dummy"
+      },
+      database: "dummy",
+      engineName: "dummy",
+      account: "my_account"
+    };
+    const firebolt = Firebolt({
+      apiEndpoint
+    });
+
+    const connection = await firebolt.connect(connectionParams);
+    await connection.execute("SET param=value");
+    await connection.execute("SELECT 1");
+    expect(searchParamsUsed.get("param")).toEqual("value");
+  });
+
+  it("handles invalid set statements correctly", async () => {
+    let searchParamsUsed = new URLSearchParams();
+    let searchParamsUsed2 = new URLSearchParams();
+    server.use(
+      rest.post(`https://some_engine.com`, async (req, res, ctx) => {
+        const body = await req.text();
+        if (body.startsWith("SELECT 1")) {
+          searchParamsUsed = req.url.searchParams;
+          return res(ctx.status(500));
+        }
+        if (body.startsWith("SELECT 2")) {
+          searchParamsUsed2 = req.url.searchParams;
+          return res(ctx.json(emptyResponse));
+        }
+      })
+    );
+
+    const connectionParams: ConnectionOptions = {
+      auth: {
+        client_id: "dummy",
+        client_secret: "dummy"
+      },
+      database: "dummy",
+      engineName: "dummy",
+      account: "my_account"
+    };
+    const firebolt = Firebolt({
+      apiEndpoint
+    });
+
+    const connection = await firebolt.connect(connectionParams);
+    await expect(connection.execute("SET param=value")).rejects.toThrow();
+    await connection.execute("SELECT 2");
+    expect(searchParamsUsed.get("param")).toEqual("value");
+    expect(searchParamsUsed2.get("param")).toEqual(null);
+  });
 });

test/unit/statement.test.ts

@@ -299,3 +299,31 @@ describe("parse values", () => {
     expect(isNaN(res["nnan"])).toBe(true);
   });
 });
+
+describe("set statements", () => {
+  it.each([
+    ["select 1", false],
+    ["set foo = 1", true],
+    ["  set   foo = 1", true],
+    ["update table set foo = 1", false]
+  ])("detects set statement %s correctly", (query, expected) => {
+    expect(new QueryFormatter().isSetStatement(query)).toBe(expected);
+  });
+
+  it.each([
+    ["set foo = bar", "foo", "bar"],
+    ["  set   foo     =   bar   ", "foo", "bar"],
+    ["\t\r set \t\nfoo  \t\r   =   \t\n bar   \r\n", "foo", "bar"],
+    ["set foo = bar;", "foo", "bar"],
+    ["set a='some 'string'", "a", "some 'string"],
+    [
+      'set query_parameters={"name":"param1","value":"Hello, world!"}',
+      "query_parameters",
+      '{"name":"param1","value":"Hello, world!"}'
+    ],
+    ["set key='val=ue'", "key", "val=ue"],
+    ["set key='val\nue'", "key", "val\nue"]
+  ])("parses set statement %j %j %j correctly", (query, key, value) => {
+    expect(new QueryFormatter().splitSetStatement(query)).toEqual([key, value]);
+  });
+});