From c5edf37cfe2f76ff1483ae7396e6a6926f7a659f Mon Sep 17 00:00:00 2001
From: Mohit Yadav <105265192+mohityadav766@users.noreply.github.com>
Date: Wed, 13 Sep 2023 18:37:36 +0530
Subject: [PATCH] Close unnecessary connections (#13177)

---
 .../jdbi3/unitofwork/JdbiHandleManager.java   |  2 ++
 .../unitofwork/JdbiTransactionManager.java    |  8 ++----
 .../LinkedRequestScopedJdbiHandleManager.java |  6 +++++
 .../ManagedHandleInvocationHandler.java       | 27 ++++++++++---------
 .../RequestScopedJdbiHandleManager.java       |  5 ++++
 5 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/JdbiHandleManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/JdbiHandleManager.java
index d4db5bd8317..77e2eba5eaa 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/JdbiHandleManager.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/JdbiHandleManager.java
@@ -9,6 +9,8 @@ public interface JdbiHandleManager {
 
   Handle get();
 
+  boolean handleExists();
+
   void clear();
 
   default ThreadFactory createThreadFactory() {
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/JdbiTransactionManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/JdbiTransactionManager.java
index 621adf0350e..8449716bea5 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/JdbiTransactionManager.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/JdbiTransactionManager.java
@@ -94,13 +94,9 @@ public class JdbiTransactionManager {
   }
 
   public void terminateHandle() {
-    if (IN_TRANSACTION_HANDLES.contains(handleManager.get().hashCode())) {
+    if (handleManager.handleExists()) {
+      IN_TRANSACTION_HANDLES.remove(handleManager.get().hashCode());
       handleManager.clear();
     }
-    IN_TRANSACTION_HANDLES.remove(handleManager.get().hashCode());
-  }
-
-  public boolean containsHandle(int hashCode) {
-    return IN_TRANSACTION_HANDLES.contains(hashCode);
   }
 }
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/LinkedRequestScopedJdbiHandleManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/LinkedRequestScopedJdbiHandleManager.java
index b5527351707..f19d52ae86a 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/LinkedRequestScopedJdbiHandleManager.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/LinkedRequestScopedJdbiHandleManager.java
@@ -57,6 +57,12 @@ class LinkedRequestScopedJdbiHandleManager implements JdbiHandleManager {
     return handle;
   }
 
+  @Override
+  public boolean handleExists() {
+    // TODO
+    return false;
+  }
+
   @Override
   public void clear() {
     String parent = getConversationId();
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/ManagedHandleInvocationHandler.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/ManagedHandleInvocationHandler.java
index 7bc80bb0dec..4ad34267746 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/ManagedHandleInvocationHandler.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/ManagedHandleInvocationHandler.java
@@ -38,23 +38,26 @@ public class ManagedHandleInvocationHandler<T> implements InvocationHandler {
   }
 
   private Object handleInvocation(Method method, Object[] args) throws Throwable {
-    Handle handle = JdbiUnitOfWorkProvider.getInstance().getHandle();
-    LOG.debug(
-        "{}.{} [{}] Thread Id [{}] with handle id [{}]",
-        method.getDeclaringClass().getSimpleName(),
-        method.getName(),
-        underlying.getSimpleName(),
-        Thread.currentThread().getId(),
-        handle.hashCode());
-
     if (CollectionDAO.class.isAssignableFrom(underlying) && method.isAnnotationPresent(CreateSqlObject.class)) {
       return getWrappedInstanceForDaoClass(method.getReturnType());
     } else {
-      if (!JdbiTransactionManager.getInstance().containsHandle(handle.hashCode())) {
+      Object dao;
+      if (JdbiUnitOfWorkProvider.getInstance().getHandleManager().handleExists()) {
+        Handle handle = JdbiUnitOfWorkProvider.getInstance().getHandle();
+        LOG.debug(
+            "{}.{} [{}] Thread Id [{}] with handle id [{}]",
+            method.getDeclaringClass().getSimpleName(),
+            method.getName(),
+            underlying.getSimpleName(),
+            Thread.currentThread().getId(),
+            handle.hashCode());
+
+        dao = handle.attach(underlying);
+      } else {
         // This is non-transactional request
-        handle.getConnection().setAutoCommit(true);
+        dao = JdbiUnitOfWorkProvider.getInstance().getHandleManager().getJdbi().onDemand(underlying);
       }
-      Object dao = handle.attach(underlying);
+
       try {
         return method.invoke(dao, args);
       } catch (Exception ex) {
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/RequestScopedJdbiHandleManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/RequestScopedJdbiHandleManager.java
index bf2ce23d24f..4858404918f 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/RequestScopedJdbiHandleManager.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/unitofwork/RequestScopedJdbiHandleManager.java
@@ -30,6 +30,11 @@ class RequestScopedJdbiHandleManager implements JdbiHandleManager {
     return handle;
   }
 
+  @Override
+  public boolean handleExists() {
+    return threadLocal.get() != null;
+  }
+
   @Override
   public void clear() {
     Handle handle = threadLocal.get();