Fix ContextVar propagation to worker threads, revert _MAX_WORKERS to 3, fix minor docstrings

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Abel Milash

examples/advanced/contextvar_thread_demo.py

@@ -0,0 +1,98 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Demonstrates that ContextVar values are NOT inherited by ThreadPoolExecutor
+worker threads — confirming the correlation ID bug in _dispatch_chunks.
+
+Run from repo root:
+    .conda/python.exe examples/advanced/contextvar_thread_demo.py
+"""
+
+import threading
+from contextvars import ContextVar, copy_context
+from concurrent.futures import ThreadPoolExecutor
+
+# Mirrors the SDK's _CALL_SCOPE_CORRELATION_ID exactly
+CORRELATION_ID: ContextVar[str | None] = ContextVar("CORRELATION_ID", default=None)
+
+def read_correlation_id(label: str) -> str:
+    """Read the ContextVar — simulates what _RequestContext.from_request() does."""
+    value = CORRELATION_ID.get()
+    print(f"  [{label}] thread={threading.current_thread().name:20s}  "
+          f"correlation_id = {value!r}")
+    return value
+
+
+# ---------------------------------------------------------------------------
+# Part 1: WITHOUT fix — plain ThreadPoolExecutor.submit()
+# ---------------------------------------------------------------------------
+print("=" * 60)
+print("PART 1: Plain submit() — no context propagation (current SDK)")
+print("=" * 60)
+
+CORRELATION_ID.set("abc-123-shared-id")
+print(f"\nMain thread sets correlation_id = 'abc-123-shared-id'")
+print(f"Dispatching 3 chunks to worker threads...\n")
+
+with ThreadPoolExecutor(max_workers=3) as pool:
+    futures = [pool.submit(read_correlation_id, f"chunk-{i}") for i in range(3)]
+    results_before = [f.result() for f in futures]
+
+print(f"\nMain thread still sees: {CORRELATION_ID.get()!r}")
+print(f"Worker results: {results_before}")
+print(f"\n=> All workers got None — correlation ID is LOST in concurrent path.\n")
+
+
+# ---------------------------------------------------------------------------
+# Part 2: WITH fix — copy_context().run()
+# ---------------------------------------------------------------------------
+print("=" * 60)
+print("PART 2: copy_context() — correct propagation (proposed fix)")
+print("=" * 60)
+
+CORRELATION_ID.set("abc-123-shared-id")
+print(f"\nMain thread sets correlation_id = 'abc-123-shared-id'")
+print(f"Dispatching 3 chunks with ctx.run()...\n")
+
+ctx = copy_context()  # snapshot the main thread's context
+with ThreadPoolExecutor(max_workers=3) as pool:
+    futures = [pool.submit(ctx.run, read_correlation_id, f"chunk-{i}") for i in range(3)]
+    results_after = [f.result() for f in futures]
+
+print(f"\nMain thread still sees: {CORRELATION_ID.get()!r}")
+print(f"Worker results: {results_after}")
+print(f"\n=> All workers got 'abc-123-shared-id' — correlation ID is preserved.\n")
+
+
+# ---------------------------------------------------------------------------
+# Part 3: Test the actual SDK _dispatch_chunks with the fix applied
+# ---------------------------------------------------------------------------
+print("=" * 60)
+print("PART 3: Real SDK _dispatch_chunks (with fix applied)")
+print("=" * 60)
+
+from PowerPlatform.Dataverse.data._odata import (
+    _dispatch_chunks,
+    _CALL_SCOPE_CORRELATION_ID,
+)
+
+def simulate_chunk_request(chunk):
+    """Reads the SDK's real ContextVar — same as _RequestContext.from_request()."""
+    corr_id = _CALL_SCOPE_CORRELATION_ID.get()
+    print(f"  chunk={chunk}  x-ms-correlation-id = {corr_id!r}  "
+          f"(thread={threading.current_thread().name})")
+    return corr_id
+
+_CALL_SCOPE_CORRELATION_ID.set("real-sdk-call-uuid-xyz")
+print(f"\nSDK: _call_scope sets correlation_id = 'real-sdk-call-uuid-xyz'")
+print(f"SDK: _dispatch_chunks dispatches 3 chunks with max_workers=3\n")
+
+chunks = ["chunk-A", "chunk-B", "chunk-C"]
+results = _dispatch_chunks(simulate_chunk_request, chunks, max_workers=3)
+
+print(f"\nResults: {results}")
+if all(r == "real-sdk-call-uuid-xyz" for r in results):
+    print("=> [PASS] All chunks received the correct correlation ID.")
+else:
+    print("=> [FAIL] Some chunks got None — fix not working.")

src/PowerPlatform/Dataverse/data/_odata.py

@@ -20,7 +20,7 @@
 from datetime import datetime, timezone
 import importlib.resources as ir
 from contextlib import contextmanager
-from contextvars import ContextVar
+from contextvars import ContextVar, copy_context
 
 from urllib.parse import quote as _url_quote, parse_qs, urlparse
 
@@ -96,7 +96,8 @@ def _execute_with_retry(chunk):
                 return fn(chunk)
             except HttpError as exc:
                 if exc.is_transient and attempt < _CHUNK_RETRY_LIMIT:
-                    wait = float(exc.details.get("retry_after") or _CHUNK_RETRY_DEFAULT_WAIT)
+                    ra = exc.details.get("retry_after")
+                    wait = float(_CHUNK_RETRY_DEFAULT_WAIT if ra is None else ra)
                     wait += random.uniform(0, _CHUNK_RETRY_JITTER_MAX)
                     time.sleep(wait)
                 else:
@@ -106,7 +107,7 @@ def _execute_with_retry(chunk):
         return [_execute_with_retry(chunk) for chunk in chunks]
 
     with ThreadPoolExecutor(max_workers=max_workers) as pool:
-        futures = [pool.submit(_execute_with_retry, chunk) for chunk in chunks]
+        futures = [pool.submit(copy_context().run, _execute_with_retry, chunk) for chunk in chunks]
         return [f.result() for f in futures]
 
 
@@ -584,11 +585,18 @@ def _upsert_multiple(
            When input exceeds ``_MULTIPLE_BATCH_SIZE`` records, the operation is
            split into multiple requests and is **not atomic** across batches.
         """
+        # Validation uses ValueError (not ValidationError) because this is a
+        # caller-facing precondition check, not a service error.  The batch path
+        # (_build_upsert_multiple) raises ValidationError for the same conditions
+        # because batch errors carry structured subcodes.
         if len(alternate_keys) != len(records):
             raise ValueError(
                 f"alternate_keys and records must have the same length " f"({len(alternate_keys)} != {len(records)})"
             )
         logical_name = table_schema_name.lower()
+        # Pre-process all targets before chunking so that validation (key
+        # conflicts, label conversion) runs eagerly.  This means all records
+        # are held in memory at once, which is acceptable for typical workloads.
         targets: List[Dict[str, Any]] = []
         for alt_key, record in zip(alternate_keys, records):
             alt_key_lower = self._lowercase_keys(alt_key)
@@ -678,8 +686,8 @@ def _delete_multiple(
     ) -> Optional[str]:
         """Delete many records by GUID list via the ``BulkDelete`` action.
 
-        :param logical_name: Logical (singular) entity name.
-        :type logical_name: ``str``
+        :param table_schema_name: Schema name of the table.
+        :type table_schema_name: ``str``
         :param ids: GUIDs of records to delete.
         :type ids: ``list[str]``
 

tests/unit/data/test_multiple_chunking.py

@@ -814,6 +814,25 @@ def test_max_workers_above_cap_is_capped(self):
         mock_pool.assert_called_once_with(max_workers=_MAX_WORKERS)
         self.assertEqual(results, chunks)
 
+    def test_contextvar_propagated_to_worker_threads(self):
+        """Worker threads see the ContextVar set by the calling thread via copy_context."""
+        from PowerPlatform.Dataverse.data._odata import _CALL_SCOPE_CORRELATION_ID
+
+        captured = []
+
+        def fn(chunk):
+            captured.append(_CALL_SCOPE_CORRELATION_ID.get())
+            return chunk
+
+        token = _CALL_SCOPE_CORRELATION_ID.set("test-correlation-id")
+        try:
+            self._dispatch(fn, ["a", "b"], max_workers=2)
+        finally:
+            _CALL_SCOPE_CORRELATION_ID.reset(token)
+
+        self.assertEqual(len(captured), 2)
+        self.assertTrue(all(c == "test-correlation-id" for c in captured))
+
 
 # ---------------------------------------------------------------------------
 # _dispatch_chunks: transient error retry with jitter