working interation with access token

Saumya Garg · gargsaumya · commit 65d19a8af4ce · 2025-05-21T20:43:48.000+05:30
diff --git a/mssql_python/connection.py b/mssql_python/connection.py
@@ -56,16 +56,9 @@ def __init__(self, connection_str: str = "", autocommit: bool = False, attrs_bef
         self.connection_str = self._construct_connection_string(
             connection_str, **kwargs
         )
-        # self._attrs_before = attrs_before
-        # self._autocommit = autocommit
-        # if self._attrs_before != {}:
-        #     self._apply_attrs_before()  # Apply pre-connection attributes
-        # if self._autocommit:
-        #     self.setautocommit(autocommit)
-        print("Connection string:", self.connection_str)
-        self._conn = ddbc_bindings.Connection(self.connection_str)
-        self._conn.connect()
-        print("Connection established")
+        self._attrs_before = attrs_before or {}
+        self._conn = ddbc_bindings.Connection(self.connection_str, autocommit)
+        self._conn.connect(self._attrs_before)
         self._autocommit = autocommit
         self.setautocommit(autocommit)
 
@@ -107,90 +100,6 @@ def _construct_connection_string(self, connection_str: str = "", **kwargs) -> st
 
         return conn_str
     
-    # 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._conn.set_attribute(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 _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)
-
     @property
     def autocommit(self) -> bool:
         """
@@ -208,8 +117,6 @@ 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.
         """
         self.setautocommit(value)
         if ENABLE_LOGGING:
@@ -243,9 +150,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:
@@ -260,10 +164,6 @@ 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
         self._conn.commit()
         if ENABLE_LOGGING:
@@ -280,10 +180,6 @@ 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
         self._conn.rollback()
         if ENABLE_LOGGING:
@@ -302,10 +198,7 @@ def close(self) -> None:
         Raises:
             DatabaseError: If there is an error while closing the connection.
         """
-        if self._is_closed():
-            # Connection is already closed
-            return
+        # Close the connection
         self._conn.close()
-
         if ENABLE_LOGGING:
             logger.info("Connection closed successfully.")
diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py
@@ -415,8 +415,7 @@ def _allocate_statement_handle(self):
         """
         Allocate the DDBC statement handle.
         """
-        self.connection._conn.alloc_statement_handle()
-        print(f"Statement handle: {self.hstmt}")
+        self.hstmt = self.connection._conn.alloc_statement_handle()
 
     def _reset_cursor(self) -> None:
         """
diff --git a/mssql_python/pybind/connection/connection.cpp b/mssql_python/pybind/connection/connection.cpp
@@ -7,24 +7,15 @@
 #include "connection.h"
 #include <iostream>
 #include <vector>
-
-<<<<<<< HEAD
 #include <pybind11/pybind11.h>
 
 //-------------------------------------------------------------------------------------------------
 // Implements the Connection class declared in connection.h.
 // This class wraps low-level ODBC operations like connect/disconnect,
 // transaction control, and autocommit configuration.
 //-------------------------------------------------------------------------------------------------
-<<<<<<< HEAD
 Connection::Connection(const std::wstring& conn_str, bool autocommit)
     : _conn_str(conn_str) , _autocommit(autocommit) {}
-=======
-Connection::Connection(const std::wstring& conn_str, bool autocommit) : _conn_str(conn_str) , _autocommit(autocommit) {}
-=======
-Connection::Connection(const std::wstring& conn_str) : _conn_str(conn_str) {}
->>>>>>> fcd64d4 (working flow with c++ connection class)
->>>>>>> 29a7a65 (working flow with c++ connection class)
 
 Connection::~Connection() {
     close();    // Ensure the connection is closed when the object is destroyed.
@@ -175,6 +166,52 @@ SqlHandlePtr Connection::allocStatementHandle() {
     return std::make_shared<SqlHandle>(SQL_HANDLE_STMT, stmt);
 }
 
+SQLRETURN Connection::set_attribute(SQLINTEGER attribute, py::object value) {
+    LOG("Setting SQL attribute {}");
+
+    SQLPOINTER ptr = nullptr;
+    SQLINTEGER length = 0;
+
+    if (py::isinstance<py::int_>(value)) {
+        int intValue = value.cast<int>();
+        ptr = reinterpret_cast<SQLPOINTER>(static_cast<uintptr_t>(intValue));
+        length = SQL_IS_INTEGER;
+    } else if (py::isinstance<py::bytes>(value) || py::isinstance<py::bytearray>(value)) {
+        static std::vector<std::string> buffers;
+        buffers.emplace_back(value.cast<std::string>());
+        ptr = const_cast<char*>(buffers.back().c_str());
+        length = static_cast<SQLINTEGER>(buffers.back().size());
+    } else {
+        LOG("Unsupported attribute value type");
+        return SQL_ERROR;
+    }
+
+    SQLRETURN ret = SQLSetConnectAttr_ptr(_dbc_handle->get(), attribute, ptr, length);
+    if (!SQL_SUCCEEDED(ret)) {
+        LOG("Failed to set attribute {}");
+    }
+    return ret;
+}
+
+void Connection::apply_attrs_before(const py::dict& attrs) {
+    for (const auto& item : attrs) {
+        int key;
+        try {
+            key = py::cast<int>(item.first);
+        } catch (...) {
+            continue;
+        }
+
+        //do not hard code the key values
+        if (key == 1256) {   
+            SQLRETURN ret = set_attribute(key, py::reinterpret_borrow<py::object>(item.second));
+            if (!SQL_SUCCEEDED(ret)) {
+                throw std::runtime_error("Failed to set access token before connect");
+            }
+        }
+    }
+}
+
 SqlHandlePtr Connection::getSharedEnvHandle() {
     static std::once_flag flag;
     static SqlHandlePtr env_handle;
diff --git a/mssql_python/pybind/connection/connection.h b/mssql_python/pybind/connection/connection.h
@@ -18,8 +18,7 @@ class Connection {
     Connection(const std::wstring& conn_str, bool autocommit = false);
     ~Connection();
 
-    // Establish the connection using the stored connection string.
-    SQLRETURN connect();
+    SQLRETURN connect(const py::dict& attrs_before = py::dict());
 
     // Close the connection and free resources.
     SQLRETURN close();
@@ -46,8 +45,10 @@ class Connection {
     std::wstring _conn_str;
     SqlHandlePtr _dbc_handle;
     bool _autocommit = false;
+    std::shared_ptr<Connection> _conn; 
 
-    static SqlHandlePtr getSharedEnvHandle();
+    SQLRETURN set_attribute(SQLINTEGER attribute, pybind11::object value);
+    void apply_attrs_before(const pybind11::dict& attrs);
 };
 
 #endif // CONNECTION_H
diff --git a/mssql_python/pybind/ddbc_bindings.cpp b/mssql_python/pybind/ddbc_bindings.cpp
diff --git a/mssql_python/pybind/ddbc_bindings.h b/mssql_python/pybind/ddbc_bindings.h
