FEAT: Access Token Login

jahnvi480 · jahnvi480 · commit fda52fff10a4 · 2025-05-13T16:23:02.000+05:30
diff --git a/mssql_python/connection.py b/mssql_python/connection.py
@@ -33,7 +33,7 @@ class Connection:
         close() -> None:
     """
 
-    def __init__(self, connection_str: str, autocommit: bool = False, **kwargs) -> None:
+    def __init__(self, connection_str: str, autocommit: bool = False, attrs_before: dict = {}, **kwargs) -> None:
         """
         Initialize the connection object with the specified connection string and parameters.
 
@@ -58,8 +58,9 @@ def __init__(self, connection_str: str, autocommit: bool = False, **kwargs) -> N
         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._autocommit = autocommit
         self.setautocommit(autocommit)
 
     def _construct_connection_string(self, connection_str: str, **kwargs) -> str:
@@ -76,23 +77,34 @@ def _construct_connection_string(self, connection_str: str, **kwargs) -> str:
         """
         # Add the driver attribute to the connection string
         conn_str = add_driver_to_connection_str(connection_str)
-        # Add additional key-value pairs to the connection string
-        for key, value in kwargs.items():
-            if key.lower() == "host":
-                key = "Server"
-            elif key.lower() == "user":
-                key = "Uid"
-            elif key.lower() == "password":
-                key = "Pwd"
-            elif key.lower() == "database":
-                key = "Database"
-            elif key.lower() == "encrypt":
-                key = "Encrypt"
-            elif key.lower() == "trust_server_certificate":
-                key = "TrustServerCertificate"
-            else:
-                continue
-            conn_str += f"{key}={value};"
+
+        # Check if access token authentication is being used
+        if "attrs_before" in kwargs:
+            # Skip adding Uid and Pwd for access token authentication
+            if ENABLE_LOGGING:
+                logger.info("Using access token authentication. Skipping Uid and Pwd.")
+        else:
+            # Add additional key-value pairs to the connection string
+            for key, value in kwargs.items():
+                if key.lower() == "host":
+                    key = "Server"
+                elif key.lower() == "user":
+                    key = "Uid"
+                elif key.lower() == "password":
+                    key = "Pwd"
+                elif key.lower() == "database":
+                    key = "Database"
+                elif key.lower() == "encrypt":
+                    key = "Encrypt"
+                elif key.lower() == "trust_server_certificate":
+                    key = "TrustServerCertificate"
+                else:
+                    continue
+                conn_str += f"{key}={value};"
+
+        if ENABLE_LOGGING:
+            logger.info("Final connection string: %s", conn_str)
+
         return conn_str
 
     def _is_closed(self) -> bool:
@@ -103,7 +115,7 @@ def _is_closed(self) -> bool:
             bool: True if the connection is closed, False otherwise.
         """
         return self.hdbc is None
-
+    
     def _initializer(self) -> None:
         """
         Initialize the environment and connection handles.
@@ -115,9 +127,41 @@ def _initializer(self) -> None:
         self._allocate_environment_handle()
         self._set_environment_attributes()
         self._allocate_connection_handle()
-        self._set_connection_attributes()
+        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 a dictionary of attributes to the database connection before connecting.
+
+        Returns:
+            bool: True if all attributes were successfully applied, False otherwise.
+        """
+        strencoding = "utf-16le"
+
+        if ENABLE_LOGGING:
+            logger.info("Applying attrs_before: %s", self._attrs_before)
+
+        for key, value in self._attrs_before.items():
+            if isinstance(key, int):
+                ikey = key
+            elif isinstance(key, str) and key.isdigit():
+                ikey = int(key)
+            else:
+                raise TypeError(f"Unsupported key type: {type(key).__name__}")
+
+            if not self._set_connection_attributes(ikey, value):
+                return False
+
+        return True
+
     def _allocate_environment_handle(self):
         """
         Allocate the environment handle.
@@ -152,18 +196,25 @@ def _allocate_connection_handle(self):
         check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, handle, ret)
         self.hdbc = handle
 
-    def _set_connection_attributes(self):
+    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.
         """
-        if self.autocommit:
-            ret = ddbc_bindings.DDBCSQLSetConnectAttr(
-                self.hdbc,  # Using the wrapper class
-                ddbc_sql_const.SQL_ATTR_AUTOCOMMIT.value,
-                ddbc_sql_const.SQL_AUTOCOMMIT_ON.value,
-                0
-            )
-            check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, self.hdbc, ret)
+
+        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:
         """
@@ -224,7 +275,6 @@ def autocommit(self, value: bool) -> None:
                 if value
                 else ddbc_sql_const.SQL_AUTOCOMMIT_OFF.value
             ),  # Value
-            0,  # String length
         )
         check_error(ddbc_sql_const.SQL_HANDLE_DBC.value, self.hdbc, ret)
         self._autocommit = value
diff --git a/mssql_python/constants.py b/mssql_python/constants.py
@@ -116,3 +116,5 @@ class ConstantsDDBC(Enum):
     SQL_C_WCHAR = -8
     SQL_NULLABLE = 1
     SQL_MAX_NUMERIC_LEN = 16
+    SQL_IS_POINTER = -4
+    
diff --git a/mssql_python/db_connection.py b/mssql_python/db_connection.py
@@ -6,7 +6,7 @@
 from mssql_python.connection import Connection
 
 
-def connect(connection_str: str, autocommit: bool = True, **kwargs) -> Connection:
+def connect(connection_str: str, autocommit: bool = True, attrs_before: dict = {}, **kwargs) -> Connection:
     """
     Constructor for creating a connection to the database.
 
@@ -34,5 +34,5 @@ def connect(connection_str: str, autocommit: bool = True, **kwargs) -> Connectio
     be used to perform database operations such as executing queries, committing
     transactions, and closing the connection.
     """
-    conn = Connection(connection_str, autocommit=autocommit, **kwargs)
+    conn = Connection(connection_str, autocommit=autocommit, attrs_before=attrs_before, **kwargs)
     return conn
diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp
@@ -692,18 +692,58 @@ SQLRETURN SQLSetEnvAttr_wrap(SqlHandlePtr EnvHandle, SQLINTEGER Attribute, intpt
 }
 
 // Wrap SQLSetConnectAttr
-SQLRETURN SQLSetConnectAttr_wrap(SqlHandlePtr ConnectionHandle, SQLINTEGER Attribute, intptr_t ValuePtr,
-                                 SQLINTEGER StringLength) {
+SQLRETURN SQLSetConnectAttr_wrap(SqlHandlePtr ConnectionHandle, SQLINTEGER Attribute, 
+                                 py::object ValuePtr) {
     LOG("Set SQL Connection Attribute");
     if (!SQLSetConnectAttr_ptr) {
         LoadDriverOrThrowException();
     }
 
-    // TODO: Does ValuePtr need to be converted from Python to C++ object?
-    SQLRETURN ret =  SQLSetConnectAttr_ptr(ConnectionHandle->get(), Attribute, reinterpret_cast<SQLPOINTER>(ValuePtr), StringLength);
+    // Print the type of ValuePtr and attribute value - helpful for debugging
+    LOG("Type of ValuePtr: {}, Attribute: {}", py::type::of(ValuePtr).attr("__name__").cast<std::string>(), Attribute);
+
+    SQLPOINTER value = 0;
+    SQLINTEGER length = 0;
+
+    if (py::isinstance<py::int_>(ValuePtr)) {
+        // Handle integer values
+        int intValue = ValuePtr.cast<int>();
+        value = reinterpret_cast<SQLPOINTER>(intValue);
+        length = SQL_IS_INTEGER;  // Integer values don't require a length
+    } else if (py::isinstance<py::str>(ValuePtr)) {
+        // Handle Unicode string values
+        static std::wstring unicodeValueBuffer;
+        unicodeValueBuffer = ValuePtr.cast<std::wstring>();
+        value = const_cast<SQLWCHAR*>(unicodeValueBuffer.c_str());
+        length = SQL_NTS;  // Indicates null-terminated string
+    } else if (py::isinstance<py::bytes>(ValuePtr) || py::isinstance<py::bytearray>(ValuePtr)) {
+        // Handle byte or bytearray values (like access tokens)
+        // Store in static buffer to ensure memory remains valid during connection
+        static std::vector<std::string> bytesBuffers;
+        bytesBuffers.push_back(ValuePtr.cast<std::string>());
+        value = const_cast<char*>(bytesBuffers.back().c_str());
+        length = SQL_IS_POINTER;  // Indicates we're passing a pointer (required for token)
+    } else if (py::isinstance<py::list>(ValuePtr) || py::isinstance<py::tuple>(ValuePtr)) {
+        // Handle list or tuple values
+        LOG("ValuePtr is a sequence (list or tuple)");
+        for (py::handle item : ValuePtr) {
+            LOG("Processing item in sequence");
+            SQLRETURN ret = SQLSetConnectAttr_wrap(ConnectionHandle, Attribute, py::reinterpret_borrow<py::object>(item));
+            if (!SQL_SUCCEEDED(ret)) {
+                LOG("Failed to set attribute for item in sequence");
+                return ret;
+            }
+        }    
+    } else {
+        LOG("Unsupported ValuePtr type");
+        return SQL_ERROR;
+    }
+
+    SQLRETURN ret = SQLSetConnectAttr_ptr(ConnectionHandle->get(), Attribute, value, length);
     if (!SQL_SUCCEEDED(ret)) {
         LOG("Failed to set Connection attribute");
     }
+    LOG("Set Connection attribute successfully");
     return ret;
 }
 
