Merge pull request #301 from Sanketika-Obsrv/query-wrapper-superset

Superset Datasource Query Wrapper
Sanketika-Obsrv · Dec 30, 2024 · 9c876ab · 9c876ab
2 parents 022d011 + 0b3b4a5
commit 9c876ab
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 16 deletions.
diff --git a/api-service/src/controllers/QueryWrapper/SqlQueryWrapper.ts b/api-service/src/controllers/QueryWrapper/SqlQueryWrapper.ts
@@ -5,6 +5,8 @@ import logger from "../../logger";
 import { ResponseHandler } from "../../helpers/ResponseHandler";
 import { ErrorObject } from "../../types/ResponseModel";
 import { druidHttpService } from "../../connections/druidConnection";
+import { getDatasourceList } from "../../services/DatasourceService";
+import { AxiosResponse } from "axios";
 
 const apiId = "api.obsrv.data.sql-query";
 const errorCode = "SQL_QUERY_FAILURE"
@@ -25,10 +27,16 @@ export const sqlQuery = async (req: Request, res: Response) => {
             } as ErrorObject, req, res);
         }
 
-        const result = await druidHttpService.post(`${config.query_api.druid.sql_query_path}`, req.body, {
-            headers: { Authorization: authorization },
-        });
-
+        const query = req.body.query as string;
+        let result: AxiosResponse;
+        if (isTableSchemaQuery(query)) {
+            const dataSources = await fetchDruidDataSources();
+            result = createMockAxiosResponse(dataSources);
+        } else {
+            result = await druidHttpService.post(`${config.query_api.druid.sql_query_path}`, req.body, {
+                headers: { Authorization: authorization },
+            });
+        }
         logger.info({ messsge: "Successfully fetched data using sql query", apiId, resmsgid })
         ResponseHandler.flatResponse(req, res, result)
     } catch (error: any) {
@@ -37,4 +45,37 @@ export const sqlQuery = async (req: Request, res: Response) => {
         logger.error(error, apiId, code, resmsgid)
         ResponseHandler.errorResponse(errorMessage, req, res);
     }
-}
+}
+
+const fetchDruidDataSources = async (): Promise<{ TABLE_NAME: string }[]> => {
+    try {
+        const dataSources = await getDatasourceList();
+        return dataSources
+            .filter((ds: any) => ds.type === "druid")
+            .map((ds: any) => ({ TABLE_NAME: ds.dataValues.datasource_ref }));
+    } catch (error) {
+        logger.error({ message: "Failed to fetch Druid data sources", error });
+        throw new Error("Failed to fetch Druid data sources");
+    }
+};
+
+const isTableSchemaQuery = (sqlQuery?: string): boolean => {
+    return (
+      sqlQuery
+        ?.trim()
+        .replace(/\s+/g, " ")
+        .toUpperCase() ===
+      "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = 'DRUID'"
+    );
+  };
+
+
+const createMockAxiosResponse = (data: any): AxiosResponse => {
+    return {
+        data,
+        status: 200,
+        statusText: "OK",
+        headers: {},
+        config: {},
+    } as AxiosResponse;
+};
diff --git a/api-service/src/services/DatasourceService.ts b/api-service/src/services/DatasourceService.ts
@@ -1,16 +1,13 @@
 import { Datasource } from "../models/Datasource";
 
-export const getDatasourceList = async (datasetId: string, raw = false) => {
-    const dataSource = await Datasource.findAll({
-        where: {
-            dataset_id: datasetId,
-        },
-        raw: raw
-    });
-    return dataSource
-}
-
-
+export const getDatasourceList = async (datasetId?: string, raw: boolean = false) => {
+    const query: any = { raw };
+    if (datasetId) {
+        query.where = { dataset_id: datasetId };
+    }
+
+    return Datasource.findAll(query);
+};