Add query stream, describe tables and columns for databricks client. (#55)

Sylvestre · web-flow · commit e350185f2a87 · 2023-01-05T15:04:17.000-05:00
* add basic query stream support * making error handling not crashing application * scope * Describe tables and columns utils function (#57) * Parameter naive implementation (#59) * fix exclusion character edge case * write schema to response even if no rows is found * update error message to be more explicit * moved runAsync to operation definition
diff --git a/lib/databricks.js b/lib/databricks.js
@@ -1,6 +1,16 @@
+import {json} from "micro";
+import {Readable} from "node:stream";
+import JSONStream from "JSONStream";
+
+import {
+  validateDescribeColumnsPayload,
+  validateQueryPayload,
+} from "./validate.js";
+import {badRequest} from "./errors.js";
+
 export class DatabricksSingleton {
   static instance = null;
-
+  static types = new Map();
   constructor() {
     throw new Error(
       "Do not use new DatabricksSingleton(). Call DatabricksSingleton.initialize() instead."
@@ -10,7 +20,65 @@ export class DatabricksSingleton {
     if (!DatabricksSingleton.instance) {
       try {
         DatabricksSingleton.instance = import("@databricks/sql").then(
-          (module) => module.default
+          (module) => {
+            const databricks = module.default;
+            const {TCLIService_types} = databricks.thrift;
+
+            const boolean = ["null", "boolean"],
+              integer = ["null", "integer"],
+              number = ["null", "number"],
+              object = ["null", "object"],
+              string = ["null", "string"];
+
+            DatabricksSingleton.types
+              .set(TCLIService_types.TTypeId.BIGINT_TYPE, {
+                type: integer,
+              })
+              .set(TCLIService_types.TTypeId.BINARY_TYPE, {
+                type: object,
+                buffer: true,
+              })
+              .set(TCLIService_types.TTypeId.BOOLEAN_TYPE, {
+                type: boolean,
+              })
+              .set(TCLIService_types.TTypeId.TINYINT_TYPE, {
+                type: integer,
+              })
+              .set(TCLIService_types.TTypeId.SMALLINT_TYPE, {
+                type: integer,
+              })
+              .set(TCLIService_types.TTypeId.INT_TYPE, {
+                type: integer,
+              })
+              .set(TCLIService_types.TTypeId.DECIMAL_TYPE, {
+                type: number,
+                decimal: true,
+              })
+              .set(TCLIService_types.TTypeId.DOUBLE_TYPE, {
+                type: number,
+              })
+              .set(TCLIService_types.TTypeId.FLOAT_TYPE, {
+                type: number,
+              })
+              .set(TCLIService_types.TTypeId.TIMESTAMP_TYPE, {
+                type: string,
+                date: true,
+              })
+              .set(TCLIService_types.TTypeId.DATE_TYPE, {
+                type: string,
+                date: true,
+              })
+              .set(TCLIService_types.TTypeId.INTERVAL_DAY_TIME_TYPE, {
+                type: string,
+                date: true,
+              })
+              .set(TCLIService_types.TTypeId.INTERVAL_YEAR_MONTH_TYPE, {
+                type: string,
+                date: true,
+              });
+
+            return databricks;
+          }
         );
       } catch (err) {
         console.error(err);
@@ -24,6 +92,214 @@ export class DatabricksSingleton {
       );
     return DatabricksSingleton.instance;
   }
+
+  get types() {
+    if (DatabricksSingleton.types.size === 0)
+      throw new Error(
+        "DatabricksSingleton not initialized. Call DatabricksSingleton.initialize() first."
+      );
+    return DatabricksSingleton.types;
+  }
+}
+
+export async function describeTables(req, res, connection) {
+  const tablesData = {};
+  const session = await new Promise((resolve, reject) => {
+    connection.on("error", reject).openSession().then(resolve);
+  });
+
+  try {
+    const catalogOperation = await session.getCatalogs({runAsync: true});
+    const catalogs = await catalogOperation.fetchAll();
+
+    for (const catalog of catalogs) {
+      const tableOperation = await session.getTables({
+        catalogName: catalog.TABLE_CAT,
+        runAsync: true,
+      });
+      const tables = await tableOperation.fetchAll();
+      tables.forEach(
+        (table) =>
+          (tablesData[
+            `${table.TABLE_CAT}.${table.TABLE_SCHEM}.${table.TABLE_NAME}`
+          ] = table)
+      );
+    }
+  } catch (error) {
+    if (!error.statusCode) error.statusCode = 400;
+    throw error;
+  } finally {
+    if (session) {
+      try {
+        await session.close();
+      } catch (err) {
+        console.error(err);
+      }
+    }
+
+    if (connection) {
+      try {
+        await connection.close();
+      } catch (err) {
+        console.error(err);
+      }
+    }
+  }
+
+  return tablesData;
+}
+
+export async function describeColumns(req, res, connection) {
+  let columns = [];
+  const session = await new Promise((resolve, reject) => {
+    connection.on("error", reject).openSession().then(resolve);
+  });
+
+  const body = await json(req);
+
+  if (!validateDescribeColumnsPayload(body)) throw badRequest();
+
+  const {table, catalog, schema} = body;
+
+  try {
+    const columnsOperation = await session.getColumns({
+      catalogName: catalog,
+      schemaName: schema,
+      tableName: table,
+      runAsync: true,
+    });
+    columns = await columnsOperation.fetchAll();
+  } catch (error) {
+    if (!error.statusCode) error.statusCode = 400;
+    throw error;
+  } finally {
+    if (session) {
+      try {
+        await session.close();
+      } catch (err) {
+        console.error(err);
+      }
+    }
+
+    if (connection) {
+      try {
+        await connection.close();
+      } catch (err) {
+        console.error(err);
+      }
+    }
+  }
+
+  return columns;
+}
+
+/*
+ * This function is running a given query and streams the results back to the client.
+ * */
+export async function queryStream(req, res, connection) {
+  let query;
+  const session = await new Promise((resolve, reject) => {
+    connection.on("error", reject).openSession().then(resolve);
+  });
+
+  const body = await json(req);
+
+  if (!validateQueryPayload(body)) throw badRequest();
+
+  res.setHeader("Content-Type", "text/plain");
+  const keepAlive = setInterval(() => res.write("\n"), 25e3);
+
+  let {sql, params = []} = body;
+
+  if (params.length) {
+    sql = params.reduce((sql, p, idx) => {
+      if (typeof p === "string") {
+        // If the parameter is a string that includes the ' (exclusion character),
+        // we need to add the r tag to treat it as raw-literal string.
+        p = p.match(/'/g) ? `r'${p}'` : `'${p}'`;
+      }
+
+      return sql.replace(`:${idx + 1}`, p);
+    }, sql);
+  }
+
+  try {
+    query = await session.executeStatement(sql, {runAsync: true});
+    const rows = await query.fetchAll();
+    const schema = await query.getSchema();
+
+    const responseSchema = {
+      type: "array",
+      items: {
+        type: "object",
+        properties: schema.columns.reduce((schema, col, idx) => {
+          return {
+            ...schema,
+            ...{
+              [col.columnName]: dataTypeSchema(
+                col.typeDesc.types[0].primitiveEntry.type
+              ),
+            },
+          };
+        }, {}),
+      },
+    };
+    res.write(`${JSON.stringify(responseSchema)}`);
+    res.write("\n");
+
+    await new Promise(async (resolve, reject) => {
+      const stream = new Readable.from(rows);
+
+      stream.once("data", () => {
+        clearInterval(keepAlive);
+      });
+
+      stream.on("close", (error) => {
+        resolve();
+        stream.destroy();
+      });
+
+      stream.on("error", reject);
+
+      stream.pipe(JSONStream.stringify("", "\n", "\n")).pipe(res);
+    });
+  } catch (error) {
+    if (!error.statusCode) error.statusCode = 400;
+    // The default error message is too generic.
+    // We want to show a more explicit error message, such as "Table not found" ...etc.
+    if (error?.response?.displayMessage)
+      error.message = error.response.displayMessage;
+
+    throw error;
+  } finally {
+    clearInterval(keepAlive);
+
+    if (query) {
+      try {
+        await query.close();
+      } catch (err) {
+        console.error(err);
+      }
+    }
+
+    if (session) {
+      try {
+        await session.close();
+      } catch (err) {
+        console.error(err);
+      }
+    }
+
+    if (connection) {
+      try {
+        await connection.close();
+      } catch (err) {
+        console.error(err);
+      }
+    }
+  }
+
+  res.end();
 }
 
 /*
@@ -58,7 +334,26 @@ export default ({token, host, path}) => {
     const client = new databricks.DBSQLClient();
     const connection = await client.connect({token, host, path});
 
-    // TODO: replace with queryStream method when implemented
-    return check(req, res, connection);
+    if (req.method === "POST") {
+      if (req.url === "/describe-tables") {
+        return describeTables(req, res, connection);
+      }
+
+      if (req.url === "/describe-columns") {
+        return describeColumns(req, res, connection);
+      }
+    }
+
+    return queryStream(req, res, connection);
   };
 };
+
+// See https://github.com/databricks/databricks-sql-nodejs/blob/main/tests/unit/result/JsonResult.test.js
+function dataTypeSchema(type) {
+  const types = DatabricksSingleton.types;
+  if (types.has(type)) {
+    return types.get(type);
+  }
+
+  return {type: ["null", "string"]};
+}
diff --git a/lib/validate.js b/lib/validate.js
@@ -11,3 +11,13 @@ export const validateQueryPayload = ajv.compile({
     params: {type: ["object", "array"]},
   },
 });
+export const validateDescribeColumnsPayload = ajv.compile({
+  type: "object",
+  additionalProperties: false,
+  required: ["table"],
+  properties: {
+    catalog: {type: "string", minLength: 1, maxLength: 32 * 1000},
+    schema: {type: "string", minLength: 1, maxLength: 32 * 1000},
+    table: {type: "string", minLength: 1, maxLength: 32 * 1000}
+  },
+})
