Cleanup Data Access (#73)

Author: akshay288

backend/package.json

@@ -35,7 +35,6 @@
     "@leoscope/openapi-response-validator": "^1.0.2",
     "@metlo/testing": "^0.0.3",
     "@types/async-retry": "^1.4.4",
-    "@types/newman": "^5.3.0",
     "@types/ssh2": "^1.11.5",
     "async-retry": "^1.3.3",
     "aws-sdk": "^2.1189.0",
@@ -53,7 +52,6 @@
     "json-source-map": "^0.6.1",
     "lodash": "^4.17.21",
     "luxon": "^3.0.3",
-    "memory-cache": "^0.2.0",
     "multer": "^1.4.5-lts.1",
     "node-schedule": "^2.1.0",
     "node-ssh": "^13.0.0",
@@ -79,7 +77,6 @@
     "@types/js-yaml": "^4.0.5",
     "@types/lodash": "^4.14.184",
     "@types/luxon": "^3.0.1",
-    "@types/memory-cache": "^0.2.2",
     "@types/multer": "^1.4.7",
     "@types/node": "^18.6.1",
     "@types/node-schedule": "^2.1.0",

backend/src/analyze-traces.ts

@@ -1,10 +1,9 @@
 import { v4 as uuidv4 } from "uuid"
 import { AppDataSource } from "data-source"
-import { ApiTrace, ApiEndpoint, DataField, Alert } from "models"
+import { ApiTrace, ApiEndpoint, DataField, Alert, OpenApiSpec } from "models"
 import { DataFieldService } from "services/data-field"
 import { SpecService } from "services/spec"
 import { AlertService } from "services/alert"
-import { DatabaseService } from "services/database"
 import { RedisClient } from "utils/redis"
 import { TRACES_QUEUE } from "~/constants"
 import { QueryRunner } from "typeorm"
@@ -18,14 +17,20 @@ import {
 import { getPathTokens } from "@common/utils"
 import { AlertType } from "@common/enums"
 import { isGraphQlEndpoint } from "services/graphql"
+import { isQueryFailedError, retryTypeormTransaction } from "utils/db"
+import { MetloContext } from "types"
+import { DatabaseService } from "services/database"
+import { getEntityManager, getQB } from "services/database/utils"
 
-const GET_ENDPOINT_QUERY = `
+const getEndpointQuery = (ctx: MetloContext) => `
 SELECT
   endpoint. *,
   CASE WHEN spec."isAutoGenerated" IS NULL THEN NULL ELSE json_build_object('isAutoGenerated', spec."isAutoGenerated") END as "openapiSpec"
 FROM
-  "api_endpoint" endpoint
-  LEFT JOIN "open_api_spec" spec ON endpoint."openapiSpecName" = spec.name
+  ${ApiEndpoint.getTableName(ctx)} endpoint
+  LEFT JOIN ${OpenApiSpec.getTableName(
+    ctx,
+  )} spec ON endpoint."openapiSpecName" = spec.name
 WHERE
   $1 ~ "pathRegex"
   AND method = $2
@@ -39,7 +44,7 @@ LIMIT
   1
 `
 
-const GET_DATA_FIELDS_QUERY = `
+const getDataFieldsQuery = (ctx: MetloContext) => `
 SELECT
   uuid,
   "dataClasses"::text[],
@@ -53,21 +58,27 @@ SELECT
   "dataPath",
   "apiEndpointUuid"
 FROM
-  data_field
+  ${DataField.getTableName(ctx)} data_field
 WHERE
   "apiEndpointUuid" = $1
 `
 
-const getQueuedApiTrace = async (): Promise<QueuedApiTrace> => {
+const getQueuedApiTrace = async (
+  ctx: MetloContext,
+): Promise<QueuedApiTrace> => {
   try {
-    const traceString = await RedisClient.popValueFromRedisList(TRACES_QUEUE)
+    const traceString = await RedisClient.popValueFromRedisList(
+      ctx,
+      TRACES_QUEUE,
+    )
     return JSON.parse(traceString)
   } catch (err) {
     return null
   }
 }
 
 const analyze = async (
+  ctx: MetloContext,
   trace: QueuedApiTrace,
   apiEndpoint: ApiEndpoint,
   queryRunner: QueryRunner,
@@ -76,11 +87,13 @@ const analyze = async (
   endpointUpdateDates(trace.createdAt, apiEndpoint)
   const dataFields = DataFieldService.findAllDataFields(trace, apiEndpoint)
   let alerts = await SpecService.findOpenApiSpecDiff(
+    ctx,
     trace,
     apiEndpoint,
     queryRunner,
   )
   const sensitiveDataAlerts = await AlertService.createDataFieldAlerts(
+    ctx,
     dataFields,
     apiEndpoint.uuid,
     apiEndpoint.path,
@@ -90,6 +103,7 @@ const analyze = async (
   alerts = alerts?.concat(sensitiveDataAlerts)
   if (newEndpoint) {
     const newEndpointAlert = await AlertService.createAlert(
+      ctx,
       AlertType.NEW_ENDPOINT,
       apiEndpoint,
     )
@@ -99,18 +113,17 @@ const analyze = async (
   }
 
   await queryRunner.startTransaction()
-  await DatabaseService.retryTypeormTransaction(
+  await retryTypeormTransaction(
     () =>
-      queryRunner.manager.insert(ApiTrace, {
+      getEntityManager(ctx, queryRunner).insert(ApiTrace, {
         ...trace,
         apiEndpointUuid: apiEndpoint.uuid,
       }),
     5,
   )
-  await DatabaseService.retryTypeormTransaction(
+  await retryTypeormTransaction(
     () =>
-      queryRunner.manager
-        .createQueryBuilder()
+      getQB(ctx, queryRunner)
         .insert()
         .into(DataField)
         .values(dataFields)
@@ -127,21 +140,19 @@ const analyze = async (
         .execute(),
     5,
   )
-  await DatabaseService.retryTypeormTransaction(
+  await retryTypeormTransaction(
     () =>
-      queryRunner.manager
-        .createQueryBuilder()
+      getQB(ctx, queryRunner)
         .insert()
         .into(Alert)
         .values(alerts)
         .orIgnore()
         .execute(),
     5,
   )
-  await DatabaseService.retryTypeormTransaction(
+  await retryTypeormTransaction(
     () =>
-      queryRunner.manager
-        .createQueryBuilder()
+      getQB(ctx, queryRunner)
         .update(ApiEndpoint)
         .set({
           firstDetected: apiEndpoint.firstDetected,
@@ -156,6 +167,7 @@ const analyze = async (
 }
 
 const generateEndpoint = async (
+  ctx: MetloContext,
   trace: QueuedApiTrace,
   queryRunner: QueryRunner,
 ): Promise<void> => {
@@ -201,36 +213,35 @@ const generateEndpoint = async (
 
     try {
       await queryRunner.startTransaction()
-      await DatabaseService.retryTypeormTransaction(
+      await retryTypeormTransaction(
         () =>
-          queryRunner.manager
-            .createQueryBuilder()
+          getQB(ctx, queryRunner)
             .insert()
             .into(ApiEndpoint)
             .values(apiEndpoint)
             .execute(),
         5,
       )
       await queryRunner.commitTransaction()
-      await analyze(trace, apiEndpoint, queryRunner, true)
+      await analyze(ctx, trace, apiEndpoint, queryRunner, true)
     } catch (err) {
       if (queryRunner.isTransactionActive) {
         await queryRunner.rollbackTransaction()
       }
-      if (DatabaseService.isQueryFailedError(err) && err.code === "23505") {
-        const existingEndpoint = await queryRunner.manager.findOne(
-          ApiEndpoint,
-          {
-            where: {
-              path: trace.path,
-              host: trace.host,
-              method: trace.method,
-            },
-            relations: { dataFields: true },
+      if (isQueryFailedError(err) && err.code === "23505") {
+        const existingEndpoint = await getEntityManager(
+          ctx,
+          queryRunner,
+        ).findOne(ApiEndpoint, {
+          where: {
+            path: trace.path,
+            host: trace.host,
+            method: trace.method,
           },
-        )
+          relations: { dataFields: true },
+        })
         if (existingEndpoint) {
-          await analyze(trace, existingEndpoint, queryRunner)
+          await analyze(ctx, trace, existingEndpoint, queryRunner)
         }
       } else {
         console.error(`Error generating new endpoint: ${err}`)
@@ -240,6 +251,8 @@ const generateEndpoint = async (
 }
 
 const analyzeTraces = async (): Promise<void> => {
+  const ctx: MetloContext = {}
+
   const datasource = await AppDataSource.initialize()
   if (!datasource.isInitialized) {
     console.error("Couldn't initialize datasource...")
@@ -251,26 +264,26 @@ const analyzeTraces = async (): Promise<void> => {
   await queryRunner.connect()
   while (true) {
     try {
-      const trace = await getQueuedApiTrace()
+      const trace = await getQueuedApiTrace(ctx)
       if (trace) {
         trace.createdAt = new Date(trace.createdAt)
         const apiEndpoint: ApiEndpoint = (
-          await queryRunner.query(GET_ENDPOINT_QUERY, [
+          await queryRunner.query(getEndpointQuery(ctx), [
             trace.path,
             trace.method,
             trace.host,
           ])
         )?.[0]
         if (apiEndpoint && !skipAutoGeneratedMatch(apiEndpoint, trace.path)) {
-          const dataFields: DataField[] = await queryRunner.query(
-            GET_DATA_FIELDS_QUERY,
+          const dataFields: DataField[] = await DatabaseService.executeRawQuery(
+            getDataFieldsQuery(ctx),
             [apiEndpoint.uuid],
           )
           apiEndpoint.dataFields = dataFields
-          await analyze(trace, apiEndpoint, queryRunner)
+          await analyze(ctx, trace, apiEndpoint, queryRunner)
         } else {
           if (trace.responseStatus !== 404 && trace.responseStatus !== 405) {
-            await generateEndpoint(trace, queryRunner)
+            await generateEndpoint(ctx, trace, queryRunner)
           }
         }
       }

backend/src/api/alert/index.ts

@@ -2,28 +2,30 @@ import { Request, Response } from "express"
 import { AlertService } from "services/alert"
 import { GetAlertParams, UpdateAlertParams } from "@common/types"
 import ApiResponseHandler from "api-response-handler"
+import { MetloRequest } from "types"
 
 export const getAlertsHandler = async (
-  req: Request,
+  req: MetloRequest,
   res: Response,
 ): Promise<void> => {
   try {
     const alertParams: GetAlertParams = req.query
-    const alerts = await AlertService.getAlerts(alertParams)
+    const alerts = await AlertService.getAlerts(req.ctx, alertParams)
     await ApiResponseHandler.success(res, alerts)
   } catch (err) {
     await ApiResponseHandler.error(res, err)
   }
 }
 
 export const updateAlertHandler = async (
-  req: Request,
+  req: MetloRequest,
   res: Response,
 ): Promise<void> => {
   try {
     const { alertId } = req.params
     const updateAlertParams: UpdateAlertParams = req.body
     const updatedAlert = await AlertService.updateAlert(
+      req.ctx,
       alertId,
       updateAlertParams,
     )

backend/src/api/alert/vulnerability.ts

@@ -1,15 +1,16 @@
-import { Request, Response } from "express"
+import { Response } from "express"
 import { GetVulnerabilityAggParams } from "@common/types"
 import ApiResponseHandler from "api-response-handler"
 import { getVulnerabilityAgg } from "services/summary/vulnerabilities"
+import { MetloRequest } from "types"
 
 export const getVulnerabilitySummaryHandler = async (
-  req: Request,
+  req: MetloRequest,
   res: Response,
 ): Promise<void> => {
   try {
     const params: GetVulnerabilityAggParams = req.query
-    const out = await getVulnerabilityAgg(params)
+    const out = await getVulnerabilityAgg(req.ctx, params)
     await ApiResponseHandler.success(res, out)
   } catch (err) {
     await ApiResponseHandler.error(res, err)