From 47d5a4499fc7e1f49bafb9f1868c00fad8353c7c Mon Sep 17 00:00:00 2001 From: aravind-segu Date: Fri, 28 Feb 2025 15:18:22 -0800 Subject: [PATCH 1/4] Ensure Invokers Rights only looks at main thread for invokers token Signed-off-by: aravind-segu --- databricks/sdk/credentials_provider.py | 4 ++-- tests/test_model_serving_auth.py | 21 ++++++++++++++++++++- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 24d01f678..2c76765a4 100644 --- a/databricks/sdk/credentials_provider.py +++ b/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"] diff --git a/tests/test_model_serving_auth.py b/tests/test_model_serving_auth.py index 49aed33a5..1adec2369 100644 --- a/tests/test_model_serving_auth.py +++ b/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): From 1e29e526b2ce91a7837f90d5ec0c77a4a52d791c Mon Sep 17 00:00:00 2001 From: aravind-segu Date: Fri, 28 Feb 2025 15:21:58 -0800 Subject: [PATCH 2/4] Formatting --- tests/test_model_serving_auth.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_model_serving_auth.py b/tests/test_model_serving_auth.py index 1adec2369..d0b75c8cf 100644 --- a/tests/test_model_serving_auth.py +++ b/tests/test_model_serving_auth.py @@ -114,6 +114,7 @@ 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,7 +146,7 @@ def test_agent_user_credentials(monkeypatch, mocker): assert headers.get("Authorization") == f'Bearer {invokers_token_val}' # Test invokers token in child thread - + successful_authentication_event = threading.Event() def authenticate(): @@ -155,14 +156,14 @@ def authenticate(): assert (cfg.host == 'x') assert headers.get("Authorization") == f'Bearer databricks_invokers_token_v2' successful_authentication_event.set() - except Exception as e: + except Exception: successful_authentication_event.clear() thread = threading.Thread(target=authenticate) thread.start() thread.join() - assert(successful_authentication_event.is_set()) + 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 From aa7997074ea8661adff8267c7e26ac2b18bfe43d Mon Sep 17 00:00:00 2001 From: aravind-segu Date: Fri, 28 Feb 2025 15:38:17 -0800 Subject: [PATCH 3/4] Fix formatting --- tests/test_model_serving_auth.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_model_serving_auth.py b/tests/test_model_serving_auth.py index d0b75c8cf..3a2621592 100644 --- a/tests/test_model_serving_auth.py +++ b/tests/test_model_serving_auth.py @@ -153,8 +153,8 @@ 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' + assert cfg.host == "x" + assert headers.get("Authorization") == f"Bearer databricks_invokers_token_v2" successful_authentication_event.set() except Exception: successful_authentication_event.clear() From da2c6a954cb42de928757801a83b143102a638c6 Mon Sep 17 00:00:00 2001 From: aravind-segu Date: Tue, 4 Mar 2025 09:47:54 -0800 Subject: [PATCH 4/4] Fix formatting and add change log --- NEXT_CHANGELOG.md | 1 + tests/test_model_serving_auth.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 44956e0ce..3f14d587e 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -11,5 +11,6 @@ ### Internal Changes * Introduce automated tagging ([#888](https://github.com/databricks/databricks-sdk-py/pull/888)) * Update Jobs GetJob API to support paginated responses ([#869](https://github.com/databricks/databricks-sdk-py/pull/869)). +* Update On Behalf Of User Authentication in Multithreaded applications ([#907](https://github.com/databricks/databricks-sdk-py/pull/907)) ### API Changes diff --git a/tests/test_model_serving_auth.py b/tests/test_model_serving_auth.py index 3a2621592..b3c6a7fc3 100644 --- a/tests/test_model_serving_auth.py +++ b/tests/test_model_serving_auth.py @@ -163,7 +163,7 @@ def authenticate(): thread.start() thread.join() - assert (successful_authentication_event.is_set()) + 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