Add database connection timeout to avoid startup hangs

postprocessing/Song/Helpers/DatabaseConnector.py

@@ -9,12 +9,15 @@ 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(
             host=self.host,
             port=self.port,
             user=self.user,
             password=self.password,
-            db=self.db
+            db=self.db,
+            connect_timeout=self.connect_timeout,
         )

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()