pooling support

main.py

@@ -1,21 +1,44 @@
 from mssql_python import connect
+from mssql_python import enable_pooling
 from mssql_python import setup_logging
 import os
 import decimal
+import time
 
 setup_logging('stdout')
 
-conn_str = os.getenv("DB_CONNECTION_STRING")
-conn = connect(conn_str)
+# conn_str = os.getenv("DB_CONNECTION_STRING")
+conn_str = "Server=Saumya;DATABASE=master;UID=sa;PWD=HappyPass1234;Trust_Connection=yes;TrustServerCertificate=yes;"
+
+enable_pooling(max_size=10, idle_timeout=300)
+conn1 = connect(conn_str)
 
 # conn.autocommit = True
 
-cursor = conn.cursor()
-cursor.execute("SELECT database_id, name from sys.databases;")
-rows = cursor.fetchall()
+cursor1 = conn1.cursor()
+cursor1.execute("SELECT database_id, name from sys.databases;")
+rows = cursor1.fetchone()
+print (rows)
+
+print(conn1._conn)
+print("First time check")
+# time.sleep(10) 
+
+# cursor1.close()
+# conn1.close()
+print("Second time check")
+# time.sleep(10)
+
+conn2 = connect(conn_str)
+cursor2 = conn2.cursor()
+cursor2.execute("SELECT database_id, name from sys.databases;") 
+row2 = cursor2.fetchone()
+print(row2)
+
+print(conn2._conn)
+print("Third time check")
+# time.sleep(10)
 
-for row in rows:
-    print(f"Database ID: {row[0]}, Name: {row[1]}")
 
-cursor.close()
-conn.close()
+cursor2.close()
+conn2.close()

mssql_python/__init__.py

@@ -38,6 +38,7 @@
 # Connection Objects
 from .connection import Connection
 from .db_connection import connect
+from .db_connection import enable_pooling
 
 # Cursor Objects
 from .cursor import Cursor

mssql_python/connection.py

@@ -33,7 +33,7 @@ class Connection:
         close() -> None:
     """
 
-    def __init__(self, connection_str: str = "", autocommit: bool = False, attrs_before: dict = None, **kwargs) -> None:
+    def __init__(self, connection_str: str = "", autocommit: bool = False, attrs_before: dict = None, use_pool: bool = False, **kwargs) -> None:
         """
         Initialize the connection object with the specified connection string and parameters.
 
@@ -53,15 +53,17 @@ def __init__(self, connection_str: str = "", autocommit: bool = False, attrs_bef
         preparing it for further operations such as connecting to the 
         database, executing queries, etc.
         """
-        self.henv = None
-        self.hdbc = None
         self.connection_str = self._construct_connection_string(
             connection_str, **kwargs
         )
-        self._attrs_before = attrs_before
-        self._autocommit = autocommit  # Initialize _autocommit before calling _initializer
-        self._initializer()
-        self.setautocommit(autocommit)
+        self._attrs_before = attrs_before or {}
+        # self._conn = ddbc_bindings.Connection(self.connection_str, autocommit, use_pool)
+        # self._conn.connect(self._attrs_before)
+        # print("Connection string: ", self.connection_str)
+        self._conn = ddbc_bindings.Connection(self.connection_str, use_pool)
+        # print("Connection object: ", self._conn)
+        # self._autocommit = autocommit
+        # self.setautocommit(autocommit)
 
     def _construct_connection_string(self, connection_str: str = "", **kwargs) -> str:
         """
@@ -100,193 +102,15 @@ def _construct_connection_string(self, connection_str: str = "", **kwargs) -> st
             logger.info("Final connection string: %s", conn_str)
 
         return conn_str
-
-    def _is_closed(self) -> bool:
-        """
-        Check if the connection is closed.
-
-        Returns:
-            bool: True if the connection is closed, False otherwise.
-        """
-        return self.hdbc is None
 
-    def _initializer(self) -> None:
-        """
-        Initialize the environment and connection handles.
-
-        This method is responsible for setting up the environment and connection
-        handles, allocating memory for them, and setting the necessary attributes.
-        It should be called before establishing a connection to the database.
-        """
-        self._allocate_environment_handle()
-        self._set_environment_attributes()
-        self._allocate_connection_handle()
-        if self._attrs_before != {}:
-            self._apply_attrs_before()  # Apply pre-connection attributes
-        if self._autocommit:
-            self._set_connection_attributes(
-                ddbc_sql_const.SQL_ATTR_AUTOCOMMIT.value,
-                ddbc_sql_const.SQL_AUTOCOMMIT_ON.value,
-            )
-        self._connect_to_db()
-
-    def _apply_attrs_before(self):
-        """
-        Apply specific pre-connection attributes.
-        Currently, this method only processes an attribute with key 1256 (e.g., SQL_COPT_SS_ACCESS_TOKEN)
-        if present in `self._attrs_before`. Other attributes are ignored.
-
-        Returns:
-            bool: True.
-        """
-
-        if ENABLE_LOGGING:
-            logger.info("Attempting to apply pre-connection attributes (attrs_before): %s", self._attrs_before)
-
-        if not isinstance(self._attrs_before, dict):
-            if self._attrs_before is not None and ENABLE_LOGGING:
-                logger.warning(
-                    f"_attrs_before is of type {type(self._attrs_before).__name__}, "
-                    f"expected dict. Skipping attribute application."
-                )
-            elif self._attrs_before is None and ENABLE_LOGGING:
-                 logger.debug("_attrs_before is None. No pre-connection attributes to apply.")
-            return True # Exit if _attrs_before is not a dictionary or is None
-
-        for key, value in self._attrs_before.items():
-            ikey = None
-            if isinstance(key, int):
-                ikey = key
-            elif isinstance(key, str) and key.isdigit():
-                try:
-                    ikey = int(key)
-                except ValueError:
-                    if ENABLE_LOGGING:
-                        logger.debug(
-                            f"Skipping attribute with key '{key}' in attrs_before: "
-                            f"could not convert string to int."
-                        )
-                    continue # Skip if string key is not a valid integer
-            else:
-                if ENABLE_LOGGING:
-                    logger.debug(
-                        f"Skipping attribute with key '{key}' in attrs_before due to "
-                        f"unsupported key type: {type(key).__name__}. Expected int or string representation of an int."
-                    )
-                continue  # Skip keys that are not int or string representation of an int
-
-            if ikey == ddbc_sql_const.SQL_COPT_SS_ACCESS_TOKEN.value:
-                if ENABLE_LOGGING:
-                    logger.info(
-                        f"Found attribute {ddbc_sql_const.SQL_COPT_SS_ACCESS_TOKEN.value}. Attempting to set it."
-                    )
-                self._set_connection_attributes(ikey, value)
-                if ENABLE_LOGGING:
-                    logger.info(
-                        f"Call to set attribute {ddbc_sql_const.SQL_COPT_SS_ACCESS_TOKEN.value} with value '{value}' completed."
-                    )
-                # If you expect only one such key, you could add 'break' here.
-            else:
-                if ENABLE_LOGGING:
-                    logger.debug(
-                        f"Ignoring attribute with key '{key}' (resolved to {ikey}) in attrs_before "
-                        f"as it is not the target attribute ({ddbc_sql_const.SQL_COPT_SS_ACCESS_TOKEN.value})."
-                    )
-        return True
-
-    def _allocate_environment_handle(self):
-        """
-        Allocate the environment handle.
-        """
-        ret, handle = ddbc_bindings.DDBCSQLAllocHandle(
-            ddbc_sql_const.SQL_HANDLE_ENV.value,  # SQL environment handle type
-            None
-        )
-        check_error(ddbc_sql_const.SQL_HANDLE_ENV.value, handle, ret)
-        self.henv = handle
-
-    def _set_environment_attributes(self):
-        """
-        Set the environment attributes.
-        """
-        ret = ddbc_bindings.DDBCSQLSetEnvAttr(
-            self.henv,  # Use the wrapper class
-            ddbc_sql_const.SQL_ATTR_DDBC_VERSION.value,  # Attribute
-            ddbc_sql_const.SQL_OV_DDBC3_80.value,  # String Length
-            0,  # Null-terminated string
-        )
-        check_error(ddbc_sql_const.SQL_HANDLE_ENV.value, self.henv, ret)
-
-    def _allocate_connection_handle(self):
-        """
-        Allocate the connection handle.
-        """
-        ret, handle = ddbc_bindings.DDBCSQLAllocHandle(
-            ddbc_sql_const.SQL_HANDLE_DBC.value,  # SQL connection handle type
-            self.henv
-        )
-        check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, handle, ret)
-        self.hdbc = handle
-
-    def _set_connection_attributes(self, ikey: int, ivalue: any) -> None:
-        """
-        Set the connection attributes before connecting.
-
-        Args:
-            ikey (int): The attribute key to set.
-            ivalue (Any): The value to set for the attribute. Can be bytes, bytearray, int, or unicode.
-            vallen (int): The length of the value.
-
-        Raises:
-            DatabaseError: If there is an error while setting the connection attribute.
-        """
-
-        ret = ddbc_bindings.DDBCSQLSetConnectAttr(
-            self.hdbc,  # Connection handle
-            ikey,  # Attribute
-            ivalue,  # Value
-        )
-        check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, self.hdbc, ret)
-
-    def _connect_to_db(self) -> None:
-        """
-        Establish a connection to the database.
-
-        This method is responsible for creating a connection to the specified database.
-        It does not take any arguments and does not return any value. The connection
-        details such as database name, user credentials, host, and port should be
-        configured within the class or passed during the class instantiation.
-
-        Raises:
-            DatabaseError: If there is an error while trying to connect to the database.
-            InterfaceError: If there is an error related to the database interface.
-        """
-        if ENABLE_LOGGING:
-            logger.info("Connecting to the database")
-        ret = ddbc_bindings.DDBCSQLDriverConnect(
-            self.hdbc,  # Connection handle (wrapper)
-            0,  # Window handle
-            self.connection_str,  # Connection string
-        )
-        check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, self.hdbc, ret)
-        if ENABLE_LOGGING:
-            logger.info("Connection established successfully.")
-
     @property
     def autocommit(self) -> bool:
         """
         Return the current autocommit mode of the connection.
         Returns:
             bool: True if autocommit is enabled, False otherwise.
         """
-        autocommit_mode = ddbc_bindings.DDBCSQLGetConnectionAttr(
-            self.hdbc,  # Connection handle (wrapper)
-            ddbc_sql_const.SQL_ATTR_AUTOCOMMIT.value,  # Attribute
-        )
-        check_error(
-            ddbc_sql_const.SQL_HANDLE_DBC.value, self.hdbc, autocommit_mode
-        )
-        return autocommit_mode == ddbc_sql_const.SQL_AUTOCOMMIT_ON.value
+        return self._conn.get_autocommit()
 
     @autocommit.setter
     def autocommit(self, value: bool) -> None:
@@ -296,20 +120,8 @@ def autocommit(self, value: bool) -> None:
             value (bool): True to enable autocommit, False to disable it.
         Returns:
             None
-        Raises:
-            DatabaseError: If there is an error while setting the autocommit mode.
         """
-        ret = ddbc_bindings.DDBCSQLSetConnectAttr(
-            self.hdbc,  # Connection handle
-            ddbc_sql_const.SQL_ATTR_AUTOCOMMIT.value,  # Attribute
-            (
-                ddbc_sql_const.SQL_AUTOCOMMIT_ON.value
-                if value
-                else ddbc_sql_const.SQL_AUTOCOMMIT_OFF.value
-            ),  # Value
-        )
-        check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, self.hdbc, ret)
-        self._autocommit = value
+        self.setautocommit(value)
         if ENABLE_LOGGING:
             logger.info("Autocommit mode set to %s.", value)
 
@@ -323,7 +135,8 @@ def setautocommit(self, value: bool = True) -> None:
         Raises:
             DatabaseError: If there is an error while setting the autocommit mode.
         """
-        self.autocommit = value
+        self._conn.set_autocommit(value)
+        self._autocommit = value
 
     def cursor(self) -> Cursor:
         """
@@ -340,9 +153,6 @@ def cursor(self) -> Cursor:
             DatabaseError: If there is an error while creating the cursor.
             InterfaceError: If there is an error related to the database interface.
         """
-        if self._is_closed():
-            # Cannot create a cursor if the connection is closed
-            raise Exception("Connection is closed. Cannot create cursor.")
         return Cursor(self)
 
     def commit(self) -> None:
@@ -357,17 +167,8 @@ def commit(self) -> None:
         Raises:
             DatabaseError: If there is an error while committing the transaction.
         """
-        if self._is_closed():
-            # Cannot commit if the connection is closed
-            raise Exception("Connection is closed. Cannot commit.")
-
         # Commit the current transaction
-        ret = ddbc_bindings.DDBCSQLEndTran(
-            ddbc_sql_const.SQL_HANDLE_DBC.value,  # Handle type
-            self.hdbc,  # Connection handle (wrapper)
-            ddbc_sql_const.SQL_COMMIT.value,  # Commit the transaction
-        )
-        check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, self.hdbc, ret)
+        self._conn.commit()
         if ENABLE_LOGGING:
             logger.info("Transaction committed successfully.")
 
@@ -382,17 +183,8 @@ def rollback(self) -> None:
         Raises:
             DatabaseError: If there is an error while rolling back the transaction.
         """
-        if self._is_closed():
-            # Cannot roll back if the connection is closed
-            raise Exception("Connection is closed. Cannot roll back.")
-
         # Roll back the current transaction
-        ret = ddbc_bindings.DDBCSQLEndTran(
-            ddbc_sql_const.SQL_HANDLE_DBC.value,  # Handle type
-            self.hdbc,  # Connection handle (wrapper)
-            ddbc_sql_const.SQL_ROLLBACK.value,  # Roll back the transaction
-        )
-        check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, self.hdbc, ret)
+        self._conn.rollback()
         if ENABLE_LOGGING:
             logger.info("Transaction rolled back successfully.")
 
@@ -409,16 +201,8 @@ def close(self) -> None:
         Raises:
             DatabaseError: If there is an error while closing the connection.
         """
-        if self._is_closed():
-            # Connection is already closed
-            return
-        # Disconnect from the database
-        ret = ddbc_bindings.DDBCSQLDisconnect(self.hdbc)
-        check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, self.hdbc, ret)
-
-        # Set the reference to None to trigger destructor
-        self.hdbc.free()
-        self.hdbc = None
-
+        # Close the connection
+        self._conn.close()
+        self._conn = None
         if ENABLE_LOGGING:
             logger.info("Connection closed successfully.")