From 2ef1586d4d520d8edc01527686045b7ce5249ce8 Mon Sep 17 00:00:00 2001
From: mike12345567 <me@michaeldrury.co.uk>
Date: Mon, 19 Apr 2021 22:36:14 +0100
Subject: [PATCH] Adding cleanup operation to dynamo client.

---
 packages/server/src/api/controllers/query.js | 18 ++++++++++++------
 packages/server/src/db/dynamoClient.js       |  1 +
 packages/server/src/integrations/dynamodb.js | 13 +++++++++++++
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/packages/server/src/api/controllers/query.js b/packages/server/src/api/controllers/query.js
index a2badb0d0d..45dc4e0d86 100644
--- a/packages/server/src/api/controllers/query.js
+++ b/packages/server/src/api/controllers/query.js
@@ -116,9 +116,8 @@ exports.preview = async function(ctx) {
 
   const enrichedQuery = await enrichQueryFields(fields, parameters)
 
-  const rows = formatResponse(
-    await new Integration(datasource.config)[queryVerb](enrichedQuery)
-  )
+  const integration = new Integration(datasource.config)
+  const rows = formatResponse(await integration[queryVerb](enrichedQuery))
 
   // get all the potential fields in the schema
   const keys = rows.flatMap(Object.keys)
@@ -127,6 +126,10 @@ exports.preview = async function(ctx) {
     rows,
     schemaFields: [...new Set(keys)],
   }
+  // cleanup
+  if (integration.end) {
+    integration.end()
+  }
 }
 
 exports.execute = async function(ctx) {
@@ -146,10 +149,13 @@ exports.execute = async function(ctx) {
     ctx.request.body.parameters
   )
 
+  const integration = new Integration(datasource.config)
   // call the relevant CRUD method on the integration class
-  ctx.body = formatResponse(
-    await new Integration(datasource.config)[query.queryVerb](enrichedQuery)
-  )
+  ctx.body = formatResponse(await integration[query.queryVerb](enrichedQuery))
+  // cleanup
+  if (integration.end) {
+    integration.end()
+  }
 }
 
 exports.destroy = async function(ctx) {
diff --git a/packages/server/src/db/dynamoClient.js b/packages/server/src/db/dynamoClient.js
index fcba726f84..3e9c8ced14 100644
--- a/packages/server/src/db/dynamoClient.js
+++ b/packages/server/src/db/dynamoClient.js
@@ -2,6 +2,7 @@ let _ = require("lodash")
 let env = require("../environment")
 
 const AWS_REGION = env.AWS_REGION ? env.AWS_REGION : "eu-west-1"
+exports.AWS_REGION = AWS_REGION
 
 const TableInfo = {
   API_KEYS: {
diff --git a/packages/server/src/integrations/dynamodb.js b/packages/server/src/integrations/dynamodb.js
index da95411b2b..aa77d6d4f0 100644
--- a/packages/server/src/integrations/dynamodb.js
+++ b/packages/server/src/integrations/dynamodb.js
@@ -1,5 +1,6 @@
 const AWS = require("aws-sdk")
 const { FIELD_TYPES, QUERY_TYPES } = require("./Integration")
+const { AWS_REGION } = require("../db/dynamoClient")
 
 const SCHEMA = {
   docs: "https://github.com/dabit3/dynamodb-documentclient-cheat-sheet",
@@ -114,10 +115,22 @@ class DynamoDBIntegration {
     })
   }
 
+  end() {
+    this.disconnect()
+  }
+
   connect() {
     AWS.config.update(this.config)
   }
 
+  disconnect() {
+    AWS.config.update({
+      secretAccessKey: undefined,
+      accessKeyId: undefined,
+      region: AWS_REGION,
+    })
+  }
+
   async create(query) {
     const params = {
       TableName: query.table,