Merge pull request #173 from bluelabsio/keepalives

Add keepalives by default to redshift db connection url
bluelabsio · Dec 13, 2022 · f2483d2 · f2483d2
2 parents 1cafea2 + 444ba80
commit f2483d2
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/records_mover/db/connect.py b/records_mover/db/connect.py
@@ -33,6 +33,8 @@
         # Please see SECURITY.md for security implications!
         "local_infile": True
     },
+    # keepalives prevent timeout errors
+    'redshift': {'keepalives': '1', 'keepalives_idle': '30'},
 }
 
 

diff --git a/tests/unit/db/test_connect.py b/tests/unit/db/test_connect.py
@@ -27,7 +27,8 @@ def test_engine_from_lpass_entry(self,
                                     password='hunter1',
                                     port=123,
                                     username='myuser',
-                                    query=None)
+                                    query={'keepalives': '1',
+                                           'keepalives_idle': '30'})
         mock_create_engine.\
             assert_called_with(mock_url.return_value)
         assert engine == mock_create_engine.return_value

diff --git a/tests/unit/db/test_sqlalchemy_driver_picking.py b/tests/unit/db/test_sqlalchemy_driver_picking.py
@@ -11,10 +11,10 @@ def test_create_sqlalchemy_url(self,
                                    mock_db_facts_from_lpass):
         expected_mappings = {
             'psql (redshift)':
-            'redshift://myuser:hunter1@myhost:123/analyticsdb',
+            'redshift://myuser:hunter1@myhost:123/analyticsdb?keepalives=1&keepalives_idle=30',
 
             'redshift':
-            'redshift://myuser:hunter1@myhost:123/analyticsdb',
+            'redshift://myuser:hunter1@myhost:123/analyticsdb?keepalives=1&keepalives_idle=30',
 
             'vertica':
             'vertica+vertica_python://myuser:hunter1@myhost:123/analyticsdb',
@@ -35,8 +35,12 @@ def test_create_sqlalchemy_url(self,
                 'port': 123,
                 'database': 'analyticsdb'
             }
+            if human_style_db_type in ['redshift', 'psql (redshift)']:
+                db_facts['query'] = {'keepalives': '1', 'keepalives_idle': '30'}
             actual_url = connect.create_sqlalchemy_url(db_facts)
-            assert str(actual_url) == expected_url
+            actual_url_str = str(actual_url)
+            self.assertEqual(actual_url_str, expected_url, "{}!={}".format(actual_url_str,
+                                                                           expected_url))
 
     @patch('records_mover.db.connect.db_facts_from_lpass')
     @patch('records_mover.db.connect.sa.create_engine')
@@ -45,10 +49,10 @@ def test_create_sqlalchemy_url_odbc_preferred(self,
                                                   mock_db_facts_from_lpass):
         expected_mappings = {
             'psql (redshift)':
-            'redshift://myuser:hunter1@myhost:123/analyticsdb',
+            'redshift://myuser:hunter1@myhost:123/analyticsdb?keepalives=1&keepalives_idle=30',
 
             'redshift':
-            'redshift://myuser:hunter1@myhost:123/analyticsdb',
+            'redshift://myuser:hunter1@myhost:123/analyticsdb?keepalives=1&keepalives_idle=30',
 
             'vertica':
             'vertica+pyodbc:///?odbc_connect=Driver'
@@ -71,5 +75,7 @@ def test_create_sqlalchemy_url_odbc_preferred(self,
                 'port': 123,
                 'database': 'analyticsdb'
             }
+            if human_style_db_type == 'redshift':
+                db_facts['query'] = {'keepalives': '1', 'keepalives_idle': '30'}
             actual_url = connect.create_sqlalchemy_url(db_facts, prefer_odbc=True)
             assert str(actual_url) == expected_url