From b5e01197f6277885da046500d0da9746a23f5633 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 11 Dec 2024 10:57:07 -0800 Subject: [PATCH] Adding a python http_request wrapper to create external tools Signed-off-by: Sunish Sheth --- src/databricks_ai_bridge/__init__.py | 3 + .../external_tool_request.py | 54 +++++++++++ .../test_external_tool_request.py | 93 +++++++++++++++++++ 3 files changed, 150 insertions(+) create mode 100644 src/databricks_ai_bridge/external_tool_request.py create mode 100644 tests/databricks_ai_bridge/test_external_tool_request.py diff --git a/src/databricks_ai_bridge/__init__.py b/src/databricks_ai_bridge/__init__.py index e69de29..25165ec 100644 --- a/src/databricks_ai_bridge/__init__.py +++ b/src/databricks_ai_bridge/__init__.py @@ -0,0 +1,3 @@ +from databricks_ai_bridge.external_tool_request import http_request + +__all__ = ["http_request"] diff --git a/src/databricks_ai_bridge/external_tool_request.py b/src/databricks_ai_bridge/external_tool_request.py new file mode 100644 index 0000000..787f3f8 --- /dev/null +++ b/src/databricks_ai_bridge/external_tool_request.py @@ -0,0 +1,54 @@ +import json as js +from typing import Any, Dict, Optional + +import requests +from databricks.sdk import WorkspaceClient + + +def http_request( + conn: str, + method: str, + path: str, + *, + json: Optional[Any] = None, + headers: Optional[Dict[str, str]] = None, + params: Optional[Dict[str, Any]] = None, +) -> requests.Response: + """ + Makes an HTTP request to an external function through the Databricks Workspace. + + Args: + conn (str): The connection name to use. This is required to identify the external connection. + method (str): The HTTP method to use (e.g., "GET", "POST"). This is required. + path (str): The relative path for the API endpoint. This is required. + json (Optional[Any]): JSON payload for the request. + headers (Optional[Dict[str, str]]): Additional headers for the request. + If not provided, only auth headers from connections would be passed. + params (Optional[Dict[str, Any]]): Query parameters for the request. + + Returns: + requests.Response: The HTTP response from the external function. + + Example Usage: + response = http_request( + conn="my_connection", + method="POST", + path="/api/v1/resource", + json={"key": "value"}, + headers={"extra_header_key": "extra_header_value"}, + params={"query": "example"} + ) + """ + workspaceConfig = WorkspaceClient().config + url = f"{workspaceConfig.host}/external-functions" + request_headers = workspaceConfig._header_factory() + payload = { + "connection_name": conn, + "method": method, + "path": path, + "json": js.dumps(json), + "header": headers, + "params": params, + } + + return requests.post(url, headers=request_headers, json=payload) diff --git a/tests/databricks_ai_bridge/test_external_tool_request.py b/tests/databricks_ai_bridge/test_external_tool_request.py new file mode 100644 index 0000000..326965f --- /dev/null +++ b/tests/databricks_ai_bridge/test_external_tool_request.py @@ -0,0 +1,93 @@ +from unittest.mock import MagicMock, patch + +from databricks_ai_bridge import http_request + + +@patch("databricks_ai_bridge.external_tool_request.WorkspaceClient") +@patch("databricks_ai_bridge.external_tool_request.requests.post") +def test_http_request_success(mock_post, mock_workspace_client): + # Mock the WorkspaceClient config + mock_workspace_config = MagicMock() + mock_workspace_config.host = "https://mock-host" + mock_workspace_config._header_factory.return_value = { + "Authorization": "Bearer mock-token" + } + mock_workspace_client.return_value.config = mock_workspace_config + + # Mock the POST request + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"success": True} + mock_post.return_value = mock_response + + # Call the function + response = http_request( + conn="mock_connection", + method="POST", + path="/mock-path", + json={"key": "value"}, + headers={"Custom-Header": "HeaderValue"}, + params={"query": "test"}, + ) + + # Assertions + assert response.status_code == 200 + assert response.json() == {"success": True} + mock_post.assert_called_once_with( + "https://mock-host/external-functions", + headers={ + "Authorization": "Bearer mock-token", + }, + json={ + "connection_name": "mock_connection", + "method": "POST", + "path": "/mock-path", + "json": '{"key": "value"}', + "header": { + "Custom-Header": "HeaderValue", + }, + "params": {"query": "test"}, + }, + ) + + +@patch("databricks_ai_bridge.external_tool_request.WorkspaceClient") +@patch("databricks_ai_bridge.external_tool_request.requests.post") +def test_http_request_error_response(mock_post, mock_workspace_client): + # Mock the WorkspaceClient config + mock_workspace_config = MagicMock() + mock_workspace_config.host = "https://mock-host" + mock_workspace_config._header_factory.return_value = { + "Authorization": "Bearer mock-token" + } + mock_workspace_client.return_value.config = mock_workspace_config + + # Mock the POST request to return an error + mock_response = MagicMock() + mock_response.status_code = 400 + mock_response.json.return_value = {"error": "Bad Request"} + mock_post.return_value = mock_response + + # Call the function + response = http_request( + conn="mock_connection", + method="POST", + path="/mock-path", + json={"key": "value"}, + ) + + # Assertions + assert response.status_code == 400 + assert response.json() == {"error": "Bad Request"} + mock_post.assert_called_once_with( + "https://mock-host/external-functions", + headers={"Authorization": "Bearer mock-token"}, + json={ + "connection_name": "mock_connection", + "method": "POST", + "path": "/mock-path", + "json": '{"key": "value"}', + "header": None, + "params": None, + }, + )