(chore) consolidate query runner raw queries into DatabaseService

metlo-labs · Sep 3, 2022 · 3a691ec · 3a691ec
1 parent 885f5ba
commit 3a691ec
Show file tree

Hide file tree

Showing 6 changed files with 102 additions and 69 deletions.
diff --git a/backend/src/services/database/index.ts b/backend/src/services/database/index.ts
@@ -3,6 +3,8 @@ import { DatabaseError } from "pg-protocol"
 import { AppDataSource } from "data-source"
 import { DatabaseModel } from "models"
 import Error500InternalServer from "errors/error-500-internal-server"
+import { getDataType } from "utils"
+import { DataType } from "@common/enums"
 
 export class DatabaseService {
   static isQueryFailedError = (
@@ -44,6 +46,30 @@ export class DatabaseService {
     return execute(1)
   }
 
+  static async executeRawQueries(rawQueries: string | string[]): Promise<any> {
+    const queryRunner = AppDataSource.createQueryRunner()
+    await queryRunner.connect()
+    const isMultiple = getDataType(rawQueries) === DataType.ARRAY
+    let res = null
+    if (isMultiple) {
+      res = []
+    }
+    try {
+      if (isMultiple) {
+        for (const query of rawQueries) {
+          res.push(await queryRunner.query(query))
+        }
+      } else {
+        res = await queryRunner.query(rawQueries as string)
+      }
+    } catch (err) {
+      console.error(`Encountered error while executing raw sql query: ${err}`)
+    } finally {
+      await queryRunner.release()
+    }
+    return res
+  }
+
   static async executeTransactions(
     saveItems: DatabaseModel[][],
     removeItems: DatabaseModel[][],

diff --git a/backend/src/services/get-endpoints/index.ts b/backend/src/services/get-endpoints/index.ts
@@ -65,6 +65,9 @@ export class GetEndpointsService {
       const res: any[] = await Promise.all(
         endpoints[0].map(async endpoint => {
           const firstDetected = await apiTraceRepository.findOne({
+            select: {
+              createdAt: true,
+            },
             where: {
               apiEndpointUuid: endpoint.uuid,
             },
@@ -73,6 +76,9 @@ export class GetEndpointsService {
             },
           })
           const lastActive = await apiTraceRepository.findOne({
+            select: {
+              createdAt: true,
+            },
             where: {
               apiEndpointUuid: endpoint.uuid,
             },

diff --git a/backend/src/services/summary/alerts.ts b/backend/src/services/summary/alerts.ts
@@ -2,17 +2,15 @@ import { Status, AlertType } from "@common/enums"
 import { AppDataSource } from "data-source"
 import { Alert } from "models"
 import cache from "memory-cache"
+import { DatabaseService } from "services/database"
 
 export const getAlertTypeAgg = async () => {
-  const queryRunner = AppDataSource.createQueryRunner()
-  await queryRunner.connect()
   const alertTypeAggRes: { type: AlertType; count: number }[] =
-    await queryRunner.manager.query(`
-    SELECT type, CAST(COUNT(*) AS INTEGER) as count
-    FROM alert WHERE status = 'Open'
-    GROUP BY 1
-  `)
-  await queryRunner.release()
+    await DatabaseService.executeRawQueries(`
+      SELECT type, CAST(COUNT(*) AS INTEGER) as count
+      FROM alert WHERE status = 'Open'
+      GROUP BY 1
+    `)
   return Object.fromEntries(alertTypeAggRes.map(e => [e.type, e.count]))
 }
 

diff --git a/backend/src/services/summary/endpoints.ts b/backend/src/services/summary/endpoints.ts
@@ -1,35 +1,35 @@
-import _ from "lodash"
+import groupBy from "lodash/groupBy"
 import { In } from "typeorm"
 import { AppDataSource } from "data-source"
 import { ApiEndpoint, ApiTrace } from "models"
 import { EndpointAndUsage } from "@common/types"
 import cache from "memory-cache"
+import { DatabaseService } from "services/database"
 
 export const getTopEndpoints = async () => {
-  const queryRunner = AppDataSource.createQueryRunner()
-  await queryRunner.connect()
   const apiTraceRepository = AppDataSource.getRepository(ApiTrace)
   const apiEndpointRepository = AppDataSource.getRepository(ApiEndpoint)
+
   const endpointStats: {
     endpoint: string
     last1MinCnt: number
     last5MinCnt: number
     last30MinCnt: number
-  }[] = await queryRunner.manager.query(`
-      SELECT
-        traces."apiEndpointUuid" as endpoint,
-        CAST(COUNT(*) AS INTEGER) as "last30MinCnt",
-        CAST(SUM(CASE WHEN traces."createdAt" > (NOW() - INTERVAL '5 minutes') THEN 1 ELSE 0 END) AS INTEGER) as "last5MinCnt",
-        CAST(SUM(CASE WHEN traces."createdAt" > (NOW() - INTERVAL '1 minutes') THEN 1 ELSE 0 END) AS INTEGER) as "last1MinCnt"
-      FROM
-        api_trace traces
-      WHERE
-        traces."apiEndpointUuid" IS NOT NULL
-        AND traces."createdAt" > (NOW() - INTERVAL '30 minutes')
-      GROUP BY 1
-      ORDER BY 4 DESC
-      LIMIT 10
-    `)
+  }[] = await DatabaseService.executeRawQueries(`
+    SELECT
+      traces."apiEndpointUuid" as endpoint,
+      CAST(COUNT(*) AS INTEGER) as "last30MinCnt",
+      CAST(SUM(CASE WHEN traces."createdAt" > (NOW() - INTERVAL '5 minutes') THEN 1 ELSE 0 END) AS INTEGER) as "last5MinCnt",
+      CAST(SUM(CASE WHEN traces."createdAt" > (NOW() - INTERVAL '1 minutes') THEN 1 ELSE 0 END) AS INTEGER) as "last1MinCnt"
+    FROM
+      api_trace traces
+    WHERE
+      traces."apiEndpointUuid" IS NOT NULL
+      AND traces."createdAt" > (NOW() - INTERVAL '30 minutes')
+    GROUP BY 1
+    ORDER BY 4 DESC
+    LIMIT 10
+  `)
 
   const endpoints = await apiEndpointRepository.find({
     where: { uuid: In(endpointStats.map(e => e.endpoint)) },
@@ -53,9 +53,7 @@ export const getTopEndpoints = async () => {
       }),
     ),
   )
-  const traceMap = _.groupBy(traces.flat(), e => e.apiEndpointUuid)
-
-  await queryRunner.release()
+  const traceMap = groupBy(traces.flat(), e => e.apiEndpointUuid)
 
   return endpointStats.map(
     stats =>

diff --git a/backend/src/services/summary/piiData.ts b/backend/src/services/summary/piiData.ts
@@ -1,17 +1,14 @@
 import { DataClass } from "@common/enums"
-import { AppDataSource } from "data-source"
 import cache from "memory-cache"
+import { DatabaseService } from "services/database"
 
 export const getPIIDataTypeCount = async () => {
-  const queryRunner = AppDataSource.createQueryRunner()
-  await queryRunner.connect()
   const piiDataTypeCountRes: { type: DataClass; cnt: number }[] =
-    await queryRunner.manager.query(`
-    SELECT data_class as type, CAST(COUNT(*) AS INTEGER) as cnt
-    FROM (SELECT UNNEST("dataClasses") as data_class FROM data_field) tbl
-    GROUP BY 1
-  `)
-  await queryRunner.release()
+    await DatabaseService.executeRawQueries(`
+      SELECT data_class as type, CAST(COUNT(*) AS INTEGER) as cnt
+      FROM (SELECT UNNEST("dataClasses") as data_class FROM data_field) tbl
+      GROUP BY 1
+    `)
   return Object.fromEntries(piiDataTypeCountRes.map(e => [e.type, e.cnt]))
 }
 

diff --git a/backend/src/services/summary/usageStats.ts b/backend/src/services/summary/usageStats.ts
@@ -1,31 +1,36 @@
-import { AppDataSource } from "data-source"
 import { UsageStats } from "@common/types"
 import cache from "memory-cache"
+import { DatabaseService } from "services/database"
 
 export const getUsageStats = async () => {
-  const queryRunner = AppDataSource.createQueryRunner()
+  const statsQuery = `
+    SELECT
+      DATE_TRUNC('day', traces."createdAt") as day,
+      COUNT(*) as cnt
+    FROM api_trace traces
+    WHERE traces."createdAt" > (NOW() - INTERVAL '15 days')
+    GROUP BY 1
+    ORDER BY 1
+  `
+  const lastNRequestsQuery = `
+    SELECT
+      CAST(SUM(CASE WHEN traces."createdAt" > (NOW() - INTERVAL '1 minutes') THEN 1 ELSE 0 END) AS INTEGER) as "last1MinCnt",
+      CAST(COUNT(*) AS INTEGER) as "last60MinCnt"
+    FROM api_trace traces
+    WHERE traces."createdAt" > (NOW() - INTERVAL '60 minutes')
+  `
+  const queryResponses = await DatabaseService.executeRawQueries([
+    statsQuery,
+    lastNRequestsQuery,
+  ])
   const stats: {
     day: string
     cnt: number
-  }[] = await queryRunner.manager.query(`
-  SELECT
-    DATE_TRUNC('day', traces."createdAt") as day,
-    COUNT(*) as cnt
-  FROM api_trace traces
-  WHERE traces."createdAt" > (NOW() - INTERVAL '15 days')
-  GROUP BY 1
-  ORDER BY 1
-  `)
+  }[] = queryResponses[0]
   const lastNRequests: {
     last1MinCnt: number
     last60MinCnt: number
-  } = await queryRunner.manager.query(`
-  SELECT
-    CAST(SUM(CASE WHEN traces."createdAt" > (NOW() - INTERVAL '1 minutes') THEN 1 ELSE 0 END) AS INTEGER) as "last1MinCnt",
-    CAST(COUNT(*) AS INTEGER) as "last60MinCnt"
-  FROM api_trace traces
-  WHERE traces."createdAt" > (NOW() - INTERVAL '60 minutes')
-  `)
+  } = queryResponses[1]
   return {
     dailyUsage: stats,
     last1MinCnt: lastNRequests[0].last1MinCnt,
@@ -52,30 +57,33 @@ interface CountsResponse {
 }
 
 export const getCounts = async () => {
-  const queryRunner = AppDataSource.createQueryRunner()
-  await queryRunner.connect()
-  const newAlertQueryRes = await queryRunner.manager.query(`
+  const newAlertQuery = `
     SELECT
       CAST(COUNT(*) AS INTEGER) as count,
       CAST(SUM(CASE WHEN "riskScore" = 'high' THEN 1 ELSE 0 END) AS INTEGER) as high_risk_count
     FROM alert WHERE status = 'Open'
-  `)
-  const newAlerts = newAlertQueryRes[0].count
-  const highRiskAlerts = newAlertQueryRes[0].high_risk_count
-  const endpointsTrackedQueryRes = await queryRunner.manager.query(`
+  `
+  const endpointsTrackedQuery = `
     SELECT
       CAST(COUNT(*) AS INTEGER) as endpoint_count,
       CAST(COUNT(DISTINCT(host)) AS INTEGER) as host_count
     FROM api_endpoint
-  `)
+  `
+  const piiDataFieldsQuery = `
+    SELECT CAST(COUNT(*) AS INTEGER) as count
+    FROM data_field WHERE "dataTag" = 'PII'
+  `
+  const [newAlertQueryRes, endpointsTrackedQueryRes, piiDataFieldsQueryRes] =
+    await DatabaseService.executeRawQueries([
+      newAlertQuery,
+      endpointsTrackedQuery,
+      piiDataFieldsQuery,
+    ])
+  const newAlerts = newAlertQueryRes[0].count
+  const highRiskAlerts = newAlertQueryRes[0].high_risk_count
   const endpointsTracked = endpointsTrackedQueryRes[0].endpoint_count
   const hostCount = endpointsTrackedQueryRes[0].host_count
-  const piiDataFieldsQueryRes = await queryRunner.manager.query(`
-      SELECT CAST(COUNT(*) AS INTEGER) as count
-      FROM data_field WHERE "dataTag" = 'PII'
-  `)
   const piiDataFields = piiDataFieldsQueryRes[0].count
-  await queryRunner.release()
   return {
     newAlerts,
     endpointsTracked,