update how we write to database

src/pg_client.py

@@ -72,18 +72,21 @@ def get_solver_rewards(
 
         # querying the prod database
         log.info("Setting tcp_keepalives_idle to 900 for prod connection")
+        # set tcp_keepalive_idle to not time out behind firewall
         with self.connections[0].connect() as connection:
-            # set tcp_keepalive_idle to not time out behind firewall
-            connection.execute(text("SET tcp_keepalives_idle = 900;"))
-            log.info("Running prod query for first connection (in get_solver_rewards)")
-            results.append(
-                self.exec_query(query=batch_reward_query_prod, engine=connection)
-            )
+            with connection.begin():
+                connection.execute(text("SET tcp_keepalives_idle = 900;"))
+        log.info("Running prod query for first connection (in get_solver_rewards)")
+        results.append(
+            self.exec_query(query=batch_reward_query_prod, engine=self.connections[0])
+        )
         # query for barn database
-        with self.connections[1].connect() as connection:
+        if len(self.connections) > 1:  # this is required due to our test setup
             log.info("Running barn query on other connections (in get_solver_rewards")
             results.append(
-                self.exec_query(query=batch_reward_query_barn, engine=connection)
+                self.exec_query(
+                    query=batch_reward_query_barn, engine=self.connections[1]
+                )
             )
 
         results_df = pd.concat(results)

tests/queries/test_batch_rewards.py

@@ -2,6 +2,7 @@
 
 import pandas.testing
 from pandas import DataFrame
+from sqlalchemy import text
 
 from src.config import RewardConfig, Network
 from src.pg_client import MultiInstanceDBFetcher
@@ -17,7 +18,9 @@ def setUp(self) -> None:
         with open(
             "./tests/queries/batch_rewards_test_db.sql", "r", encoding="utf-8"
         ) as file:
-            self.fetcher.connections[0].execute(file.read())
+            with self.fetcher.connections[0].connect() as connection:
+                with connection.begin():
+                    connection.execute(text(file.read()))
 
     def test_get_batch_rewards(self):
         start_block, end_block = "0", "100"

tests/queries/test_quote_rewards.py

@@ -2,6 +2,7 @@
 
 import pandas.testing
 from pandas import DataFrame
+from sqlalchemy import text
 
 from src.pg_client import MultiInstanceDBFetcher
 
@@ -13,7 +14,9 @@ def setUp(self) -> None:
         with open(
             "./tests/queries/quote_rewards_test_db.sql", "r", encoding="utf-8"
         ) as file:
-            self.fetcher.connections[0].execute(file.read())
+            with self.fetcher.connections[0].connect() as connection:
+                with connection.begin():
+                    connection.execute(text(file.read()))
 
     def test_get_quote_rewards(self):
         start_block, end_block = "0", "100"