From 38c31d48d6ccc27d08ad34d8b24f37be7fa98af4 Mon Sep 17 00:00:00 2001
From: tsshadow <teunschriks@gmail.com>
Date: Sun, 14 Sep 2025 17:54:14 +0200
Subject: [PATCH] Add database connection timeout

---
 frontend/dist/.gitkeep                        |  0
 .../Song/Helpers/DatabaseConnector.py         |  5 ++-
 tests/DatabaseConnectorTimeoutTest.py         | 34 +++++++++++++++++++
 3 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 frontend/dist/.gitkeep
 create mode 100644 tests/DatabaseConnectorTimeoutTest.py

diff --git a/frontend/dist/.gitkeep b/frontend/dist/.gitkeep
new file mode 100644
index 0000000..e69de29
diff --git a/postprocessing/Song/Helpers/DatabaseConnector.py b/postprocessing/Song/Helpers/DatabaseConnector.py
index c629d40..2da1c7c 100644
--- a/postprocessing/Song/Helpers/DatabaseConnector.py
+++ b/postprocessing/Song/Helpers/DatabaseConnector.py
@@ -9,6 +9,8 @@ def __init__(self):
         self.port = int(os.getenv('DB_PORT'))
         self.password = os.getenv('DB_PASS')
         self.db = os.getenv('DB_DB')
+        # Fail fast if the database cannot be reached
+        self.connect_timeout = int(os.getenv('DB_CONNECT_TIMEOUT', '5'))
 
     def connect(self):
         return pymysql.connect(
@@ -16,5 +18,6 @@ def connect(self):
             port=self.port,
             user=self.user,
             password=self.password,
-            db=self.db
+            db=self.db,
+            connect_timeout=self.connect_timeout,
         )
diff --git a/tests/DatabaseConnectorTimeoutTest.py b/tests/DatabaseConnectorTimeoutTest.py
new file mode 100644
index 0000000..13cc37c
--- /dev/null
+++ b/tests/DatabaseConnectorTimeoutTest.py
@@ -0,0 +1,34 @@
+import os
+import unittest
+from unittest.mock import patch
+
+from postprocessing.Song.Helpers.DatabaseConnector import DatabaseConnector
+
+
+class DatabaseConnectorTimeoutTest(unittest.TestCase):
+    def setUp(self):
+        os.environ.setdefault('DB_HOST', 'localhost')
+        os.environ.setdefault('DB_USER', 'user')
+        os.environ.setdefault('DB_PORT', '0')
+        os.environ.setdefault('DB_PASS', '')
+        os.environ.setdefault('DB_DB', 'db')
+
+    def test_connect_uses_default_timeout(self):
+        import postprocessing.Song.Helpers.DatabaseConnector as db_module
+        with patch.object(db_module, 'pymysql') as mock_pymysql:
+            DatabaseConnector().connect()
+            _, kwargs = mock_pymysql.connect.call_args
+            self.assertEqual(kwargs.get('connect_timeout'), 5)
+
+    def test_env_can_override_timeout(self):
+        os.environ['DB_CONNECT_TIMEOUT'] = '1'
+        import postprocessing.Song.Helpers.DatabaseConnector as db_module
+        with patch.object(db_module, 'pymysql') as mock_pymysql:
+            DatabaseConnector().connect()
+            _, kwargs = mock_pymysql.connect.call_args
+            self.assertEqual(kwargs.get('connect_timeout'), 1)
+        del os.environ['DB_CONNECT_TIMEOUT']
+
+
+if __name__ == '__main__':
+    unittest.main()