flyte 0.0.1b0__py3-none-any.whl

union/remote/_client/auth/_keyring.py

@@ -0,0 +1,154 @@
+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 union._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 | None = ""
+    refresh_token: str | None = None
+    expires_in: int | None = None
+    id_token: str | 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 or id_token."""
+        if self.access_token:
+            self.id = hashlib.md5(self.access_token.encode()).hexdigest()
+        elif self.id_token:
+            self.id = hashlib.md5(self.id_token.encode()).hexdigest()
+        return self
+
+
+class KeyringStore:
+    """
+    Methods to access Keyring Store.
+    """
+
+    _access_token_key = "access_token"
+    _refresh_token_key = "refresh_token"
+    _id_token_key = "id_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,
+            )
+            if credentials.id_token:
+                keyring.set_password(
+                    credentials.for_endpoint,
+                    KeyringStore._id_token_key,
+                    credentials.id_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)
+            id_token = keyring.get_password(for_endpoint, KeyringStore._id_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 and not id_token:
+            return None
+        return Credentials(
+            access_token=access_token,
+            refresh_token=refresh_token,
+            for_endpoint=for_endpoint,
+            id_token=id_token,
+            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)
+        _delete_key(KeyringStore._id_token_key)

union/remote/_client/auth/_token_client.py

@@ -0,0 +1,258 @@
+import asyncio
+import base64
+import enum
+import typing
+import urllib.parse
+from datetime import datetime, timedelta
+
+import httpx
+import pydantic
+
+from union._logging import logger
+from union.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"]
+
+    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!")

union/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
+    """

union/remote/_client/controlplane.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import grpc
+from flyteidl.service import admin_pb2_grpc, dataproxy_pb2_grpc
+
+from union._protos.secret import secret_pb2_grpc
+from union._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, data_proxy_channel: grpc.aio.Channel | None = None):
+        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))
+
+    @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)

union/remote/_data.py

@@ -0,0 +1,149 @@
+import asyncio
+import hashlib
+import os
+import typing
+import uuid
+from base64 import b64encode
+from datetime import timedelta
+from functools import lru_cache
+from pathlib import Path
+from typing import Tuple
+
+import aiofiles
+import grpc
+import httpx
+from flyteidl.service import dataproxy_pb2
+from google.protobuf import duration_pb2
+
+from union._initialize import CommonInit, get_client, get_common_config, requires_client
+from union.errors import RuntimeSystemError
+
+_UPLOAD_EXPIRES_IN = timedelta(seconds=60)
+
+
+def get_extra_headers_for_protocol(native_url: str) -> typing.Dict[str, str]:
+    """
+    For Azure Blob Storage, we need to set certain headers for http request.
+    This is used when we work with signed urls.
+    :param native_url:
+    :return:
+    """
+    if native_url.startswith("abfs://"):
+        return {"x-ms-blob-type": "BlockBlob"}
+    return {}
+
+
+@lru_cache
+def hash_file(file_path: typing.Union[os.PathLike, str]) -> (bytes, str, int):
+    """
+    Hash a file and produce a digest to be used as a version
+    """
+    h = hashlib.md5()
+    size = 0
+
+    with open(file_path, "rb") as file:
+        while True:
+            # Reading is buffered, so we can read smaller chunks.
+            chunk = file.read(h.block_size)
+            if not chunk:
+                break
+            h.update(chunk)
+            size += len(chunk)
+
+    return h.digest(), h.hexdigest(), size
+
+
+async def _upload_single_file(
+    cfg: CommonInit, fp: Path, verify: bool = True, basedir: str | None = None
+) -> Tuple[str, str]:
+    md5_bytes, str_digest, _ = hash_file(fp)
+
+    try:
+        expires_in_pb = duration_pb2.Duration()
+        expires_in_pb.FromTimedelta(_UPLOAD_EXPIRES_IN)
+        client = get_client()
+        resp = await client.dataproxy_service.CreateUploadLocation(
+            dataproxy_pb2.CreateUploadLocationRequest(
+                project=cfg.project,
+                domain=cfg.domain,
+                content_md5=md5_bytes,
+                filename=fp.name,
+                expires_in=expires_in_pb,
+                filename_root=basedir,
+                add_content_md5_metadata=True,
+            )
+        )
+    except grpc.aio.AioRpcError as e:
+        if e.code() == grpc.StatusCode.NOT_FOUND:
+            raise RuntimeSystemError(
+                "NotFound", f"Failed to get signed url for {fp}, please check your project and domain."
+            )
+        elif e.code() == grpc.StatusCode.PERMISSION_DENIED:
+            raise RuntimeSystemError(
+                "PermissionDenied", f"Failed to get signed url for {fp}, please check your permissions."
+            )
+        else:
+            raise RuntimeSystemError(e.code().value, f"Failed to get signed url for {fp}.")
+    except Exception as e:
+        raise RuntimeSystemError(type(e).__name__, f"Failed to get signed url for {fp}.") from e
+
+    extra_headers = get_extra_headers_for_protocol(resp.native_url)
+    extra_headers.update(resp.headers)
+    encoded_md5 = b64encode(md5_bytes)
+    content_length = fp.stat().st_size
+
+    async with aiofiles.open(str(fp), "rb") as file:
+        extra_headers.update({"Content-Length": str(content_length), "Content-MD5": encoded_md5})
+        async with httpx.AsyncClient(verify=verify) as client:
+            await client.put(resp.signed_url, headers=extra_headers, content=file)
+        # TODO in old code we did this
+        #             if self._config.platform.insecure_skip_verify is True
+        #             else self._config.platform.ca_cert_file_path,
+
+    return str_digest, resp.native_url
+
+
+@requires_client
+async def upload_file(fp: Path, verify: bool = True) -> Tuple[str, str]:
+    """
+    Uploads a file to a remote location and returns the remote URI.
+
+    :param fp: The file path to upload.
+    :param verify: Whether to verify the certificate for HTTPS requests.
+    :return: A tuple containing the MD5 digest and the remote URI.
+    """
+    # This is a placeholder implementation. Replace with actual upload logic.
+    cfg = get_common_config()
+    if not fp.is_file():
+        raise ValueError(f"{fp} is not a single file, upload arg must be a single file.")
+    return await _upload_single_file(cfg, fp, verify=verify)
+
+
+@requires_client
+async def upload_dir(dir_path: Path, verify: bool = True) -> str:
+    """
+    Uploads a directory to a remote location and returns the remote URI.
+
+    :param dir_path: The directory path to upload.
+    :param verify: Whether to verify the certificate for HTTPS requests.
+    :return: The remote URI of the uploaded directory.
+    """
+    # This is a placeholder implementation. Replace with actual upload logic.
+    cfg = get_common_config()
+    if not dir_path.is_dir():
+        raise ValueError(f"{dir_path} is not a directory, upload arg must be a directory.")
+
+    prefix = uuid.uuid4().hex
+
+    files = dir_path.rglob("*")
+    uploaded_files = []
+    for file in files:
+        if file.is_file():
+            uploaded_files.append(_upload_single_file(cfg, file, verify=verify, basedir=prefix))
+
+    urls = await asyncio.gather(*uploaded_files)
+    native_url = urls[0][1]  # Assuming all files are uploaded to the same prefix
+    # native_url is of the form s3://my-s3-bucket/flytesnacks/development/{prefix}/source/empty.md
+    uri = native_url.split(prefix)[0] + "/" + prefix
+
+    return uri