test(e2e): authenticate e2e via OAuth M2M so staging tests match DATABRICKS_USER

The e2e suite connected via a PAT (DATABRICKS_TOKEN). The Personal Staging
Location tests PUT/GET/REMOVE against stage://tmp/<DATABRICKS_USER>/..., where
DATABRICKS_USER is the PECO service principal (TEST_PECO_SP_ID). A personal
stage is identity-scoped by design (there is even a test asserting you cannot
touch another user's stage), so the connecting identity MUST equal
DATABRICKS_USER. When DATABRICKS_TOKEN authenticates as a different identity,
those tests fail with `PERMISSION_DENIED: <user> does not have access to
Personal Stage`.

Switch the e2e connection to OAuth M2M as the service principal via
credentials_provider (conftest.auth_connect_kwargs), so the connecting identity
IS the SP == DATABRICKS_USER. Falls back to the PAT when SP OAuth creds aren't
set, so local PAT runs are unaffected. Wires DATABRICKS_CLIENT_ID /
DATABRICKS_CLIENT_SECRET (TEST_PECO_SP_ID / TEST_PECO_SP_OAUTH_SECRET, already
in azure-prod) into code-coverage.yml.

Verified locally against the PECO workspace: all 9 staging_ingestion e2e tests
pass via the real M2M path (including fails_to_modify_another_staging_user,
which validates the identity scoping). Kernel e2e files are unchanged (they run
in kernel-e2e.yml, ignored by code-coverage.yml).

Co-authored-by: Isaac
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

.github/workflows/code-coverage.yml

@@ -45,6 +45,14 @@ jobs:
       DATABRICKS_SERVER_HOSTNAME: ${{ secrets.DATABRICKS_HOST }}
       DATABRICKS_HTTP_PATH: ${{ secrets.TEST_PECO_WAREHOUSE_HTTP_PATH }}
       DATABRICKS_TOKEN: ${{ secrets.DATABRICKS_TOKEN }}
+      # OAuth M2M as the PECO service principal. The e2e suite prefers this
+      # over the PAT (see conftest.auth_connect_kwargs) so the connecting
+      # identity IS the service principal -- required for the Personal Staging
+      # Location tests, whose stage path is stage://tmp/<DATABRICKS_USER>/ and
+      # DATABRICKS_USER is this same SP. A PAT for a different identity fails
+      # those with PERMISSION_DENIED on the personal stage.
+      DATABRICKS_CLIENT_ID: ${{ secrets.TEST_PECO_SP_ID }}
+      DATABRICKS_CLIENT_SECRET: ${{ secrets.TEST_PECO_SP_OAUTH_SECRET }}
       DATABRICKS_CATALOG: peco
       DATABRICKS_USER: ${{ secrets.TEST_PECO_SP_ID }}
     steps:

conftest.py

@@ -17,6 +17,17 @@ def access_token():
     return os.getenv("DATABRICKS_TOKEN")
 
 
+@pytest.fixture(scope="session")
+def client_id():
+    # OAuth M2M service-principal client id (application id).
+    return os.getenv("DATABRICKS_CLIENT_ID")
+
+
+@pytest.fixture(scope="session")
+def client_secret():
+    return os.getenv("DATABRICKS_CLIENT_SECRET")
+
+
 @pytest.fixture(scope="session")
 def ingestion_user():
     return os.getenv("DATABRICKS_USER")
@@ -33,12 +44,58 @@ def schema():
 
 
 @pytest.fixture(scope="session", autouse=True)
-def connection_details(host, http_path, access_token, ingestion_user, catalog, schema):
+def connection_details(
+    host,
+    http_path,
+    access_token,
+    client_id,
+    client_secret,
+    ingestion_user,
+    catalog,
+    schema,
+):
     return {
         "host": host,
         "http_path": http_path,
         "access_token": access_token,
+        "client_id": client_id,
+        "client_secret": client_secret,
         "ingestion_user": ingestion_user,
         "catalog": catalog,
         "schema": schema,
     }
+
+
+def auth_connect_kwargs(details):
+    """Return the sql.connect auth kwargs from connection_details.
+
+    Prefers OAuth M2M (service principal) when DATABRICKS_CLIENT_ID /
+    DATABRICKS_CLIENT_SECRET are set, otherwise falls back to a PAT
+    (DATABRICKS_TOKEN). M2M is required for identity-scoped operations such as
+    Personal Staging Location tests (stage://tmp/<DATABRICKS_USER>/...), where
+    the connecting identity must equal DATABRICKS_USER -- the service principal.
+    A PAT belonging to a different identity cannot access that stage.
+    """
+    client_id = details.get("client_id")
+    client_secret = details.get("client_secret")
+    if client_id and client_secret:
+        host = details["host"]
+        host_url = host if host.startswith("http") else f"https://{host}"
+
+        def credential_provider():
+            # Imported lazily so a PAT-only environment doesn't require the SDK.
+            from databricks.sdk.core import Config, oauth_service_principal
+
+            return oauth_service_principal(
+                Config(
+                    host=host_url,
+                    client_id=client_id,
+                    client_secret=client_secret,
+                    # Explicit so an ambient DATABRICKS_TOKEN doesn't collide
+                    # ("more than one authorization method configured").
+                    auth_type="oauth-m2m",
+                )
+            )
+
+        return {"credentials_provider": credential_provider}
+    return {"access_token": details.get("access_token")}

tests/e2e/test_circuit_breaker.py

@@ -21,6 +21,7 @@
 
 import databricks.sql as sql
 from databricks.sql.telemetry.circuit_breaker_manager import CircuitBreakerManager
+from conftest import auth_connect_kwargs
 
 
 def wait_for_circuit_state(circuit_breaker, expected_states, timeout=5):
@@ -120,7 +121,7 @@ def mock_request(*args, **kwargs):
             with sql.connect(
                 server_hostname=self.arguments["host"],
                 http_path=self.arguments["http_path"],
-                access_token=self.arguments.get("access_token"),
+                **auth_connect_kwargs(self.arguments),
                 force_enable_telemetry=True,
                 telemetry_batch_size=1,
                 _telemetry_circuit_breaker_enabled=True,
@@ -181,7 +182,7 @@ def mock_rate_limited_request(*args, **kwargs):
             with sql.connect(
                 server_hostname=self.arguments["host"],
                 http_path=self.arguments["http_path"],
-                access_token=self.arguments.get("access_token"),
+                **auth_connect_kwargs(self.arguments),
                 force_enable_telemetry=True,
                 telemetry_batch_size=1,
                 _telemetry_circuit_breaker_enabled=False,  # Disabled
@@ -215,7 +216,7 @@ def mock_conditional_request(*args, **kwargs):
             with sql.connect(
                 server_hostname=self.arguments["host"],
                 http_path=self.arguments["http_path"],
-                access_token=self.arguments.get("access_token"),
+                **auth_connect_kwargs(self.arguments),
                 force_enable_telemetry=True,
                 telemetry_batch_size=1,
                 _telemetry_circuit_breaker_enabled=True,

tests/e2e/test_driver.py

@@ -99,9 +99,9 @@ def connection_params(self):
         return params
 
     def auth_params(self):
-        return {
-            "access_token": self.arguments.get("access_token"),
-        }
+        from conftest import auth_connect_kwargs
+
+        return auth_connect_kwargs(self.arguments)
 
     @contextmanager
     def connection(self, extra_params=()):

tests/e2e/test_telemetry_e2e.py

@@ -24,10 +24,12 @@ def get_details(self, connection_details):
         self.arguments = connection_details.copy()
 
     def connection_params(self):
+        from conftest import auth_connect_kwargs
+
         return {
             "server_hostname": self.arguments["host"],
             "http_path": self.arguments["http_path"],
-            "access_token": self.arguments.get("access_token"),
+            **auth_connect_kwargs(self.arguments),
         }
 
     @contextmanager

tests/e2e/test_transactions.py

@@ -56,10 +56,12 @@ def _unique_table_name_raw(suffix):
 @pytest.fixture
 def mst_conn_params(connection_details):
     """Connection parameters with MST enabled."""
+    from conftest import auth_connect_kwargs
+
     return {
         "server_hostname": connection_details["host"],
         "http_path": connection_details["http_path"],
-        "access_token": connection_details.get("access_token"),
+        **auth_connect_kwargs(connection_details),
         "ignore_transactions": False,
     }