flyte 0.0.1b0__py3-none-any.whl

flyte/remote/_client/auth/_grpc_utils/default_metadata_interceptor.py

@@ -0,0 +1,151 @@
+import typing
+
+import grpc.aio
+from grpc.aio import ClientCallDetails, Metadata
+
+_default_metadata = Metadata(("accept", "application/grpc"))
+
+
+def with_metadata(call_details: ClientCallDetails, new_metadata: Metadata) -> ClientCallDetails:
+    metadata = Metadata()
+    for k, v in call_details.metadata.keys():
+        # Add existing metadata to the new metadata object
+        metadata.add(key=k, value=v)
+    for k, v in new_metadata.keys():
+        metadata.add(key=k, value=v)
+
+    # return call_details._replace(metadata=metadata), None
+    return ClientCallDetails(
+        method=call_details.method,
+        timeout=call_details.timeout,
+        metadata=metadata,
+        credentials=call_details.credentials,
+        wait_for_ready=call_details.wait_for_ready,
+    )
+
+
+class _BaseDefaultMetadataInterceptor:
+    """
+    Base class for all default metadata interceptors that provides common functionality.
+    """
+
+    async def _inject_default_metadata(self, call_details: grpc.aio.ClientCallDetails):
+        """
+        Injects default metadata into the client call details.
+
+        This method adds all key-value pairs from the default metadata dictionary to the
+        client call details metadata. If the client call details don't have metadata,
+        a new Metadata object is created.
+
+        :param call_details: The client call details to inject metadata into
+        :return: A new ClientCallDetails object with the injected metadata
+        """
+        return with_metadata(call_details, _default_metadata)
+
+
+class DefaultMetadataUnaryUnaryInterceptor(_BaseDefaultMetadataInterceptor, grpc.aio.UnaryUnaryClientInterceptor):
+    """
+    Interceptor for unary-unary RPC calls that adds default metadata.
+    """
+
+    async def intercept_unary_unary(
+        self,
+        continuation: typing.Callable,
+        client_call_details: grpc.aio.ClientCallDetails,
+        request: typing.Any,
+    ):
+        """
+        Intercepts unary-unary calls and injects default metadata.
+
+        This method adds default metadata to the client call details before continuing the RPC call chain.
+
+        :param continuation: Function to continue the RPC call chain with the updated call details
+        :param client_call_details: Details about the RPC call including method, timeout, metadata, credentials,
+        and wait_for_ready
+        :param request: The request message to be sent to the server
+        :return: The response from the RPC call
+        """
+        updated_call_details = await self._inject_default_metadata(client_call_details)
+        return await (await continuation(updated_call_details, request))
+
+
+class DefaultMetadataUnaryStreamInterceptor(_BaseDefaultMetadataInterceptor, grpc.aio.UnaryStreamClientInterceptor):
+    """
+    Interceptor for unary-stream RPC calls that adds default metadata.
+    """
+
+    async def intercept_unary_stream(
+        self,
+        continuation: typing.Callable,
+        client_call_details: grpc.aio.ClientCallDetails,
+        request: typing.Any,
+    ):
+        """
+        Intercepts unary-stream calls and injects default metadata.
+
+        This method adds default metadata to the client call details before continuing the RPC call chain.
+
+        :param continuation: Function to continue the RPC call chain with the updated call details
+        :param client_call_details: Details about the RPC call including method, timeout, metadata, credentials,
+        and wait_for_ready
+        :param request: The request message to be sent to the server
+        :return: A stream of responses from the RPC call
+        """
+        updated_call_details = await self._inject_default_metadata(client_call_details)
+        return await continuation(updated_call_details, request)
+
+
+class DefaultMetadataStreamUnaryInterceptor(_BaseDefaultMetadataInterceptor, grpc.aio.StreamUnaryClientInterceptor):
+    """
+    Interceptor for stream-unary RPC calls that adds default metadata.
+    """
+
+    async def intercept_stream_unary(
+        self,
+        continuation: typing.Callable,
+        client_call_details: grpc.aio.ClientCallDetails,
+        request_iterator: typing.Any,
+    ):
+        """
+        Intercepts stream-unary calls and injects default metadata.
+
+        This method adds default metadata to the client call details before continuing the RPC call chain.
+
+        :param continuation: Function to continue the RPC call chain with the updated call details
+        :param client_call_details: Details about the RPC call including method, timeout, metadata, credentials,
+        and wait_for_ready
+        :param request_iterator: An iterator of request messages to be sent to the server
+        :return: The response from the RPC call
+        """
+        updated_call_details = await self._inject_default_metadata(client_call_details)
+        return await continuation(updated_call_details, request_iterator)
+
+
+class DefaultMetadataStreamStreamInterceptor(_BaseDefaultMetadataInterceptor, grpc.aio.StreamStreamClientInterceptor):
+    """
+    Interceptor for stream-stream RPC calls that adds default metadata.
+    """
+
+    async def intercept_stream_stream(
+        self,
+        continuation: typing.Callable,
+        client_call_details: grpc.aio.ClientCallDetails,
+        request_iterator: typing.Any,
+    ):
+        """
+        Intercepts stream-stream calls and injects default metadata.
+
+        This method adds default metadata to the client call details before continuing the RPC call chain.
+
+        :param continuation: Function to continue the RPC call chain with the updated call details
+        :param client_call_details: Details about the RPC call including method, timeout, metadata, credentials, and
+         wait_for_ready
+        :param request_iterator: An iterator of request messages to be sent to the server
+        :return: A stream of responses from the RPC call
+        """
+        updated_call_details = await self._inject_default_metadata(client_call_details)
+        return await continuation(updated_call_details, request_iterator)
+
+
+# For backward compatibility, maintain the original class name but as a type alias
+DefaultMetadataInterceptor = DefaultMetadataUnaryUnaryInterceptor

flyte/remote/_client/auth/_keyring.py

@@ -0,0 +1,143 @@
+import hashlib  # Added import for hashing
+import typing
+from urllib.parse import urlparse  # Added import
+
+import keyring
+import pydantic
+from keyring.errors import NoKeyringError, PasswordDeleteError
+
+from flyte._logging import logger
+
+
+def strip_scheme(url: str) -> str:
+    """
+    Strips the scheme from a URL.
+    Handles cases like:
+    - dns:///foo.com -> foo.com
+    - https://foo.com -> foo.com
+    - https://foo.com/blah -> foo.com/blah
+    """
+    parsed_url = urlparse(url)
+    if parsed_url.scheme == "dns":
+        return parsed_url.path.lstrip("/")
+    return f"{parsed_url.netloc}{parsed_url.path}" if parsed_url.netloc else url
+
+
+class Credentials(pydantic.BaseModel):
+    """
+    Stores the credentials together
+    """
+
+    access_token: str
+    for_endpoint: str = "flyte-default"
+    id: str = ""
+    refresh_token: str | None = None
+    expires_in: int | None = None
+
+    @pydantic.field_validator("for_endpoint", mode="after")
+    @classmethod
+    def validate_endpoint(cls, v: str) -> str:
+        return strip_scheme(v)
+
+    @pydantic.model_validator(mode="after")
+    def compute_id(self) -> "Credentials":
+        """Computes the id field as a hash of the access_token."""
+        if self.access_token:
+            self.id = hashlib.md5(self.access_token.encode()).hexdigest()
+        return self
+
+
+class KeyringStore:
+    """
+    Methods to access Keyring Store.
+    """
+
+    _access_token_key = "access_token"
+    _refresh_token_key = "refresh_token"
+
+    @staticmethod
+    def store(credentials: Credentials) -> Credentials:
+        """
+        Stores the provided credentials in the system keyring.
+
+        This method stores the access token, refresh token (if available), and ID token (if available)
+        in the system keyring, using the endpoint as the service name and specific key names for each token type.
+
+        :param credentials: The credentials object containing tokens to store
+        :return: The same credentials object that was passed in
+        :raises: Logs but does not raise NoKeyringError if the system keyring is not available
+        """
+        try:
+            if credentials.refresh_token:
+                keyring.set_password(
+                    credentials.for_endpoint,
+                    KeyringStore._refresh_token_key,
+                    credentials.refresh_token,
+                )
+            keyring.set_password(
+                credentials.for_endpoint,
+                KeyringStore._access_token_key,
+                credentials.access_token,
+            )
+        except NoKeyringError as e:
+            logger.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
+        return credentials
+
+    @staticmethod
+    def retrieve(for_endpoint: str) -> typing.Optional[Credentials]:
+        """
+        Retrieves stored credentials from the system keyring for the specified endpoint.
+
+        This method attempts to retrieve the access token, refresh token, and ID token from the system keyring
+        using the endpoint as the service name. The endpoint URL scheme is stripped before lookup.
+
+        :param for_endpoint: The endpoint URL to retrieve credentials for
+        :return: A Credentials object containing the retrieved tokens, or None if no tokens were found
+                 or if the system keyring is not available
+        """
+        for_endpoint = strip_scheme(for_endpoint)
+        try:
+            refresh_token = keyring.get_password(for_endpoint, KeyringStore._refresh_token_key)
+            access_token = keyring.get_password(for_endpoint, KeyringStore._access_token_key)
+        except NoKeyringError as e:
+            logger.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
+            return None
+
+        if not access_token:
+            logger.debug("No access token found in keyring.")
+            return None
+
+        return Credentials(
+            access_token=access_token,
+            refresh_token=refresh_token,
+            for_endpoint=for_endpoint,
+            expires_in=None,
+        )
+
+    @staticmethod
+    def delete(for_endpoint: str):
+        """
+        Deletes all stored credentials for the specified endpoint from the system keyring.
+
+        This method attempts to delete the access token, refresh token, and ID token from the system keyring
+        using the endpoint as the service name. The endpoint URL scheme is stripped before lookup.
+
+        :param for_endpoint: The endpoint URL to delete credentials for
+        """
+        for_endpoint = strip_scheme(for_endpoint)
+
+        def _delete_key(key):
+            """
+            Helper function to delete a specific key from the keyring.
+
+            :param key: The key name to delete
+            """
+            try:
+                keyring.delete_password(for_endpoint, key)
+            except PasswordDeleteError as e:
+                logger.debug(f"Key {key} not found in key store, Ignoring. Error: {e}")
+            except NoKeyringError as e:
+                logger.debug(f"KeyRing not available, Key {key} deletion failed. Error: {e}")
+
+        _delete_key(KeyringStore._access_token_key)
+        _delete_key(KeyringStore._refresh_token_key)

flyte/remote/_client/auth/_token_client.py

@@ -0,0 +1,260 @@
+import asyncio
+import base64
+import enum
+import typing
+import urllib.parse
+from datetime import datetime, timedelta
+
+import httpx
+import pydantic
+
+from flyte._logging import logger
+from flyte.remote._client.auth.errors import AuthenticationError, AuthenticationPending
+
+utf_8 = "utf-8"
+
+# Errors that Token endpoint will return
+error_slow_down = "slow_down"
+error_auth_pending = "authorization_pending"
+
+
+# Grant Types
+class GrantType(str, enum.Enum):
+    CLIENT_CREDS = "client_credentials"
+    DEVICE_CODE = "urn:ietf:params:oauth:grant-type:device_code"
+    REFRESH_TOKEN = "refresh_token"
+
+
+class DeviceCodeResponse(pydantic.BaseModel):
+    """
+    Response from device auth flow endpoint
+    {
+        'device_code': 'code',
+         'user_code': 'BNDJJFXL',
+         'verification_uri': 'url',
+         'expires_in': 600,
+         'interval': 5
+    }
+
+    Attributes:
+        device_code (str): The device verification code.
+        user_code (str): The user-facing code that should be entered on the verification page.
+        verification_uri (str): The URL where the user should enter the user_code.
+        expires_in (int): The lifetime in seconds of the device code and user code.
+        interval (int): The minimum amount of time in seconds to wait between polling requests.
+    """
+
+    device_code: str
+    user_code: str
+    verification_uri: str
+    expires_in: int
+    interval: int
+
+    @classmethod
+    def from_json_response(cls, j: typing.Dict) -> "DeviceCodeResponse":
+        """
+        Create a DeviceCodeResponse instance from a JSON response dictionary.
+
+        :param j: The JSON response dictionary containing device code information
+        :return: A new instance with values from the JSON response
+        """
+        return cls(
+            device_code=j["device_code"],
+            user_code=j["user_code"],
+            verification_uri=j["verification_uri"],
+            expires_in=j["expires_in"],
+            interval=j["interval"],
+        )
+
+
+def get_basic_authorization_header(client_id: str, client_secret: str) -> str:
+    """
+    This function transforms the client id and the client secret into a header that conforms with http basic auth.
+    It joins the id and the secret with a : then base64 encodes it, then adds the appropriate text. Secrets are
+    first URL encoded to escape illegal characters.
+
+    :param client_id: The client ID for authentication
+    :param client_secret: The client secret for authentication
+    :rtype: str
+    """
+    encoded = urllib.parse.quote_plus(client_secret)
+    concatenated = "{}:{}".format(client_id, encoded)
+    return "Basic {}".format(base64.b64encode(concatenated.encode(utf_8)).decode(utf_8))
+
+
+async def get_token(
+    token_endpoint: str,
+    http_session: httpx.AsyncClient,
+    scopes: typing.Optional[typing.List[str]] = None,
+    authorization_header: typing.Optional[str] = None,
+    client_id: typing.Optional[str] = None,
+    device_code: typing.Optional[str] = None,
+    audience: typing.Optional[str] = None,
+    grant_type: GrantType = GrantType.CLIENT_CREDS,
+    http_proxy_url: typing.Optional[str] = None,
+    verify: typing.Optional[typing.Union[bool, str]] = None,
+    refresh_token: typing.Optional[str] = None,
+) -> typing.Tuple[str, str, int]:
+    """
+    Retrieves an access token from the specified token endpoint.
+
+    :param token_endpoint: The endpoint URL for token retrieval
+    :param http_session: HTTP session to use for requests
+    :param scopes: Optional list of scopes to request during authentication
+    :param authorization_header: Optional authorization header value
+    :param client_id: Optional client ID for authentication
+    :param device_code: Optional device code for device flow authentication
+    :param audience: Optional audience for the token
+    :param grant_type: The grant type to use (default: CLIENT_CREDS)
+    :param http_proxy_url: Optional HTTP proxy URL
+    :param verify: Whether to verify SSL certificates (bool or path to cert)
+    :param refresh_token: Optional refresh token for token refresh
+    :return: A tuple of (access_token, refresh_token, expires_in)
+
+    :param token_endpoint: The URL of the token endpoint
+    :param scopes: List of scopes to request
+    :param authorization_header: Authorization header value if using client credentials
+    :param client_id: The client ID to use for authentication
+    :param device_code: The device code when using device code flow
+    :param audience: The audience value to request
+    :param grant_type: The OAuth grant type to use
+    :param http_proxy_url: HTTP proxy URL if needed
+    :param verify: SSL verification mode
+    :param http_session: An existing HTTP client session
+    :param refresh_token: Refresh token for refresh token flow
+    :return: A tuple containing (access_token, refresh_token, expires_in)
+
+    Raises:
+        AuthenticationPending: When authentication is still pending (for device code flow).
+        AuthenticationError: When authentication fails for any reason.
+    """
+    headers = {
+        "Cache-Control": "no-cache",
+        "Accept": "application/json",
+        "Content-Type": "application/x-www-form-urlencoded",
+    }
+    if authorization_header:
+        headers["Authorization"] = authorization_header
+    body = {
+        "grant_type": grant_type.value,
+    }
+    if client_id:
+        body["client_id"] = client_id
+    if device_code:
+        body["device_code"] = device_code
+    if scopes is not None:
+        body["scope"] = " ".join(s.strip("' ") for s in scopes).strip("[]'")
+    if audience:
+        body["audience"] = audience
+    if refresh_token:
+        body["refresh_token"] = refresh_token
+
+    response = await http_session.post(token_endpoint, data=body, headers=headers)
+
+    if not response.is_success:
+        j = response.json()
+        if "error" in j:
+            err = j["error"]
+            if err == error_auth_pending or err == error_slow_down:
+                raise AuthenticationPending(f"Token not yet available, try again in some time {err}")
+        logger.error("Status Code ({}) received from IDP: {}".format(response.status_code, response.text))
+        raise AuthenticationError("Status Code ({}) received from IDP: {}".format(response.status_code, response.text))
+
+    j = response.json()
+    new_refresh_token = None
+    if "refresh_token" in j:
+        new_refresh_token = j["refresh_token"]
+    else:
+        raise AuthenticationError("Token not yet available, try again in some time")
+
+    return j["access_token"], new_refresh_token, j["expires_in"]
+
+
+async def get_device_code(
+    device_auth_endpoint: str,
+    client_id: str,
+    http_session: httpx.AsyncClient,
+    *,
+    audience: typing.Optional[str] = None,
+    scopes: typing.Optional[typing.List[str]] = None,
+) -> DeviceCodeResponse:
+    """
+    Retrieves the device authentication code that can be used to authenticate the request using a browser on a
+    separate device.
+
+    :param device_auth_endpoint: The URL of the device authorization endpoint
+    :param client_id: The client ID to use for authentication
+    :param audience: The audience value to request
+    :param scopes: List of scopes to request
+    :param http_proxy_url: HTTP proxy URL if needed
+    :param verify: SSL verification mode
+    :param http_session: An existing HTTP client session
+    :return: An object containing the device code and related information
+    :raises AuthenticationError: When device code retrieval fails
+    """
+    _scope = " ".join(s.strip("' ") for s in scopes).strip("[]'") if scopes is not None else ""
+    payload = {"client_id": client_id, "scope": _scope, "audience": audience}
+    resp = await http_session.post(device_auth_endpoint, data=payload)
+    if not resp.is_success:
+        raise AuthenticationError(
+            f"Unable to retrieve Device Authentication Code for {payload},"
+            f" Status Code {resp.status_code} Reason {resp.json()}"
+        )
+    return DeviceCodeResponse.from_json_response(resp.json())
+
+
+async def poll_token_endpoint(
+    resp: DeviceCodeResponse,
+    *,
+    token_endpoint: str,
+    client_id: str,
+    http_session: httpx.AsyncClient,
+    audience: typing.Optional[str] = None,
+    scopes: typing.Optional[typing.List[str]] = None,
+    http_proxy_url: typing.Optional[str] = None,
+    verify: typing.Optional[typing.Union[bool, str]] = None,
+) -> typing.Tuple[str, str, int]:
+    """
+    Polls the token endpoint until authentication is complete or times out.
+
+    This function repeatedly calls the token endpoint at the specified interval until either:
+    1. Authentication is successful and a token is returned
+    2. The device code expires (as specified in the DeviceCodeResponse)
+
+    :param resp: The device code response from a previous call to get_device_code
+    :param token_endpoint: The URL of the token endpoint
+    :param client_id: The client ID to use for authentication
+    :param audience: The audience value to request
+    :param scopes: Space-separated list of scopes to request
+    :param http_proxy_url: HTTP proxy URL if needed
+    :param verify: SSL verification mode
+    :return: A tuple containing (access_token, refresh_token, expires_in)
+    :raises AuthenticationError: When authentication fails or times out
+    """
+    tick = datetime.now()
+    interval = timedelta(seconds=resp.interval)
+    end_time = tick + timedelta(seconds=resp.expires_in)
+    while tick < end_time:
+        try:
+            access_token, refresh_token, expires_in = await get_token(
+                token_endpoint,
+                grant_type=GrantType.DEVICE_CODE,
+                client_id=client_id,
+                audience=audience,
+                scopes=scopes,
+                device_code=resp.device_code,
+                http_proxy_url=http_proxy_url,
+                verify=verify,
+                http_session=http_session,
+            )
+            logger.debug(f"Authentication successful, access token received, expires in {expires_in} seconds")
+            return access_token, refresh_token, expires_in
+        except AuthenticationPending:
+            ...
+        except Exception as e:
+            logger.warning(f"Authentication failed, reason {e}")
+            raise e
+        logger.debug(f"Authentication pending, ..., waiting for {resp.interval} seconds")
+        await asyncio.sleep(interval.total_seconds())
+        tick = tick + interval
+    raise AuthenticationError("Authentication failed!")

flyte/remote/_client/auth/errors.py

@@ -0,0 +1,16 @@
+class AccessTokenNotFoundError(RuntimeError):
+    """
+    This error is raised with Access token is not found or if Refreshing the token fails
+    """
+
+
+class AuthenticationError(RuntimeError):
+    """
+    This is raised for any AuthenticationError
+    """
+
+
+class AuthenticationPending(RuntimeError):
+    """
+    This is raised if the token endpoint returns authentication pending
+    """

flyte/remote/_client/controlplane.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import grpc
+from flyteidl.service import admin_pb2_grpc, dataproxy_pb2_grpc
+
+from flyte._protos.secret import secret_pb2_grpc
+from flyte._protos.workflow import run_logs_service_pb2_grpc, run_service_pb2_grpc, task_service_pb2_grpc
+
+from ._protocols import (
+    DataProxyService,
+    MetadataServiceProtocol,
+    ProjectDomainService,
+    RunLogsService,
+    RunService,
+    SecretService,
+    TaskService,
+)
+from .auth import create_channel
+
+
+class ClientSet:
+    def __init__(
+        self,
+        channel: grpc.aio.Channel,
+        endpoint: str,
+        insecure: bool = False,
+        data_proxy_channel: grpc.aio.Channel | None = None,
+        **kwargs,
+    ):
+        self.endpoint = endpoint
+        self.insecure = insecure
+        self._channel = channel
+        self._admin_client = admin_pb2_grpc.AdminServiceStub(channel=channel)
+        self._task_service = task_service_pb2_grpc.TaskServiceStub(channel=channel)
+        self._run_service = run_service_pb2_grpc.RunServiceStub(channel=channel)
+        self._dataproxy = dataproxy_pb2_grpc.DataProxyServiceStub(channel=channel)
+        self._log_service = run_logs_service_pb2_grpc.RunLogsServiceStub(channel=channel)
+        self._secrets_service = secret_pb2_grpc.SecretServiceStub(channel=channel)
+
+    @classmethod
+    async def for_endpoint(cls, endpoint: str, *, insecure: bool = False, **kwargs) -> ClientSet:
+        if insecure:
+            del kwargs["api_key"]
+            del kwargs["auth_type"]
+            del kwargs["headless"]
+            del kwargs["command"]
+            del kwargs["client_id"]
+            del kwargs["client_credentials_secret"]
+            del kwargs["client_config"]
+            del kwargs["rpc_retries"]
+            del kwargs["http_proxy_url"]
+        return cls(await create_channel(endpoint, insecure=insecure, **kwargs), endpoint, insecure=insecure, **kwargs)
+
+    @classmethod
+    async def for_api_key(cls, api_key: str, **kwargs) -> ClientSet:
+        raise NotImplementedError
+
+    @classmethod
+    async def for_serverless(cls) -> ClientSet:
+        raise NotImplementedError
+
+    @classmethod
+    async def from_env(cls) -> ClientSet:
+        raise NotImplementedError
+
+    @property
+    def metadata_service(self) -> MetadataServiceProtocol:
+        return self._admin_client
+
+    @property
+    def project_domain_service(self) -> ProjectDomainService:
+        return self._admin_client
+
+    @property
+    def task_service(self) -> TaskService:
+        return self._task_service
+
+    @property
+    def run_service(self) -> RunService:
+        return self._run_service
+
+    @property
+    def dataproxy_service(self) -> DataProxyService:
+        return self._dataproxy
+
+    @property
+    def logs_service(self) -> RunLogsService:
+        return self._log_service
+
+    @property
+    def secrets_service(self) -> SecretService:
+        return self._secrets_service
+
+    async def close(self, grace: float | None = None):
+        return await self._channel.close(grace=grace)

flyte/remote/_console.py

@@ -0,0 +1,18 @@
+from urllib.parse import urlparse
+
+
+def _get_http_domain(endpoint: str, insecure: bool) -> str:
+    scheme = "http" if insecure else "https"
+    parsed = urlparse(endpoint)
+    if parsed.scheme == "dns":
+        domain = parsed.path.lstrip("/")
+    else:
+        domain = parsed.netloc
+    # TODO: make console url configurable
+    if domain.split(":")[0] == "localhost":
+        domain = "localhost:8080"
+    return f"{scheme}://{domain}"
+
+
+def get_run_url(endpoint: str, insecure: bool, project: str, domain: str, run_name: str) -> str:
+    return f"{_get_http_domain(endpoint, insecure)}/v2/runs/project/{project}/domain/{domain}/{run_name}"