prefect_databricks.rest

This is a module containing generic REST tasks.

HTTPMethod

Bases: Enum

Available HTTP request methods.

Source code in prefect_databricks/rest.py

class HTTPMethod(Enum):
    """
    Available HTTP request methods.
    """

    GET = "get"
    POST = "post"
    PUT = "put"
    DELETE = "delete"
    PATCH = "patch"

execute_endpoint(endpoint, databricks_credentials, http_method=HTTPMethod.GET, params=None, json=None, **kwargs) async

Generic function for executing REST endpoints.

Parameters:

Name	Type	Description	Default
endpoint	str	The endpoint route.	required
databricks_credentials	DatabricksCredentials	Credentials to use for authentication with Databricks.	required
http_method	HTTPMethod	Either GET, POST, PUT, DELETE, or PATCH.	GET
params	Dict[str, Any]	URL query parameters in the request.	None
json	Dict[str, Any]	JSON serializable object to include in the body of the request.	None
**kwargs	Dict[str, Any]	Additional keyword arguments to pass.	{}

Returns:

Type	Description
Response	The httpx.Response from interacting with the endpoint.

Examples:

Lists jobs on the Databricks instance.

from prefect import flow
from prefect_databricks import DatabricksCredentials
from prefect_databricks.rest import execute_endpoint
@flow
def example_execute_endpoint_flow():
    endpoint = "/2.1/jobs/list"
    databricks_credentials = DatabricksCredentials.load("my-block")
    params = {
        "limit": 5,
        "offset": None,
        "expand_tasks": True,
    }
    response = execute_endpoint(
        endpoint,
        databricks_credentials,
        params=params
    )
    return response.json()

Source code in prefect_databricks/rest.py

@task
async def execute_endpoint(
    endpoint: str,
    databricks_credentials: "DatabricksCredentials",
    http_method: HTTPMethod = HTTPMethod.GET,
    params: Dict[str, Any] = None,
    json: Dict[str, Any] = None,
    **kwargs: Dict[str, Any],
) -> httpx.Response:
    """
    Generic function for executing REST endpoints.

    Args:
        endpoint: The endpoint route.
        databricks_credentials: Credentials to use for authentication with Databricks.
        http_method: Either GET, POST, PUT, DELETE, or PATCH.
        params: URL query parameters in the request.
        json: JSON serializable object to include in the body of the request.
        **kwargs: Additional keyword arguments to pass.

    Returns:
        The httpx.Response from interacting with the endpoint.

    Examples:
        Lists jobs on the Databricks instance.
        ```python
        from prefect import flow
        from prefect_databricks import DatabricksCredentials
        from prefect_databricks.rest import execute_endpoint
        @flow
        def example_execute_endpoint_flow():
            endpoint = "/2.1/jobs/list"
            databricks_credentials = DatabricksCredentials.load("my-block")
            params = {
                "limit": 5,
                "offset": None,
                "expand_tasks": True,
            }
            response = execute_endpoint(
                endpoint,
                databricks_credentials,
                params=params
            )
            return response.json()
        ```
    """
    if isinstance(http_method, HTTPMethod):
        http_method = http_method.value

    if params is not None:
        stripped_params = strip_kwargs(**params)
    else:
        stripped_params = None

    if json is not None:
        kwargs["json"] = strip_kwargs(**json)

    async with databricks_credentials.get_client() as client:
        response = await getattr(client, http_method)(
            endpoint, params=stripped_params, **kwargs
        )

    return response

serialize_model(obj)

Recursively serializes pydantic.BaseModel into JSON; returns original obj if not a BaseModel.

Parameters:

Name	Type	Description	Default
obj	Any	Input object to serialize.	required

Returns:

Type	Description
Any	Serialized version of object.

Source code in prefect_databricks/rest.py

def serialize_model(obj: Any) -> Any:
    """
    Recursively serializes `pydantic.BaseModel` into JSON;
    returns original obj if not a `BaseModel`.

    Args:
        obj: Input object to serialize.

    Returns:
        Serialized version of object.
    """
    if isinstance(obj, list):
        return [serialize_model(o) for o in obj]
    elif isinstance(obj, Dict):
        return {k: serialize_model(v) for k, v in obj.items()}

    if isinstance(obj, BaseModel):
        return obj.model_dump(mode="json")
    elif isinstance(obj, Enum):
        return obj.value
    return obj

strip_kwargs(**kwargs)

Recursively drops keyword arguments if value is None, and serializes any pydantic.BaseModel types.

Parameters:

Name	Type	Description	Default
**kwargs	Dict	Input keyword arguments.	{}

Returns:

Type	Description
Dict	Stripped version of kwargs.

Source code in prefect_databricks/rest.py

def strip_kwargs(**kwargs: Dict) -> Dict:
    """
    Recursively drops keyword arguments if value is None,
    and serializes any `pydantic.BaseModel` types.

    Args:
        **kwargs: Input keyword arguments.

    Returns:
        Stripped version of kwargs.
    """
    stripped_dict = {}
    for k, v in kwargs.items():
        v = serialize_model(v)
        if isinstance(v, dict):
            v = strip_kwargs(**v)
        if v is not None:
            stripped_dict[k] = v
    return stripped_dict or {}