wcp-library 1.5.8__py3-none-any.whl → 1.6.1__py3-none-any.whl

wcp_library/sql/postgres.py

@@ -1,6 +1,5 @@
 import logging
-from contextlib import _GeneratorContextManager
-from typing import Optional, Any
+from typing import Optional
 
 import numpy as np
 import pandas as pd
@@ -128,7 +127,7 @@ class PostgresConnection(object):
         else:
             self._connection = self.connection
 
-    def _get_connection(self) -> _GeneratorContextManager[Any] | Connection:
+    def _get_connection(self) -> Connection:
         """
         Get the connection object
 
@@ -136,7 +135,7 @@ class PostgresConnection(object):
         """
 
         if self.use_pool:
-            connection = self._session_pool.connection()
+            connection = self._session_pool.getconn()
             return connection
         else:
             if self._connection is None or self._connection.closed:
@@ -186,6 +185,9 @@ class PostgresConnection(object):
         with connection:
             connection.execute(query)
 
+        if self.use_pool:
+            self._session_pool.putconn(connection)
+
     @retry
     def safe_execute(self, query: SQL | str, packed_values: dict) -> None:
         """
@@ -200,6 +202,9 @@ class PostgresConnection(object):
         with connection:
             connection.execute(query, packed_values)
 
+        if self.use_pool:
+            self._session_pool.putconn(connection)
+
     @retry
     def execute_multiple(self, queries: list[tuple[SQL | str, dict]]) -> None:
         """
@@ -219,6 +224,9 @@ class PostgresConnection(object):
                 else:
                     connection.execute(query)
 
+        if self.use_pool:
+            self._session_pool.putconn(connection)
+
     @retry
     def execute_many(self, query: SQL | str, dictionary: list[dict]) -> None:
         """
@@ -234,6 +242,9 @@ class PostgresConnection(object):
             cursor = connection.cursor()
             cursor.executemany(query, dictionary)
 
+        if self.use_pool:
+            self._session_pool.putconn(connection)
+
     @retry
     def fetch_data(self, query: SQL | str, packed_data=None) -> list[tuple]:
         """
@@ -252,6 +263,9 @@ class PostgresConnection(object):
             else:
                 cursor.execute(query)
             rows = cursor.fetchall()
+
+        if self.use_pool:
+            self._session_pool.putconn(connection)
         return rows
 
     @retry

{wcp_library-1.5.8.dist-info → wcp_library-1.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: wcp-library
-Version: 1.5.8
+Version: 1.6.1
 Summary: Common utilites for internal development at WCP
 Author: Mitch-Petersen
 Author-email: mitch.petersen@wcap.ca

{wcp_library-1.5.8.dist-info → wcp_library-1.6.1.dist-info}/RECORD

@@ -18,8 +18,8 @@ wcp_library/selenium/_selenium_driver.py,sha256=Es83ISNryOVhYGDo_xrr8nCLsbnVBeXs
 wcp_library/selenium/selenium_helper.py,sha256=OJvC0RlblfMWhmjNxiY582yuzkDgjzFZMuRz_8to3PY,2000
 wcp_library/sql/__init__.py,sha256=f0JyVyFyUjueQ1XdVmCgmRjzkFENmFKl3BOzrvv3i1I,2422
 wcp_library/sql/oracle.py,sha256=L2UG-0jQNYSvTy9OAxRjWLIEIpfbLrIXeddYfuC3ipA,18812
-wcp_library/sql/postgres.py,sha256=OH0p7sjQDyv1zqFfnvrsFh6R82LDh-fB5HW_FDfOMn8,18329
+wcp_library/sql/postgres.py,sha256=I6Mw7QnXHu-aIOe8CiLcqgJ8wNzyjrTy3_bhfEV-2pI,18631
 wcp_library/time.py,sha256=ktSzhK7SZnGPNXobFexnhFGQAcriLCJQKxnO0fed8fQ,1740
-wcp_library-1.5.8.dist-info/METADATA,sha256=Kul1UZxy1MNUSUpNVUh0qvUWg5ohIsL6tA2f0Y7zkgo,1393
-wcp_library-1.5.8.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-wcp_library-1.5.8.dist-info/RECORD,,
+wcp_library-1.6.1.dist-info/METADATA,sha256=POlu-hZub7JLGST9pp2bS4wv4YQGgi2PbAMT76kJXUE,1393
+wcp_library-1.6.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+wcp_library-1.6.1.dist-info/RECORD,,