fix: move all DataServices to use GramConnectionPool and transaction …

…instead of the pg.Pool

api/src/healthchecks/postgresCheck.ts

@@ -33,7 +33,7 @@ export async function postgresAvailableConnectionsCheck(dal: DataAccessLayer) {
     const check: any = {
       name: "@gram/api-postgres-available-connections",
       actionable: true,
-      healthy: dal.pool.waitingCount === 0,
+      healthy: dal.pool._pool.waitingCount === 0,
       dependentOn: "postgres",
       type: physical.type.EXTERNAL_DEPENDENCY,
       severity: physical.severity.WARNING,

api/src/resources/gram/v1/models/create.spec.ts

@@ -1,17 +1,9 @@
-import { jest } from "@jest/globals";
-import pg from "pg";
-import request from "supertest";
-import * as jwt from "@gram/core/dist/auth/jwt.js";
 import { DataAccessLayer } from "@gram/core/dist/data/dal.js";
 import Model, { Component } from "@gram/core/dist/data/models/Model.js";
 import { _deleteAllTheThings } from "@gram/core/dist/data/utils.js";
+import request from "supertest";
 import { createTestApp } from "../../../../test-util/app.js";
 import { sampleOwnedSystem } from "../../../../test-util/sampleOwnedSystem.js";
-import {
-  sampleAdmin,
-  sampleOtherUser,
-  sampleUser,
-} from "../../../../test-util/sampleUser.js";
 import {
   sampleAdminToken,
   sampleOtherUserToken,
@@ -23,19 +15,18 @@ describe("models.create", () => {
   const componentId2 = "fe93572e-9d0c-4afe-b042-e02c1c459999";
   const dataFlowId = "fe93572e-9d0c-4afe-b042-e02c1cstonks";
   let app: any;
-  let pool: pg.Pool;
   let dal: DataAccessLayer;
   let token = "";
   let adminToken = "";
 
   beforeAll(async () => {
     adminToken = await sampleAdminToken();
     token = await sampleUserToken();
-    ({ app, dal, pool } = await createTestApp());
+    ({ app, dal } = await createTestApp());
   });
 
   beforeEach(async () => {
-    await _deleteAllTheThings(pool);
+    await _deleteAllTheThings(dal.pool);
   });
 
   it("should return 401 on un-authenticated request", async () => {

api/src/test-util/app.ts

@@ -4,5 +4,5 @@ import { bootstrap } from "@gram/core/dist/bootstrap.js";
 export async function createTestApp() {
   const dal = await bootstrap();
   const app = await createApp(dal);
-  return { pool: dal.pool, app, dal };
+  return { pool: dal.pool._pool, app, dal };
 }

api/src/ws/index.ts

@@ -101,6 +101,7 @@ export function attachWebsocketServer(server: Server, dal: DataAccessLayer) {
     if (!server) return;
     server.tellClientsToRefetch("threats", { modelId, componentId });
   });
+
   dal.threatService.on("deleted-for", ({ modelId, componentId }) => {
     const server = wssRegistry.get(modelId);
     log.debug(`threat was deleted via api ${modelId} ${componentId}`);

core/package.json

@@ -41,7 +41,7 @@
     "handlebars": "^4.7.7",
     "jsonwebtoken": "^9.0.0",
     "log4js": "^6.6.1",
-    "pg": "^8.11.1",
+    "pg": "^8.11.3",
     "postgres-migrations": "^5.3.0",
     "prom-client": "^14.0.1",
     "tslib": "^2.3.1"

core/src/data/banners/BannerDataService.ts

@@ -1,6 +1,6 @@
-import pg from "pg";
 import log4js from "log4js";
 import { DataAccessLayer } from "../dal.js";
+import { GramConnectionPool } from "../postgres.js";
 
 interface Banner {
   id: number;
@@ -10,7 +10,11 @@ interface Banner {
 }
 
 export class BannerDataService {
-  constructor(private pool: pg.Pool, private dal: DataAccessLayer) {}
+  constructor(private dal: DataAccessLayer) {
+    this.pool = dal.pool;
+  }
+
+  private pool: GramConnectionPool;
 
   log = log4js.getLogger("BannerDataService");
 

core/src/data/controls/ControlDataService.spec.ts

@@ -12,25 +12,23 @@ import { ControlDataService } from "./ControlDataService.js";
 describe("ControlDataService implementation", () => {
   let data: ControlDataService;
   let dal: DataAccessLayer;
-  let pool: pg.Pool;
   let model: Model;
 
   beforeAll(async () => {
-    pool = await createPostgresPool();
+    const pool = await createPostgresPool();
     dal = new DataAccessLayer(pool);
-    data = new ControlDataService(pool, dal);
-    await _deleteAllTheThings(pool);
+    data = new ControlDataService(dal);
   });
 
   beforeEach(async () => {
-    await _deleteAllTheThings(pool);
+    await _deleteAllTheThings(dal);
     model = new Model("some-system-id", "some-version", "root");
     model.data = { components: [], dataFlows: [] };
     model.id = await dal.modelService.create(model);
   });
 
   afterAll(async () => {
-    await pool.end();
+    await dal.pool.end();
   });
 
   describe("getById control", () => {

core/src/data/controls/ControlDataService.ts

@@ -5,6 +5,7 @@ import { SuggestionID } from "../../suggestions/models.js";
 import { DataAccessLayer } from "../dal.js";
 import { SuggestionStatus } from "../suggestions/Suggestion.js";
 import Control from "./Control.js";
+import { GramConnectionPool } from "../postgres.js";
 
 export function convertToControl(row: any) {
   const control = new Control(
@@ -23,15 +24,13 @@ export function convertToControl(row: any) {
 }
 
 export class ControlDataService extends EventEmitter {
-  constructor(pool: pg.Pool, dal: DataAccessLayer) {
+  constructor(private dal: DataAccessLayer) {
     super();
-    this.pool = pool;
-    this.dal = dal;
+    this.pool = dal.pool;
     this.log = log4js.getLogger("ControlDataService");
   }
 
-  private pool: pg.Pool;
-  private dal: DataAccessLayer;
+  private pool: GramConnectionPool;
   log: any;
   /**
    * Create a control object of specified id
@@ -173,12 +172,9 @@ export class ControlDataService extends EventEmitter {
       WHERE c.model_id = $1::uuid AND m.control_id = c.id AND c.id ${filter}      
    `;
 
-    const client = await this.pool.connect();
-    let result = false;
-    try {
-      await client.query("BEGIN");
+    const result = await this.pool.runTransaction(async (client) => {
       const res = await client.query(query, [modelId, ...ids]);
-      result = res.rowCount > 0;
+      const result = res.rowCount > 0;
 
       if (result) {
         const suggestionIds = res.rows
@@ -200,13 +196,9 @@ export class ControlDataService extends EventEmitter {
           componentId: res.rows[0].component_id,
         });
       }
-      await client.query("COMMIT");
-    } catch (e) {
-      this.log.error("Failed to delete control", e);
-      await client.query("ROLLBACK");
-    } finally {
-      client.release();
-    }
+
+      return result;
+    });
     return result;
   }
 

core/src/data/dal.ts

@@ -19,14 +19,18 @@ import { authzProvider } from "../auth/authorization.js";
 import { systemProvider } from "./systems/systems.js";
 import { SystemProvider } from "./systems/SystemProvider.js";
 import { ReviewerHandler } from "./reviews/ReviewerHandler.js";
-import { createPostgresPool, getDatabaseName } from "./postgres.js";
+import {
+  GramConnectionPool,
+  createPostgresPool,
+  getDatabaseName,
+} from "./postgres.js";
 
 /**
  * Class that carries access to all DataServices, useful for passing dependencies.
  */
 export class DataAccessLayer {
   // Database Connection Pool for direct access to postgres
-  pool: pg.Pool;
+  pool: GramConnectionPool;
 
   // DataServices - specific logic to handle database interactions
   modelService: ModelDataService;
@@ -62,7 +66,7 @@ export class DataAccessLayer {
   }
 
   constructor(pool: pg.Pool) {
-    this.pool = pool;
+    this.pool = new GramConnectionPool(pool);
     this.sysPropHandler = new SystemPropertyHandler();
     this.ccHandler = new ComponentClassHandler();
     this.templateHandler = new TemplateHandler();
@@ -71,15 +75,15 @@ export class DataAccessLayer {
     this.reviewerHandler = new ReviewerHandler();
 
     // Initialize Data Services
-    this.modelService = new ModelDataService(pool, this);
-    this.controlService = new ControlDataService(pool, this);
-    this.threatService = new ThreatDataService(pool, this);
-    this.mitigationService = new MitigationDataService(pool);
-    this.notificationService = new NotificationDataService(pool, this);
-    this.reviewService = new ReviewDataService(pool, this);
-    this.suggestionService = new SuggestionDataService(pool, this);
+    this.modelService = new ModelDataService(this);
+    this.controlService = new ControlDataService(this);
+    this.threatService = new ThreatDataService(this);
+    this.mitigationService = new MitigationDataService(this);
+    this.notificationService = new NotificationDataService(this);
+    this.reviewService = new ReviewDataService(this);
+    this.suggestionService = new SuggestionDataService(this);
     this.suggestionEngine = new SuggestionEngine(this);
-    this.reportService = new ReportDataService(pool, this);
-    this.bannerService = new BannerDataService(pool, this);
+    this.reportService = new ReportDataService(this);
+    this.bannerService = new BannerDataService(this);
   }
 }

core/src/data/mitigations/MitigationDataService.spec.ts

@@ -1,5 +1,4 @@
 import { randomUUID } from "crypto";
-import pg from "pg";
 import Control from "../controls/Control.js";
 import { DataAccessLayer } from "../dal.js";
 import Model from "../models/Model.js";
@@ -12,25 +11,24 @@ import { MitigationDataService } from "./MitigationDataService.js";
 describe("MitigationDataService implementation", () => {
   let data: MitigationDataService;
   let dal: DataAccessLayer;
-  let pool: pg.Pool;
+
   let model: Model;
 
   beforeAll(async () => {
-    pool = await createPostgresPool();
-    data = new MitigationDataService(pool);
+    const pool = await createPostgresPool();
     dal = new DataAccessLayer(pool);
-    await _deleteAllTheThings(pool);
+    data = new MitigationDataService(dal);
   });
 
   beforeEach(async () => {
-    await _deleteAllTheThings(pool);
+    await _deleteAllTheThings(dal);
     model = new Model("some-system-id", "some-version", "root");
     model.data = { components: [], dataFlows: [] };
     model.id = await dal.modelService.create(model);
   });
 
   afterAll(async () => {
-    await pool.end();
+    await dal.pool.end();
   });
 
   describe("getById mitigation", () => {

core/src/data/mitigations/MitigationDataService.ts

@@ -1,6 +1,7 @@
-import pg from "pg";
-import { EventEmitter } from "node:events";
 import log4js from "log4js";
+import { EventEmitter } from "node:events";
+import { DataAccessLayer } from "../dal.js";
+import { GramConnectionPool } from "../postgres.js";
 import Mitigation from "./Mitigation.js";
 
 function convertToMitigation(row: any) {
@@ -15,12 +16,12 @@ function convertToMitigation(row: any) {
 }
 
 export class MitigationDataService extends EventEmitter {
-  constructor(pool: pg.Pool) {
+  constructor(private dal: DataAccessLayer) {
     super();
-    this.pool = pool;
+    this.pool = dal.pool;
     this.log = log4js.getLogger("MitigationDataService");
   }
-  private pool: pg.Pool;
+  private pool: GramConnectionPool;
   log: any;
 
   /**
@@ -45,26 +46,17 @@ export class MitigationDataService extends EventEmitter {
       WHERE id = $1::uuid
     `;
 
-    const client = await this.pool.connect();
-    try {
-      await client.query("BEGIN");
+    await this.pool.runTransaction(async (client) => {
       await client.query(query, [threatId, controlId, createdBy]);
       const res_threats = await client.query(queryThreats, [threatId]);
-      await client.query("COMMIT");
 
       this.emit("updated-for", {
         modelId: res_threats.rows[0].model_id,
         componentId: res_threats.rows[0].component_id,
       });
-      return { threatId, controlId };
-    } catch (e) {
-      await client.query("ROLLBACK");
-      this.log.error("Failed to create mitigation", e);
-    } finally {
-      client.release();
-    }
+    });
 
-    return false;
+    return { threatId, controlId };
   }
 
   /**

core/src/data/models/ModelDataService.spec.ts

@@ -17,21 +17,19 @@ import { sampleUser } from "../../test-util/sampleUser.js";
 describe("ModelDataService implementation", () => {
   let data: ModelDataService;
   let dal: DataAccessLayer;
-  let pool: pg.Pool;
 
   beforeAll(async () => {
-    pool = await createPostgresPool();
+    const pool = await createPostgresPool();
     dal = new DataAccessLayer(pool);
-    data = new ModelDataService(pool, dal);
-    await _deleteAllTheThings(pool);
+    data = new ModelDataService(dal);
   });
 
   afterEach(async () => {
-    await _deleteAllTheThings(pool);
+    await _deleteAllTheThings(dal);
   });
 
   afterAll(async () => {
-    await pool.end();
+    await dal.pool.end();
   });
 
   describe("getById model", () => {
@@ -461,7 +459,6 @@ describe("ModelDataService implementation", () => {
       const mitigationsCopy = await dal.mitigationService.list(
         resModelCopy!.id!
       );
-      console.log(mitigationsCopy);
 
       expect(mitigationsCopy.length).toEqual(3);