Allow custom client instance to be passed to agent

flapili commented 1 year ago

First check

[X] I added a descriptive title to this issue.
[X] I used the GitHub search to find a similar request and didn't find it.
[X] I searched the Prefect documentation for this feature.

Prefect Version

2.x

Describe the current behavior

for now (https://github.com/PrefectHQ/prefect/blob/a334ec7626861fd07eb5c378c875943a806f7822/src/prefect/agent.py#L633) the prefect agent get a fresh client without passing any httpx settings

Describe the proposed behavior

allow the possibility to pass httpx settings to prefect Agent from code

Example Use

give more freedom to people who want supercharge prefect and bypass the CLI
- add http proxy
- add headers to requests (my case)
- custom SSL cert support
- etc ...

Additional context

I'm building an app which use prefect to schedule jobs and I would like to secure the api, I wrote a POC for the api: https://prefect-community.slack.com/archives/CL09KU1K7/p1681831679514849?thread_ts=1681818646.342989&cid=CL09KU1K7

# coding: utf-8
from secrets import compare_digest

from fastapi import Depends, FastAPI, status, Security, HTTPException
from fastapi.security.api_key import APIKeyQuery, APIKeyCookie, APIKeyHeader
from fastapi.openapi.utils import get_openapi
from prefect.server.api.server import create_orion_api, create_ui_app

def get_token(
    access_token_query: str | None = Security(APIKeyQuery(name="token", auto_error=False)),
    access_token_header: str | None = Security(APIKeyHeader(name="token", auto_error=False)),
    access_token_cookie: str | None = Security(APIKeyCookie(name="token", auto_error=False)),
):
    if access_token_query is not None:
        return access_token_query

    elif access_token_header is not None:
        return access_token_header

    elif access_token_cookie is not None:
        return access_token_cookie

    else:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="token is missing")

def verify_token(token=Depends(get_token)):
    if compare_digest(token, "secret") is False:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="bad token")

custom_app = FastAPI()
api_app = create_orion_api(dependencies=[Depends(verify_token)])
custom_app.mount("/api", api_app)
custom_app.mount("/", create_ui_app(ephemeral=False))

def openapi():
    partial_schema = get_openapi(title="Prefect API customized", version="0.1.0", routes=api_app.routes)
    new_schema = partial_schema.copy()
    new_schema["paths"] = {}
    for path, value in partial_schema["paths"].items():
        new_schema["paths"][f"/api{path}"] = value

    new_schema["info"]["x-logo"] = {"url": "static/prefect-logo-mark-gradient.png"}
    return new_schema

custom_app.openapi = openapi

but that mean I need to bypass the cli and handle the agent by myself, pretty simple according to https://github.com/PrefectHQ/prefect/blob/a334ec7626861fd07eb5c378c875943a806f7822/src/prefect/cli/agent.py#L164

flapili commented 1 year ago

I thinks it should be pretty easy to add this feature, I can submit a PR if needed (I'm just a bit rusted about testings, if 100% codecov is targeted I'll need onboarding / help 😅 )

billpalombi commented 1 year ago

Thanks for the thorough explanation of the motivation behind this request, @flapili! While we don't want to expose this type of config via the CLI, we think that exposing the ability to pass a full client object is a good idea for the reasons you listed.

PrefectHQ / prefect