From 977ffa1286ea049050be335060f56436729945f2 Mon Sep 17 00:00:00 2001
From: Kavin <20838718+FireMasterK@users.noreply.github.com>
Date: Sun, 16 Apr 2023 22:56:33 +0100
Subject: [PATCH] Add shutdown thread to close session factory

---
 src/main/java/me/kavin/piped/Main.java        |  3 +++
 .../piped/utils/DatabaseSessionFactory.java   | 21 ++++++++++++-------
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/main/java/me/kavin/piped/Main.java b/src/main/java/me/kavin/piped/Main.java
index 6aedabb..c4f48d4 100644
--- a/src/main/java/me/kavin/piped/Main.java
+++ b/src/main/java/me/kavin/piped/Main.java
@@ -74,6 +74,9 @@ public class Main {
             System.exit(1);
         }
 
+        // Close the HikariCP connection pool
+        Runtime.getRuntime().addShutdownHook(new Thread(DatabaseSessionFactory::close));
+
         if (Constants.DISABLE_TIMERS)
             return;
 
diff --git a/src/main/java/me/kavin/piped/utils/DatabaseSessionFactory.java b/src/main/java/me/kavin/piped/utils/DatabaseSessionFactory.java
index 21ef8f3..edc6089 100644
--- a/src/main/java/me/kavin/piped/utils/DatabaseSessionFactory.java
+++ b/src/main/java/me/kavin/piped/utils/DatabaseSessionFactory.java
@@ -12,15 +12,18 @@ public class DatabaseSessionFactory {
     private static final SessionFactory sessionFactory;
 
     static {
+        try {
+            final Configuration configuration = new Configuration();
 
-        final Configuration configuration = new Configuration();
+            Constants.hibernateProperties.forEach(configuration::setProperty);
+            configuration.configure();
 
-        Constants.hibernateProperties.forEach(configuration::setProperty);
-        configuration.configure();
-
-        sessionFactory = configuration.addAnnotatedClass(User.class).addAnnotatedClass(Channel.class)
-                .addAnnotatedClass(Video.class).addAnnotatedClass(PubSub.class).addAnnotatedClass(Playlist.class)
-                .addAnnotatedClass(PlaylistVideo.class).addAnnotatedClass(UnauthenticatedSubscription.class).buildSessionFactory();
+            sessionFactory = configuration.addAnnotatedClass(User.class).addAnnotatedClass(Channel.class)
+                    .addAnnotatedClass(Video.class).addAnnotatedClass(PubSub.class).addAnnotatedClass(Playlist.class)
+                    .addAnnotatedClass(PlaylistVideo.class).addAnnotatedClass(UnauthenticatedSubscription.class).buildSessionFactory();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
     }
 
     public static Session createSession() {
@@ -30,4 +33,8 @@ public class DatabaseSessionFactory {
     public static StatelessSession createStatelessSession() {
         return sessionFactory.openStatelessSession();
     }
+
+    public static void close() {
+        sessionFactory.close();
+    }
 }