Merge branch 'main' into saumya/stress-tests-gh

Author: gargsaumya

OneBranchPipelines/build-release-package-pipeline.yml

@@ -201,8 +201,8 @@ extends:
   parameters:
     # Pool Configuration
     # Different platforms use different agent pools:
-    # - Windows: Custom 1ES pool (Django-1ES-pool) with WIN22-SQL22 image (Windows Server 2022 + SQL Server 2022)
-    # - Linux: Custom 1ES pool (Django-1ES-pool) with ADO-UB22-SQL22 image (Ubuntu 22.04 + SQL Server 2022)
+    # - Windows: Custom 1ES pool (Python-1ES-pool) with PYTHON-1ES-MMS2022 image (Windows Server 2022 + SQL Server 2022)
+    # - Linux: Custom 1ES pool (Python-1ES-pool) with PYTHON-1ES-UB2404 image (Ubuntu 24.04 + SQL Server 2022)
     # - macOS: Microsoft-hosted pool (Azure Pipelines) with macOS-14 image (macOS Sonoma)
     # Note: Container definitions section present but unused (pools configured in individual stage templates)
 

OneBranchPipelines/dummy-release-pipeline.yml

@@ -96,9 +96,9 @@ extends:
             pool:
               type: windows
               isCustom: true
-              name: Django-1ES-pool
+              name: Python-1ES-pool
               demands:
-              - imageOverride -equals WIN22-SQL22
+              - imageOverride -equals PYTHON-1ES-MMS2022
 
             variables:
               ob_outputDirectory: '$(Build.ArtifactStagingDirectory)'

OneBranchPipelines/official-release-pipeline.yml

@@ -99,9 +99,9 @@ extends:
             pool:
               type: windows
               isCustom: true
-              name: Django-1ES-pool
+              name: Python-1ES-pool
               demands:
-              - imageOverride -equals WIN22-SQL22
+              - imageOverride -equals PYTHON-1ES-MMS2022
 
             variables:
               ob_outputDirectory: '$(Build.ArtifactStagingDirectory)'

OneBranchPipelines/stages/build-linux-single-stage.yml

@@ -35,9 +35,9 @@ stages:
         pool:
           type: linux
           isCustom: true
-          name: Django-1ES-pool
+          name: Python-1ES-pool
           demands:
-          - imageOverride -equals ADO-UB22-SQL22
+          - imageOverride -equals PYTHON-1ES-UB2404
         # Extended timeout for multi-version builds + testing (5 Python versions × build + test time)
         timeoutInMinutes: 120
 
@@ -65,12 +65,8 @@ stages:
           - checkout: self
             fetchDepth: 0
 
-          # Install Docker
-          - task: DockerInstaller@0
-            inputs:
-              dockerVersion: '20.10.21'
-            displayName: 'Install Docker'
-          
+          # Docker is pre-installed now
+          # Verify it's working and start the daemon
           - bash: |
               set -e
               echo "Verifying we're on Linux..."
@@ -82,9 +78,14 @@ stages:
               
               uname -a
               
-              # Start dockerd
-              sudo dockerd > docker.log 2>&1 &
-              sleep 10
+              # Check if Docker daemon is already running
+              if ! docker info > /dev/null 2>&1; then
+                echo "Docker daemon not running, starting it..."
+                sudo dockerd > docker.log 2>&1 &
+                sleep 10
+              else
+                echo "Docker daemon already running"
+              fi
               
               # Verify Docker works
               docker --version

OneBranchPipelines/stages/build-windows-single-stage.yml

@@ -38,9 +38,9 @@ stages:
         pool:
           type: windows
           isCustom: true
-          name: Django-1ES-pool
+          name: Python-1ES-pool
           demands:
-          - imageOverride -equals WIN22-SQL22
+          - imageOverride -equals PYTHON-1ES-MMS2022
         # Extended timeout for downloads, builds, and testing
         timeoutInMinutes: 120
 

mssql_python/auth.py

@@ -9,7 +9,7 @@
 from typing import Tuple, Dict, Optional, List
 
 from mssql_python.logging import logger
-from mssql_python.constants import AuthType
+from mssql_python.constants import AuthType, ConstantsDDBC
 
 
 class AADAuth:
@@ -30,7 +30,25 @@ def get_token_struct(token: str) -> bytes:
 
     @staticmethod
     def get_token(auth_type: str) -> bytes:
-        """Get token using the specified authentication type"""
+        """Get DDBC token struct for the specified authentication type."""
+        token_struct, _ = AADAuth._acquire_token(auth_type)
+        return token_struct
+
+    @staticmethod
+    def get_raw_token(auth_type: str) -> str:
+        """Acquire a fresh raw JWT for the mssql-py-core connection (bulk copy).
+
+        This deliberately does NOT cache the credential or token — each call
+        creates a new Azure Identity credential instance and requests a token.
+        A fresh acquisition avoids expired-token errors when bulkcopy() is
+        called long after the original DDBC connect().
+        """
+        _, raw_token = AADAuth._acquire_token(auth_type)
+        return raw_token
+
+    @staticmethod
+    def _acquire_token(auth_type: str) -> Tuple[bytes, str]:
+        """Internal: acquire token and return (ddbc_struct, raw_jwt)."""
         # Import Azure libraries inside method to support test mocking
         # pylint: disable=import-outside-toplevel
         try:
@@ -53,30 +71,27 @@ def get_token(auth_type: str) -> bytes:
             "interactive": InteractiveBrowserCredential,
         }
 
-        credential_class = credential_map[auth_type]
+        credential_class = credential_map.get(auth_type)
+        if not credential_class:
+            raise ValueError(
+                f"Unsupported auth_type '{auth_type}'. " f"Supported: {', '.join(credential_map)}"
+            )
         logger.info(
             "get_token: Starting Azure AD authentication - auth_type=%s, credential_class=%s",
             auth_type,
             credential_class.__name__,
         )
 
         try:
-            logger.debug(
-                "get_token: Creating credential instance - credential_class=%s",
-                credential_class.__name__,
-            )
             credential = credential_class()
-            logger.debug(
-                "get_token: Requesting token from Azure AD - scope=https://database.windows.net/.default"
-            )
-            token = credential.get_token("https://database.windows.net/.default").token
+            raw_token = credential.get_token("https://database.windows.net/.default").token
             logger.info(
                 "get_token: Azure AD token acquired successfully - token_length=%d chars",
-                len(token),
+                len(raw_token),
             )
-            return AADAuth.get_token_struct(token)
+            token_struct = AADAuth.get_token_struct(raw_token)
+            return token_struct, raw_token
         except ClientAuthenticationError as e:
-            # Re-raise with more specific context about Azure AD authentication failure
             logger.error(
                 "get_token: Azure AD authentication failed - credential_class=%s, error=%s",
                 credential_class.__name__,
@@ -88,7 +103,6 @@ def get_token(auth_type: str) -> bytes:
                 f"user cancellation, network issues, or unsupported configuration."
             ) from e
         except Exception as e:
-            # Catch any other unexpected exceptions
             logger.error(
                 "get_token: Unexpected error during credential creation - credential_class=%s, error=%s",
                 credential_class.__name__,
@@ -180,7 +194,7 @@ def remove_sensitive_params(parameters: List[str]) -> List[str]:
 
 
 def get_auth_token(auth_type: str) -> Optional[bytes]:
-    """Get authentication token based on auth type"""
+    """Get DDBC authentication token struct based on auth type."""
     logger.debug("get_auth_token: Starting - auth_type=%s", auth_type)
     if not auth_type:
         logger.debug("get_auth_token: No auth_type specified, returning None")
@@ -202,17 +216,37 @@ def get_auth_token(auth_type: str) -> Optional[bytes]:
         return None
 
 
+def extract_auth_type(connection_string: str) -> Optional[str]:
+    """Extract Entra ID auth type from a connection string.
+
+    Used as a fallback when process_connection_string does not propagate
+    auth_type (e.g. Windows Interactive where DDBC handles auth natively).
+    Bulkcopy still needs the auth type to acquire a token via Azure Identity.
+    """
+    auth_map = {
+        AuthType.INTERACTIVE.value: "interactive",
+        AuthType.DEVICE_CODE.value: "devicecode",
+        AuthType.DEFAULT.value: "default",
+    }
+    for part in connection_string.split(";"):
+        key, _, value = part.strip().partition("=")
+        if key.strip().lower() == "authentication":
+            return auth_map.get(value.strip().lower())
+    return None
+
+
 def process_connection_string(
     connection_string: str,
-) -> Tuple[str, Optional[Dict[int, bytes]]]:
+) -> Tuple[str, Optional[Dict[int, bytes]], Optional[str]]:
     """
     Process connection string and handle authentication.
 
     Args:
         connection_string: The connection string to process
 
     Returns:
-        Tuple[str, Optional[Dict]]: Processed connection string and attrs_before dict if needed
+        Tuple[str, Optional[Dict], Optional[str]]: Processed connection string,
+            attrs_before dict if needed, and auth_type string for bulk copy token acquisition
 
     Raises:
         ValueError: If the connection string is invalid or empty
@@ -259,7 +293,11 @@ def process_connection_string(
                 "process_connection_string: Token authentication configured successfully - auth_type=%s",
                 auth_type,
             )
-            return ";".join(modified_parameters) + ";", {1256: token_struct}
+            return (
+                ";".join(modified_parameters) + ";",
+                {ConstantsDDBC.SQL_COPT_SS_ACCESS_TOKEN.value: token_struct},
+                auth_type,
+            )
         else:
             logger.warning(
                 "process_connection_string: Token acquisition failed, proceeding without token"
@@ -269,4 +307,4 @@ def process_connection_string(
         "process_connection_string: Connection string processing complete - has_auth=%s",
         bool(auth_type),
     )
-    return ";".join(modified_parameters) + ";", None
+    return ";".join(modified_parameters) + ";", None, auth_type

mssql_python/connection.py

@@ -39,7 +39,7 @@
     ProgrammingError,
     NotSupportedError,
 )
-from mssql_python.auth import process_connection_string
+from mssql_python.auth import extract_auth_type, process_connection_string
 from mssql_python.constants import ConstantsDDBC, GetInfoConstants
 from mssql_python.connection_string_parser import _ConnectionStringParser
 from mssql_python.connection_string_builder import _ConnectionStringBuilder
@@ -263,6 +263,11 @@ def __init__(
             },
         }
 
+        # Auth type for acquiring fresh tokens at bulk copy time.
+        # We intentionally do NOT cache the token — a fresh one is acquired
+        # each time bulkcopy() is called to avoid expired-token errors.
+        self._auth_type = None
+
         # Check if the connection string contains authentication parameters
         # This is important for processing the connection string correctly.
         # If authentication is specified, it will be processed to handle
@@ -272,6 +277,10 @@ def __init__(
             self.connection_str = connection_result[0]
             if connection_result[1]:
                 self._attrs_before.update(connection_result[1])
+            # Store auth type so bulkcopy() can acquire a fresh token later.
+            # On Windows Interactive, process_connection_string returns None
+            # (DDBC handles auth natively), so fall back to the connection string.
+            self._auth_type = connection_result[2] or extract_auth_type(self.connection_str)
 
         self._closed = False
         self._timeout = timeout

mssql_python/constants.py

@@ -158,6 +158,9 @@ class ConstantsDDBC(Enum):
     SQL_ATTR_SERVER_NAME = 13
     SQL_ATTR_RESET_CONNECTION = 116
 
+    # SQL Server-specific connection option constants
+    SQL_COPT_SS_ACCESS_TOKEN = 1256
+
     # Transaction Isolation Level Constants
     SQL_TXN_READ_UNCOMMITTED = 1
     SQL_TXN_READ_COMMITTED = 2

mssql_python/cursor.py

@@ -2607,15 +2607,36 @@ def _bulkcopy(
         context = {
             "server": params.get("server"),
             "database": params.get("database"),
-            "user_name": params.get("uid", ""),
             "trust_server_certificate": trust_cert,
             "encryption": encryption,
         }
 
-        # Extract password separately to avoid storing it in generic context that may be logged
-        password = params.get("pwd", "")
+        # Build pycore_context with appropriate authentication.
+        # For Azure AD: acquire a FRESH token right now instead of reusing
+        # the one from connect() time — avoids expired-token errors when
+        # bulkcopy() is called long after the original connection.
         pycore_context = dict(context)
-        pycore_context["password"] = password
+
+        if self.connection._auth_type:
+            # Fresh token acquisition for mssql-py-core connection
+            from mssql_python.auth import AADAuth
+
+            try:
+                raw_token = AADAuth.get_raw_token(self.connection._auth_type)
+            except (RuntimeError, ValueError) as e:
+                raise RuntimeError(
+                    f"Bulk copy failed: unable to acquire Azure AD token "
+                    f"for auth_type '{self.connection._auth_type}': {e}"
+                ) from e
+            pycore_context["access_token"] = raw_token
+            logger.debug(
+                "Bulk copy: acquired fresh Azure AD token for auth_type=%s",
+                self.connection._auth_type,
+            )
+        else:
+            # SQL Server authentication — use uid/password from connection string
+            pycore_context["user_name"] = params.get("uid", "")
+            pycore_context["password"] = params.get("pwd", "")
 
         pycore_connection = None
         pycore_cursor = None
@@ -2653,10 +2674,10 @@ def _bulkcopy(
 
         finally:
             # Clear sensitive data to minimize memory exposure
-            password = ""
             if pycore_context:
-                pycore_context["password"] = ""
-                pycore_context["user_name"] = ""
+                pycore_context.pop("password", None)
+                pycore_context.pop("user_name", None)
+                pycore_context.pop("access_token", None)
             # Clean up bulk copy resources
             for resource in (pycore_cursor, pycore_connection):
                 if resource and hasattr(resource, "close"):