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

wcp_library/sql/__init__.py

@@ -9,55 +9,60 @@ import psycopg
 logger = logging.getLogger(__name__)
 
 
-def retry(f: callable) -> callable:
+def retry(func: callable) -> callable:
     """
     Decorator to retry a function
 
-    :param f: function
+    :param func: function
     :return: function
     """
 
-    @wraps(f)
+    @wraps(func)
     def wrapper(self, *args, **kwargs):
         self._retry_count = 0
         while True:
             try:
-                return f(self, *args, **kwargs)
+                return func(self, *args, **kwargs)
             except (oracledb.OperationalError, oracledb.DatabaseError, psycopg.OperationalError) as e:
-                error_obj, = e.args
-                if error_obj.full_code in self.retry_error_codes and self._retry_count < self.retry_limit:
-                    self._retry_count += 1
-                    logger.debug("Oracle connection error")
-                    logger.debug(error_obj.message)
-                    logger.info("Waiting 5 minutes before retrying Oracle connection")
-                    sleep(300)
-                else:
-                    raise e
+                if isinstance(e, (oracledb.OperationalError, oracledb.DatabaseError, psycopg.OperationalError, psycopg.DatabaseError)):
+                    error_obj, = e.args
+                    if error_obj.full_code in self.retry_error_codes and self._retry_count < self.retry_limit:
+                        self._retry_count += 1
+                        logger.debug("Oracle connection error")
+                        logger.debug(error_obj.message)
+                        logger.info("Waiting 5 minutes before retrying Oracle connection")
+                        sleep(300)
+                    else:
+                        raise e
+                raise e
     return wrapper
 
 
-def async_retry(f: callable) -> callable:
+def async_retry(func: callable) -> callable:
     """
     Decorator to retry a function
 
-    :param f: function
+    :param func: function
     :return: function
     """
 
-    @wraps(f)
+    @wraps(func)
     async def wrapper(self, *args, **kwargs):
         self._retry_count = 0
         while True:
             try:
-                return await f(self, *args, **kwargs)
+                return await func(self, *args, **kwargs)
             except (oracledb.OperationalError, oracledb.DatabaseError, psycopg.OperationalError) as e:
-                error_obj, = e.args
-                if error_obj.full_code in self.retry_error_codes and self._retry_count < self.retry_limit:
-                    self._retry_count += 1
-                    logger.debug(f"{self._db_service} connection error")
-                    logger.debug(error_obj.message)
-                    logger.info("Waiting 5 minutes before retrying Oracle connection")
-                    await asyncio.sleep(300)
+                if isinstance(e, (oracledb.OperationalError, oracledb.DatabaseError, psycopg.OperationalError, psycopg.DatabaseError)):
+                    error_obj, = e.args
+                    if error_obj.full_code in self.retry_error_codes and self._retry_count < self.retry_limit:
+                        self._retry_count += 1
+                        logger.debug(f"{self._db_service} connection error")
+                        logger.debug(error_obj.message)
+                        logger.info("Waiting 5 minutes before retrying Oracle connection")
+                        await asyncio.sleep(300)
+                    else:
+                        raise e
                 else:
                     raise e
     return wrapper

wcp_library/sql/oracle.py

@@ -4,7 +4,7 @@ from typing import Optional
 import numpy as np
 import pandas as pd
 import oracledb
-from oracledb import ConnectionPool, AsyncConnectionPool
+from oracledb import ConnectionPool, AsyncConnectionPool, Connection, AsyncConnection
 
 from wcp_library.sql import retry, async_retry
 
@@ -12,7 +12,7 @@ logger = logging.getLogger(__name__)
 
 
 def _connect_warehouse(username: str, password: str, hostname: str, port: int, database: str, min_connections: int,
-                       max_connections: int) -> ConnectionPool:
+                       max_connections: int, use_pool: bool) -> ConnectionPool | Connection:
     """
     Create Warehouse Connection
 
@@ -23,23 +23,34 @@ def _connect_warehouse(username: str, password: str, hostname: str, port: int, d
     :param database: database
     :param min_connections:
     :param max_connections:
-    :return: session_pool
+    :param use_pool: use connection pool
+    :return: session_pool | connection
     """
 
-    dsn = oracledb.makedsn(hostname, port, sid=database)
-    session_pool = oracledb.create_pool(
-        user=username,
-        password=password,
-        dsn=dsn,
-        min=min_connections,
-        max=max_connections,
-        increment=1,
-    )
-    return session_pool
+    if use_pool:
+        logger.debug(f"Creating connection pool with min size {min_connections} and max size {max_connections}")
+        dsn = oracledb.makedsn(hostname, port, sid=database)
+        session_pool = oracledb.create_pool(
+            user=username,
+            password=password,
+            dsn=dsn,
+            min=min_connections,
+            max=max_connections,
+            increment=1,
+        )
+        return session_pool
+    else:
+        logger.debug("Creating single connection")
+        connection = oracledb.connect(
+            user=username,
+            password=password,
+            dsn=oracledb.makedsn(hostname, port, sid=database)
+        )
+        return connection
 
 
 async def _async_connect_warehouse(username: str, password: str, hostname: str, port: int, database: str,
-                                   min_connections: int, max_connections: int) -> AsyncConnectionPool:
+                                   min_connections: int, max_connections: int, use_pool: bool) -> AsyncConnectionPool | AsyncConnection:
     """
     Create Warehouse Connection
 
@@ -50,19 +61,30 @@ async def _async_connect_warehouse(username: str, password: str, hostname: str,
     :param database: database
     :param min_connections:
     :param max_connections:
-    :return: session_pool
+    :param use_pool: use connection pool
+    :return: session_pool | connection
     """
 
-    dsn = oracledb.makedsn(hostname, port, sid=database)
-    session_pool = oracledb.create_pool_async(
-        user=username,
-        password=password,
-        dsn=dsn,
-        min=min_connections,
-        max=max_connections,
-        increment=1
-    )
-    return session_pool
+    if use_pool:
+        logger.debug(f"Creating async connection pool with min size {min_connections} and max size {max_connections}")
+        dsn = oracledb.makedsn(hostname, port, sid=database)
+        session_pool = oracledb.create_pool_async(
+            user=username,
+            password=password,
+            dsn=dsn,
+            min=min_connections,
+            max=max_connections,
+            increment=1
+        )
+        return session_pool
+    else:
+        logger.debug("Creating single async connection")
+        connection = await oracledb.connect_async(
+            user=username,
+            password=password,
+            dsn=oracledb.makedsn(hostname, port, sid=database)
+        )
+        return connection
 
 
 """~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"""
@@ -75,15 +97,17 @@ class OracleConnection(object):
     :return: None
     """
 
-    def __init__(self, min_connections: int = 2, max_connections: int = 5):
+    def __init__(self, use_pool: bool = False, min_connections: int = 2, max_connections: int = 5):
         self._username: Optional[str] = None
         self._password: Optional[str] = None
         self._hostname: Optional[str] = None
         self._port: Optional[int] = None
         self._database: Optional[str] = None
         self._sid: Optional[str] = None
+        self._connection: Optional[Connection] = None
         self._session_pool: Optional[ConnectionPool] = None
 
+        self.use_pool = use_pool
         self.min_connections = min_connections
         self.max_connections = max_connections
 
@@ -101,8 +125,27 @@ class OracleConnection(object):
 
         sid_or_service = self._database if self._database else self._sid
 
-        self._session_pool = _connect_warehouse(self._username, self._password, self._hostname, self._port,
-                                                sid_or_service, self.min_connections, self.max_connections)
+        connection = _connect_warehouse(self._username, self._password, self._hostname, self._port,
+                                        sid_or_service, self.min_connections, self.max_connections, self.use_pool)
+
+        if self.use_pool:
+            self._session_pool = connection
+        else:
+            self._connection = connection
+
+    def _get_connection(self) -> Connection:
+        """
+        Get the connection, either from the pool or create a new one
+
+        :return: Connection
+        """
+
+        if self.use_pool:
+            return self._session_pool.acquire()
+        else:
+            if not self._connection or not self._connection.is_healthy():
+                self._connect()
+            return self._connection
 
     def set_user(self, credentials_dict: dict) -> None:
         """
@@ -131,7 +174,12 @@ class OracleConnection(object):
         :return: None
         """
 
-        self._session_pool.close()
+        if self.use_pool:
+            self._session_pool.close()
+        else:
+            if self._connection and self._connection.is_healthy():
+                self._connection.close()
+            self._connection = None
 
     @retry
     def execute(self, query: str) -> None:
@@ -142,11 +190,11 @@ class OracleConnection(object):
         :return: None
         """
 
-        connection = self._session_pool.acquire()
-        cursor = connection.cursor()
-        cursor.execute(query)
-        connection.commit()
-        self._session_pool.release(connection)
+        connection = self._get_connection()
+        with connection:
+            cursor = connection.cursor()
+            cursor.execute(query)
+            connection.commit()
 
     @retry
     def safe_execute(self, query: str, packed_values: dict) -> None:
@@ -158,14 +206,14 @@ class OracleConnection(object):
         :return: None
         """
 
-        connection = self._session_pool.acquire()
-        cursor = connection.cursor()
-        cursor.execute(query, packed_values)
-        connection.commit()
-        self._session_pool.release(connection)
+        connection = self._get_connection()
+        with connection:
+            cursor = connection.cursor()
+            cursor.execute(query, packed_values)
+            connection.commit()
 
     @retry
-    def execute_multiple(self, queries: list[list[str, dict]]) -> None:
+    def execute_multiple(self, queries: list[tuple[str, dict]]) -> None:
         """
         Execute multiple queries
 
@@ -173,17 +221,17 @@ class OracleConnection(object):
         :return: None
         """
 
-        connection = self._session_pool.acquire()
-        cursor = connection.cursor()
-        for item in queries:
-            query = item[0]
-            packed_values = item[1]
-            if packed_values:
-                cursor.execute(query, packed_values)
-            else:
-                cursor.execute(query)
-        connection.commit()
-        self._session_pool.release(connection)
+        connection = self._get_connection()
+        with connection:
+            cursor = connection.cursor()
+            for item in queries:
+                query = item[0]
+                packed_values = item[1]
+                if packed_values:
+                    cursor.execute(query, packed_values)
+                else:
+                    cursor.execute(query)
+            connection.commit()
 
     @retry
     def execute_many(self, query: str, dictionary: list[dict]) -> None:
@@ -195,11 +243,11 @@ class OracleConnection(object):
         :return: None
         """
 
-        connection = self._session_pool.acquire()
-        cursor = connection.cursor()
-        cursor.executemany(query, dictionary)
-        connection.commit()
-        self._session_pool.release(connection)
+        connection = self._get_connection()
+        with connection:
+            cursor = connection.cursor()
+            cursor.executemany(query, dictionary)
+            connection.commit()
 
     @retry
     def fetch_data(self, query: str, packed_data=None) -> list:
@@ -211,14 +259,14 @@ class OracleConnection(object):
         :return: rows
         """
 
-        connection = self._session_pool.acquire()
-        cursor = connection.cursor()
-        if packed_data:
-            cursor.execute(query, packed_data)
-        else:
-            cursor.execute(query)
-        rows = cursor.fetchall()
-        self._session_pool.release(connection)
+        connection = self._get_connection()
+        with connection:
+            cursor = connection.cursor()
+            if packed_data:
+                cursor.execute(query, packed_data)
+            else:
+                cursor.execute(query)
+            rows = cursor.fetchall()
         return rows
 
     @retry
@@ -267,7 +315,7 @@ class OracleConnection(object):
             dfObj = dfObj.replace({np.nan: None})
         main_dict = dfObj.to_dict('records')
 
-        query = """INSERT INTO {} ({}) VALUES ({})""".format(outputTableName, col, bind)
+        query = f"""INSERT INTO {outputTableName} ({col}) VALUES ({bind})"""
         self.execute_many(query, main_dict)
 
     @retry
@@ -279,7 +327,7 @@ class OracleConnection(object):
         :return: None
         """
 
-        truncateQuery = """TRUNCATE TABLE {}""".format(tableName)
+        truncateQuery = f"""TRUNCATE TABLE {tableName}"""
         self.execute(truncateQuery)
 
     @retry
@@ -291,7 +339,7 @@ class OracleConnection(object):
         :return: None
         """
 
-        deleteQuery = """DELETE FROM {}""".format(tableName)
+        deleteQuery = f"""DELETE FROM {tableName}"""
         self.execute(deleteQuery)
 
     def __del__(self) -> None:
@@ -301,7 +349,12 @@ class OracleConnection(object):
         :return: None
         """
 
-        self._session_pool.close()
+        if self.use_pool:
+            self._session_pool.close()
+        else:
+            if self._connection and self._connection.is_healthy():
+                self._connection.close()
+            self._connection = None
 
 
 class AsyncOracleConnection(object):
@@ -311,7 +364,7 @@ class AsyncOracleConnection(object):
     :return: None
     """
 
-    def __init__(self, min_connections: int = 2, max_connections: int = 5):
+    def __init__(self, use_pool: bool = False, min_connections: int = 2, max_connections: int = 5):
         self._db_service: str = "Oracle"
         self._username: Optional[str] = None
         self._password: Optional[str] = None
@@ -319,8 +372,10 @@ class AsyncOracleConnection(object):
         self._port: Optional[int] = None
         self._database: Optional[str] = None
         self._sid: Optional[str] = None
+        self._connection: Optional[AsyncConnection] = None
         self._session_pool: Optional[AsyncConnectionPool] = None
 
+        self.use_pool = use_pool
         self.min_connections = min_connections
         self.max_connections = max_connections
 
@@ -338,8 +393,28 @@ class AsyncOracleConnection(object):
 
         sid_or_service = self._database if self._database else self._sid
 
-        self._session_pool = await _async_connect_warehouse(self._username, self._password, self._hostname, self._port,
-                                                            sid_or_service, self.min_connections, self.max_connections)
+        connection = await _async_connect_warehouse(self._username, self._password, self._hostname, self._port,
+                                                    sid_or_service, self.min_connections, self.max_connections,
+                                                    self.use_pool)
+
+        if self.use_pool:
+            self._session_pool = connection
+        else:
+            self._connection = connection
+
+    async def _get_connection(self) -> AsyncConnection:
+        """
+        Get the connection, either from the pool or create a new one
+
+        :return: AsyncConnection
+        """
+
+        if self.use_pool:
+            return await self._session_pool.acquire()
+        else:
+            if not self._connection or not self._connection.is_healthy():
+                await self._connect()
+            return self._connection
 
     async def set_user(self, credentials_dict: dict) -> None:
         """
@@ -368,7 +443,12 @@ class AsyncOracleConnection(object):
         :return: None
         """
 
-        await self._session_pool.close()
+        if self.use_pool:
+            await self._session_pool.close()
+        else:
+            if self._connection and self._connection.is_healthy():
+                await self._connection.close()
+            self._connection = None
 
     @async_retry
     async def execute(self, query: str) -> None:
@@ -379,7 +459,8 @@ class AsyncOracleConnection(object):
         :return: None
         """
 
-        async with self._session_pool.acquire() as connection:
+        connection = await self._get_connection()
+        async with connection:
             with connection.cursor() as cursor:
                 await cursor.execute(query)
                 await connection.commit()
@@ -394,13 +475,14 @@ class AsyncOracleConnection(object):
         :return: None
         """
 
-        async with self._session_pool.acquire() as connection:
+        connection = await self._get_connection()
+        async with connection:
             with connection.cursor() as cursor:
                 await cursor.execute(query, packed_values)
                 await connection.commit()
 
     @async_retry
-    async def execute_multiple(self, queries: list[list[str, dict]]) -> None:
+    async def execute_multiple(self, queries: list[tuple[str, dict]]) -> None:
         """
         Execute multiple queries
 
@@ -408,7 +490,8 @@ class AsyncOracleConnection(object):
         :return: None
         """
 
-        async with self._session_pool.acquire() as connection:
+        connection = await self._get_connection()
+        async with connection:
             with connection.cursor() as cursor:
                 for item in queries:
                     query = item[0]
@@ -429,7 +512,8 @@ class AsyncOracleConnection(object):
         :return: None
         """
 
-        async with self._session_pool.acquire() as connection:
+        connection = await self._get_connection()
+        async with connection:
             with connection.cursor() as cursor:
                 await cursor.executemany(query, dictionary)
                 await connection.commit()
@@ -444,7 +528,8 @@ class AsyncOracleConnection(object):
         :return: rows
         """
 
-        async with self._session_pool.acquire() as connection:
+        connection = await self._get_connection()
+        async with connection:
             with connection.cursor() as cursor:
                 if packed_data:
                     await cursor.execute(query, packed_data)
@@ -499,7 +584,7 @@ class AsyncOracleConnection(object):
             dfObj = dfObj.replace({np.nan: None})
         main_dict = dfObj.to_dict('records')
 
-        query = """INSERT INTO {} ({}) VALUES ({})""".format(outputTableName, col, bind)
+        query = f"""INSERT INTO {outputTableName} ({col}) VALUES ({bind})"""
         await self.execute_many(query, main_dict)
 
     @async_retry
@@ -511,7 +596,7 @@ class AsyncOracleConnection(object):
         :return: None
         """
 
-        truncateQuery = """TRUNCATE TABLE {}""".format(tableName)
+        truncateQuery = f"""TRUNCATE TABLE {tableName}"""
         await self.execute(truncateQuery)
 
     @async_retry
@@ -523,5 +608,5 @@ class AsyncOracleConnection(object):
         :return: None
         """
 
-        deleteQuery = """DELETE FROM {}""".format(tableName)
+        deleteQuery = f"""DELETE FROM {tableName}"""
         await self.execute(deleteQuery)

wcp_library/sql/postgres.py

@@ -1,8 +1,11 @@
 import logging
-from typing import Optional
+from contextlib import _GeneratorContextManager
+from typing import Optional, Any
 
 import numpy as np
 import pandas as pd
+import psycopg
+from psycopg import AsyncConnection, Connection
 from psycopg.conninfo import make_conninfo
 from psycopg.sql import SQL
 from psycopg_pool import AsyncConnectionPool, ConnectionPool
@@ -13,7 +16,7 @@ logger = logging.getLogger(__name__)
 
 
 def _connect_warehouse(username: str, password: str, hostname: str, port: int, database: str, min_connections: int,
-                       max_connections: int) -> ConnectionPool:
+                       max_connections: int, use_pool: bool) -> Connection | ConnectionPool:
     """
     Create Warehouse Connection
 
@@ -30,18 +33,24 @@ def _connect_warehouse(username: str, password: str, hostname: str, port: int, d
     conn_string = f"dbname={database} user={username} password={password} host={hostname} port={port}"
     conninfo = make_conninfo(conn_string)
 
-    session_pool = ConnectionPool(
-        conninfo=conninfo,
-        min_size=min_connections,
-        max_size=max_connections,
-        kwargs={'options': '-c datestyle=ISO,YMD'},
-        open=True
-    )
-    return session_pool
+    if use_pool:
+        logger.debug(f"Creating connection pool with min size {min_connections} and max size {max_connections}")
+        session_pool = ConnectionPool(
+            conninfo=conninfo,
+            min_size=min_connections,
+            max_size=max_connections,
+            kwargs={'options': '-c datestyle=ISO,YMD'},
+            open=True
+        )
+        return session_pool
+    else:
+        logger.debug("Creating single connection")
+        connection = psycopg.connect(conninfo=conninfo, options='-c datestyle=ISO,YMD')
+        return connection
 
 
 async def _async_connect_warehouse(username: str, password: str, hostname: str, port: int, database: str, min_connections: int,
-                             max_connections: int) -> AsyncConnectionPool:
+                                   max_connections: int, use_pool: bool) -> AsyncConnection | AsyncConnectionPool:
     """
     Create Warehouse Connection
 
@@ -58,14 +67,20 @@ async def _async_connect_warehouse(username: str, password: str, hostname: str,
     conn_string = f"dbname={database} user={username} password={password} host={hostname} port={port}"
     conninfo = make_conninfo(conn_string)
 
-    session_pool = AsyncConnectionPool(
-        conninfo=conninfo,
-        min_size=min_connections,
-        max_size=max_connections,
-        kwargs={"options": "-c datestyle=ISO,YMD"},
-        open=False
-    )
-    return session_pool
+    if use_pool:
+        logger.debug(f"Creating async connection pool with min size {min_connections} and max size {max_connections}")
+        session_pool = AsyncConnectionPool(
+            conninfo=conninfo,
+            min_size=min_connections,
+            max_size=max_connections,
+            kwargs={"options": "-c datestyle=ISO,YMD"},
+            open=False
+        )
+        return session_pool
+    else:
+        logger.debug("Creating single async connection")
+        connection = await AsyncConnection.connect(conninfo=conninfo, options='-c datestyle=ISO,YMD')
+        return connection
 
 
 """~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"""
@@ -78,14 +93,16 @@ class PostgresConnection(object):
     :return: None
     """
 
-    def __init__(self, min_connections: int = 2, max_connections: int = 5):
+    def __init__(self, use_pool: bool = False, min_connections: int = 2, max_connections: int = 5):
         self._username: Optional[str] = None
         self._password: Optional[str] = None
         self._hostname: Optional[str] = None
         self._port: Optional[int] = None
         self._database: Optional[str] = None
+        self._connection: Optional[Connection] = None
         self._session_pool: Optional[ConnectionPool] = None
 
+        self.use_pool = use_pool
         self.min_connections = min_connections
         self.max_connections = max_connections
 
@@ -101,8 +118,30 @@ class PostgresConnection(object):
         :return: None
         """
 
-        self._session_pool = _connect_warehouse(self._username, self._password, self._hostname, self._port,
-                                               self._database, self.min_connections, self.max_connections)
+
+        self.connection = _connect_warehouse(self._username, self._password, self._hostname, self._port,
+                                             self._database, self.min_connections, self.max_connections, self.use_pool)
+
+        if self.use_pool:
+            self._session_pool = self.connection
+            self._session_pool.open()
+        else:
+            self._connection = self.connection
+
+    def _get_connection(self) -> _GeneratorContextManager[Any] | Connection:
+        """
+        Get the connection object
+
+        :return: connection
+        """
+
+        if self.use_pool:
+            connection = self._session_pool.connection()
+            return connection
+        else:
+            if self._connection is None or self._connection.closed:
+                self._connect()
+            return self._connection
 
     def set_user(self, credentials_dict: dict) -> None:
         """
@@ -127,7 +166,12 @@ class PostgresConnection(object):
         :return: None
         """
 
-        self._session_pool.close()
+        if self.use_pool:
+            self._session_pool.close()
+        else:
+            if self._connection is not None and not self._connection.closed:
+                self._connection.close()
+            self._connection = None
 
     @retry
     def execute(self, query: SQL | str) -> None:
@@ -138,7 +182,8 @@ class PostgresConnection(object):
         :return: None
         """
 
-        with self._session_pool.connection() as connection:
+        connection = self._get_connection()
+        with connection:
             connection.execute(query)
 
     @retry
@@ -151,11 +196,12 @@ class PostgresConnection(object):
         :return: None
         """
 
-        with self._session_pool.connection() as connection:
+        connection = self._get_connection()
+        with connection:
             connection.execute(query, packed_values)
 
     @retry
-    def execute_multiple(self, queries: list[list[SQL | str, dict]]) -> None:
+    def execute_multiple(self, queries: list[tuple[SQL | str, dict]]) -> None:
         """
         Execute multiple queries
 
@@ -163,7 +209,8 @@ class PostgresConnection(object):
         :return: None
         """
 
-        with self._session_pool.connection() as connection:
+        connection = self._get_connection()
+        with connection:
             for item in queries:
                 query = item[0]
                 packed_values = item[1]
@@ -182,7 +229,8 @@ class PostgresConnection(object):
         :return: None
         """
 
-        with self._session_pool.connection() as connection:
+        connection = self._get_connection()
+        with connection:
             cursor = connection.cursor()
             cursor.executemany(query, dictionary)
 
@@ -196,7 +244,8 @@ class PostgresConnection(object):
         :return: rows
         """
 
-        with self._session_pool.connection() as connection:
+        connection = self._get_connection()
+        with connection:
             cursor = connection.cursor()
             if packed_data:
                 cursor.execute(query, packed_data)
@@ -255,7 +304,7 @@ class PostgresConnection(object):
                 if record[key] == '':
                     record[key] = None
 
-        query = """INSERT INTO {} ({}) VALUES ({})""".format(outputTableName, col, params)
+        query = f"INSERT INTO {outputTableName} ({col}) VALUES ({params})"
         self.execute_many(query, main_dict)
 
     @retry
@@ -267,7 +316,7 @@ class PostgresConnection(object):
         :return: None
         """
 
-        truncateQuery = """TRUNCATE TABLE {}""".format(tableName)
+        truncateQuery = f"TRUNCATE TABLE {tableName}"
         self.execute(truncateQuery)
 
     @retry
@@ -279,7 +328,7 @@ class PostgresConnection(object):
         :return: None
         """
 
-        deleteQuery = """DELETE FROM {}""".format(tableName)
+        deleteQuery = f"DELETE FROM {tableName}"
         self.execute(deleteQuery)
 
     def __del__(self) -> None:
@@ -289,7 +338,12 @@ class PostgresConnection(object):
         :return: None
         """
 
-        self._session_pool.close()
+        if self._session_pool is not None:
+            self._session_pool.close()
+        else:
+            if self._connection is not None and not self._connection.closed:
+                self._connection.close()
+            self._connection = None
 
 
 class AsyncPostgresConnection(object):
@@ -299,14 +353,16 @@ class AsyncPostgresConnection(object):
     :return: None
     """
 
-    def __init__(self, min_connections: int = 2, max_connections: int = 5):
+    def __init__(self, use_pool: bool = False, min_connections: int = 2, max_connections: int = 5):
         self._username: Optional[str] = None
         self._password: Optional[str] = None
         self._hostname: Optional[str] = None
         self._port: Optional[int] = None
         self._database: Optional[str] = None
+        self._connection: Optional[AsyncConnection] = None
         self._session_pool: Optional[AsyncConnectionPool] = None
 
+        self.use_pool = use_pool
         self.min_connections = min_connections
         self.max_connections = max_connections
 
@@ -322,9 +378,30 @@ class AsyncPostgresConnection(object):
         :return: None
         """
 
-        self._session_pool = await _async_connect_warehouse(self._username, self._password, self._hostname, self._port,
-                                                            self._database, self.min_connections, self.max_connections)
-        await self._session_pool.open()
+        connection = await _async_connect_warehouse(self._username, self._password, self._hostname, self._port,
+                                                    self._database, self.min_connections, self.max_connections,
+                                                    self.use_pool)
+        if self.use_pool:
+            self._session_pool = connection
+            await self._session_pool.open()
+        else:
+            self._connection = connection
+
+    async def _get_connection(self) -> AsyncConnection:
+        """
+        Get the connection object
+
+        :return: connection
+        """
+
+        if self.use_pool:
+            connection = await self._session_pool.getconn()
+            return connection
+        else:
+            if self._connection is None or self._connection.closed:
+                await self._connect()
+            return self._connection
+
 
     async def set_user(self, credentials_dict: dict) -> None:
         """
@@ -349,7 +426,12 @@ class AsyncPostgresConnection(object):
         :return: None
         """
 
-        await self._session_pool.close()
+        if self.use_pool:
+            await self._session_pool.close()
+        else:
+            if self._connection is not None and not self._connection.closed:
+                await self._connection.close()
+            self._connection = None
 
     @async_retry
     async def execute(self, query: SQL | str) -> None:
@@ -360,9 +442,13 @@ class AsyncPostgresConnection(object):
         :return: None
         """
 
-        async with self._session_pool.connection() as connection:
+        connection = await self._get_connection()
+        async with connection:
             await connection.execute(query)
 
+        if self.use_pool:
+            await self._session_pool.putconn(connection)
+
     @async_retry
     async def safe_execute(self, query: SQL | str, packed_values: dict) -> None:
         """
@@ -373,11 +459,15 @@ class AsyncPostgresConnection(object):
         :return: None
         """
 
-        async with self._session_pool.connection() as connection:
+        connection = await self._get_connection()
+        async with connection:
             await connection.execute(query, packed_values)
 
+        if self.use_pool:
+            await self._session_pool.putconn(connection)
+
     @async_retry
-    async def execute_multiple(self, queries: list[list[SQL | str, dict]]) -> None:
+    async def execute_multiple(self, queries: list[tuple[SQL | str, dict]]) -> None:
         """
         Execute multiple queries
 
@@ -385,7 +475,8 @@ class AsyncPostgresConnection(object):
         :return: None
         """
 
-        async with self._session_pool.connection() as connection:
+        connection = await self._get_connection()
+        async with connection:
             for item in queries:
                 query = item[0]
                 packed_values = item[1]
@@ -394,6 +485,9 @@ class AsyncPostgresConnection(object):
                 else:
                     await connection.execute(query)
 
+        if self.use_pool:
+            await self._session_pool.putconn(connection)
+
     @async_retry
     async def execute_many(self, query: SQL | str, dictionary: list[dict]) -> None:
         """
@@ -404,10 +498,14 @@ class AsyncPostgresConnection(object):
         :return: None
         """
 
-        async with self._session_pool.connection() as connection:
+        connection = await self._get_connection()
+        async with connection:
             cursor = connection.cursor()
             await cursor.executemany(query, dictionary)
 
+        if self.use_pool:
+            await self._session_pool.putconn(connection)
+
     @async_retry
     async def fetch_data(self, query: SQL | str, packed_data=None) -> list[tuple]:
         """
@@ -418,13 +516,17 @@ class AsyncPostgresConnection(object):
         :return: rows
         """
 
-        async with self._session_pool.connection() as connection:
+        connection = await self._get_connection()
+        async with connection:
             cursor = connection.cursor()
             if packed_data:
                 await cursor.execute(query, packed_data)
             else:
                 await cursor.execute(query)
             rows = await cursor.fetchall()
+
+        if self.use_pool:
+            await self._session_pool.putconn(connection)
         return rows
 
     @async_retry
@@ -448,7 +550,7 @@ class AsyncPostgresConnection(object):
             params = param_list[0]
 
         main_dict = df.to_dict('records')
-        query = """DELETE FROM {} WHERE {}""".format(outputTableName, params)
+        query = f"DELETE FROM {outputTableName} WHERE {params}"
         await self.execute_many(query, main_dict)
 
     @async_retry
@@ -477,7 +579,7 @@ class AsyncPostgresConnection(object):
                 if record[key] == '':
                     record[key] = None
 
-        query = """INSERT INTO {} ({}) VALUES ({})""".format(outputTableName, col, params)
+        query = f"INSERT INTO {outputTableName} ({col}) VALUES ({params})"
         await self.execute_many(query, main_dict)
 
     @async_retry
@@ -489,7 +591,7 @@ class AsyncPostgresConnection(object):
         :return: None
         """
 
-        truncateQuery = """TRUNCATE TABLE {}""".format(tableName)
+        truncateQuery = f"TRUNCATE TABLE {tableName}"
         await self.execute(truncateQuery)
 
     @async_retry
@@ -501,5 +603,5 @@ class AsyncPostgresConnection(object):
         :return: None
         """
 
-        deleteQuery = """DELETE FROM {}""".format(tableName)
+        deleteQuery = f"DELETE FROM {tableName}"
         await self.execute(deleteQuery)

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

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

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

@@ -16,10 +16,10 @@ wcp_library/logging.py,sha256=e6gG7HFgUrMajUZs4Gs0atFfOJJmdmxN0GerfynNWlY,2061
 wcp_library/selenium/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wcp_library/selenium/_selenium_driver.py,sha256=Es83ISNryOVhYGDo_xrr8nCLsbnVBeXsaz1sRiz7kLI,2961
 wcp_library/selenium/selenium_helper.py,sha256=OJvC0RlblfMWhmjNxiY582yuzkDgjzFZMuRz_8to3PY,2000
-wcp_library/sql/__init__.py,sha256=ocKH9-1HPDF5OhUKkpby3uNvQ6RBt2YD81nwwS90PgE,1980
-wcp_library/sql/oracle.py,sha256=jcGoH2iGkgxW-RJV5FRf_CwNzilkxPN6_pJ27te5Qe0,15815
-wcp_library/sql/postgres.py,sha256=lPcsmX0ejEgCRmNKWL9Wp_7R8sop5_KcWmpLieTGT1I,14938
+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/time.py,sha256=ktSzhK7SZnGPNXobFexnhFGQAcriLCJQKxnO0fed8fQ,1740
-wcp_library-1.5.6.dist-info/METADATA,sha256=7JyNWin_He1VyxYXAw8MKRVto7xv3kP7cVpY8ZInU0U,1393
-wcp_library-1.5.6.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-wcp_library-1.5.6.dist-info/RECORD,,
+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,,