Ensure Invokers Rights only looks at main thread for invokers token

Signed-off-by: aravind-segu <[email protected]>

Author: aravind-segu

databricks/sdk/credentials_provider.py

@@ -769,8 +769,8 @@ def _get_model_dependency_oauth_token(self, should_retry=True) -> str:
         return self.current_token
 
     def _get_invokers_token(self):
-        current_thread = threading.current_thread()
-        thread_data = current_thread.__dict__
+        main_thread = threading.main_thread()
+        thread_data = main_thread.__dict__
         invokers_token = None
         if "invokers_token" in thread_data:
             invokers_token = thread_data["invokers_token"]

tests/test_model_serving_auth.py

@@ -114,7 +114,6 @@ def test_model_serving_auth_refresh(monkeypatch, mocker):
     # Read V2 now
     assert headers.get("Authorization") == 'Bearer databricks_sdk_unit_test_token_v2'
 
-
 def test_agent_user_credentials(monkeypatch, mocker):
     monkeypatch.setenv('IS_IN_DB_MODEL_SERVING_ENV', 'true')
     monkeypatch.setenv('DB_MODEL_SERVING_HOST_URL', 'x')
@@ -145,6 +144,26 @@ def test_agent_user_credentials(monkeypatch, mocker):
     assert (cfg.host == 'x')
     assert headers.get("Authorization") == f'Bearer {invokers_token_val}'
 
+    # Test invokers token in child thread
+    
+    successful_authentication_event = threading.Event()
+
+    def authenticate():
+        try:
+            cfg = Config(credentials_strategy=ModelServingUserCredentials())
+            headers = cfg.authenticate()
+            assert (cfg.host == 'x')
+            assert headers.get("Authorization") == f'Bearer databricks_invokers_token_v2'
+            successful_authentication_event.set()
+        except Exception as e:
+            successful_authentication_event.clear()
+
+    thread = threading.Thread(target=authenticate)
+
+    thread.start()
+    thread.join()
+    assert(successful_authentication_event.is_set())
+
 
 # If this credential strategy is being used in a non model serving environments then use default credential strategy instead
 def test_agent_user_credentials_in_non_model_serving_environments(monkeypatch):