flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/remote/_client/auth/_authenticators/pkce.py

@@ -17,7 +17,6 @@ from urllib.parse import urlencode as _urlencode
 import click
 import httpx
 import pydantic
-from h11 import Response
 
 from flyte._logging import logger
 from flyte.remote._client.auth._authenticators.base import Authenticator
@@ -123,7 +122,7 @@ class PKCEAuthenticator(Authenticator):
             try:
                 return await self._auth_client.refresh_access_token(self._creds)
             except AccessTokenNotFoundError:
-                logger.warning("Failed to refresh token. Kicking off a full authorization flow.")
+                logger.warning("Logging in...")
 
         return await self._auth_client.get_creds_from_remote()
 
@@ -363,19 +362,17 @@ class AuthorizationClient(object):
 
         data.update(self._refresh_access_token_params)
 
-        async with typing.cast(
-            typing.AsyncContextManager[Response],
-            self._http_session.post(
-                url=self._token_endpoint,
-                data=data,
-                headers=self._headers,
-                follow_redirects=False,
-            ),
-        ) as resp:
-            if resp.status_code != _StatusCodes.OK:
-                raise AccessTokenNotFoundError(f"Non-200 returned from refresh token endpoint {resp.status_code}")
+        resp: httpx.Response = await self._http_session.post(
+            url=self._token_endpoint,
+            data=data,
+            headers=self._headers,
+            follow_redirects=False,
+        )
 
-            return await self._credentials_from_response(resp)
+        if resp.status_code != _StatusCodes.OK:
+            raise AccessTokenNotFoundError(f"Non-200 returned from refresh token endpoint {resp.status_code}")
+
+        return await self._credentials_from_response(resp)
 
 
 class OAuthCallbackHandler:
@@ -410,10 +407,12 @@ class OAuthCallbackHandler:
         :param reader: The StreamReader for reading the incoming request
         :param writer: The StreamWriter for writing the response
         """
-        data = await reader.read(1024)
-        message = data.decode()
-        headers = message.split("\r\n")
-        path = headers[0].split(" ")[1]
+        # Read only the first line of the HTTP request (e.g., "GET /callback?code=...&state=... HTTP/1.1")
+        # Using readline() instead of read() because read() waits for EOF, which won't come
+        # until the client closes the connection - but the client is waiting for our response first.
+        request_line = await reader.readline()
+        # request_line looks like this: 'GET /callback?code=ABC&state=FOO HTTP/1.1'
+        path = request_line.decode().split(" ")[1]
         url = _urlparse.urlparse(path)
         if url.path.strip("/") == self.redirect_path.strip("/"):
             response = f"HTTP/1.1 {_StatusCodes.OK.value} {_StatusCodes.OK.phrase}\r\n"

flyte/remote/_client/auth/_channel.py

@@ -7,6 +7,7 @@ import httpx
 from grpc.experimental.aio import init_grpc_aio
 
 from flyte._logging import logger
+from flyte._utils.org_discovery import hostname_from_url
 
 from ._authenticators.base import get_async_session
 from ._authenticators.factory import (
@@ -30,21 +31,26 @@ def bootstrap_ssl_from_server(endpoint: str) -> grpc.ChannelCredentials:
     :param endpoint: The endpoint URL to retrieve the SSL certificate from, may include port number
     :return: gRPC channel credentials created from the retrieved certificate
     """
+    hostname = hostname_from_url(endpoint)
+
     # Get port from endpoint or use 443
-    endpoint_parts = endpoint.rsplit(":", 1)
+    endpoint_parts = hostname.rsplit(":", 1)
     if len(endpoint_parts) == 2 and endpoint_parts[1].isdigit():
         server_address = (endpoint_parts[0], int(endpoint_parts[1]))
     else:
-        logger.warning(f"Unrecognized port in endpoint [{endpoint}], defaulting to 443.")
-        server_address = (endpoint, 443)
+        logger.warning(f"Unrecognized port in endpoint [{hostname}], defaulting to 443.")
+        server_address = (hostname, 443)
 
-    # Run the blocking SSL certificate retrieval in a thread pool
-    cert = ssl.get_server_certificate(server_address)
+    # Run the blocking SSL certificate retrieval with a timeout
+    logger.debug(f"Retrieving SSL certificate from {server_address}")
+    cert = ssl.get_server_certificate(server_address, timeout=10)
     return grpc.ssl_channel_credentials(str.encode(cert))
 
 
 async def create_channel(
-    endpoint: str,
+    endpoint: str | None,
+    api_key: str | None = None,
+    /,
     insecure: typing.Optional[bool] = None,
     insecure_skip_verify: typing.Optional[bool] = False,
     ca_cert_file_path: typing.Optional[str] = None,
@@ -66,6 +72,7 @@ async def create_channel(
     and create authentication interceptors that perform async operations.
 
     :param endpoint: The endpoint URL for the gRPC channel
+    :param api_key: API key for authentication; if provided, it will be used to detect the endpoint and credentials.
     :param insecure: Whether to use an insecure channel (no SSL)
     :param insecure_skip_verify: Whether to skip SSL certificate verification
     :param ca_cert_file_path: Path to CA certificate file for SSL verification
@@ -104,23 +111,40 @@ async def create_channel(
             - refresh_access_token_params: Parameters to add when refreshing access token
     :return: grpc.aio.Channel with authentication interceptors configured
     """
+    assert endpoint or api_key, "Either endpoint or api_key must be specified"
+
+    if api_key:
+        from flyte.remote._client.auth._auth_utils import decode_api_key
 
-    if not ssl_credentials:
-        if insecure_skip_verify:
-            ssl_credentials = bootstrap_ssl_from_server(endpoint)
-        elif ca_cert_file_path:
-            import aiofiles
+        endpoint, client_id, client_secret, _org = decode_api_key(api_key)
+        kwargs["auth_type"] = "ClientSecret"
+        kwargs["client_id"] = client_id
+        kwargs["client_secret"] = client_secret
+        kwargs["client_credentials_secret"] = client_secret
 
-            async with aiofiles.open(ca_cert_file_path, "rb") as f:
-                st_cert = f.read()
-            ssl_credentials = grpc.ssl_channel_credentials(st_cert)
-        else:
-            ssl_credentials = grpc.ssl_channel_credentials()
+    assert endpoint, "Endpoint must be specified by this point"
 
     # Create an unauthenticated channel first to use to get the server metadata
     if insecure:
-        unauthenticated_channel = grpc.aio.insecure_channel(endpoint, **kwargs)
+        insecure_kwargs = {}
+        if kw_opts := kwargs.get("options"):
+            insecure_kwargs["options"] = kw_opts
+        if compression:
+            insecure_kwargs["compression"] = compression
+        unauthenticated_channel = grpc.aio.insecure_channel(endpoint, **insecure_kwargs)
     else:
+        # Only create SSL credentials if not provided and also only when using secure channel.
+        if not ssl_credentials:
+            if insecure_skip_verify:
+                ssl_credentials = bootstrap_ssl_from_server(endpoint)
+            elif ca_cert_file_path:
+                import aiofiles
+
+                async with aiofiles.open(ca_cert_file_path, "rb") as f:
+                    st_cert = await f.read()
+                ssl_credentials = grpc.ssl_channel_credentials(st_cert)
+            else:
+                ssl_credentials = grpc.ssl_channel_credentials()
         unauthenticated_channel = grpc.aio.secure_channel(
             target=endpoint,
             credentials=ssl_credentials,
@@ -173,7 +197,12 @@ async def create_channel(
     interceptors.extend(auth_interceptors)
 
     if insecure:
-        return grpc.aio.insecure_channel(endpoint, interceptors=interceptors, **kwargs)
+        insecure_kwargs = {}
+        if kw_opts := kwargs.get("options"):
+            insecure_kwargs["options"] = kw_opts
+        if compression:
+            insecure_kwargs["compression"] = compression
+        return grpc.aio.insecure_channel(endpoint, interceptors=interceptors, **insecure_kwargs)
 
     return grpc.aio.secure_channel(
         target=endpoint,

flyte/remote/_client/auth/_client_config.py

@@ -1,3 +1,4 @@
+import asyncio
 import typing
 from abc import abstractmethod
 
@@ -6,7 +7,7 @@ import pydantic
 from flyteidl.service.auth_pb2 import OAuth2MetadataRequest, PublicClientAuthConfigRequest
 from flyteidl.service.auth_pb2_grpc import AuthMetadataServiceStub
 
-AuthType = typing.Literal["ClientSecret", "Pkce", "ExternalCommand", "DeviceFlow"]
+AuthType = typing.Literal["ClientSecret", "Pkce", "ExternalCommand", "DeviceFlow", "Passthrough"]
 
 
 class ClientConfig(pydantic.BaseModel):
@@ -69,8 +70,9 @@ class RemoteClientConfigStore(ClientConfigStore):
         Retrieves the ClientConfig from the given grpc.Channel assuming  AuthMetadataService is available
         """
         metadata_service = AuthMetadataServiceStub(self._unauthenticated_channel)
-        public_client_config = await metadata_service.GetPublicClientConfig(PublicClientAuthConfigRequest())
-        oauth2_metadata = await metadata_service.GetOAuth2Metadata(OAuth2MetadataRequest())
+        oauth2_metadata_task = metadata_service.GetOAuth2Metadata(OAuth2MetadataRequest())
+        public_client_config_task = metadata_service.GetPublicClientConfig(PublicClientAuthConfigRequest())
+        oauth2_metadata, public_client_config = await asyncio.gather(oauth2_metadata_task, public_client_config_task)
         return ClientConfig(
             token_endpoint=oauth2_metadata.token_endpoint,
             authorization_endpoint=oauth2_metadata.authorization_endpoint,

flyte/remote/_client/auth/_keyring.py

@@ -81,6 +81,8 @@ class KeyringStore:
             )
         except NoKeyringError as e:
             logger.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
+        except Exception as e:
+            logger.debug(f"Failed to store tokens in keyring. Error: {e}")
         return credentials
 
     @staticmethod
@@ -96,16 +98,23 @@ class KeyringStore:
                  or if the system keyring is not available
         """
         for_endpoint = strip_scheme(for_endpoint)
+        access_token: str | None = None
         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
+        except Exception as e:
+            logger.debug(f"Failed to retrieve tokens from keyring. Error: {e}")
+            return None
 
         if not access_token:
-            logger.debug("No access token found in keyring.")
-            return None
+            if not refresh_token:
+                logger.debug("No access token found in keyring.")
+                return None
+            else:
+                access_token = ""
 
         return Credentials(
             access_token=access_token,
@@ -138,6 +147,10 @@ class KeyringStore:
                 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}")
+            except NotImplementedError as e:
+                logger.debug(f"Key {key} deletion not implemented in keyring backend. Error: {e}")
+            except Exception as e:
+                logger.debug(f"Failed to delete key {key} from keyring. Error: {e}")
 
         _delete_key(KeyringStore._access_token_key)
         _delete_key(KeyringStore._refresh_token_key)

flyte/remote/_client/auth/_token_client.py

@@ -94,7 +94,7 @@ async def get_token(
     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]:
+) -> typing.Tuple[str, str | None, int]:
     """
     Retrieves an access token from the specified token endpoint.
 
@@ -165,7 +165,7 @@ async def get_token(
     if "refresh_token" in j:
         new_refresh_token = j["refresh_token"]
     else:
-        raise AuthenticationError("Token not yet available, try again in some time")
+        logger.info("No refresh token received, this is expected for client credentials flow")
 
     return j["access_token"], new_refresh_token, j["expires_in"]
 
@@ -213,7 +213,7 @@ async def poll_token_endpoint(
     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]:
+) -> typing.Tuple[str, str | None, int]:
     """
     Polls the token endpoint until authentication is complete or times out.
 

flyte/remote/_client/controlplane.py

@@ -1,59 +1,218 @@
 from __future__ import annotations
 
-import grpc
-from flyteidl.service import admin_pb2_grpc, dataproxy_pb2_grpc
+import os
+from urllib.parse import urlparse
+
+# Set environment variables for gRPC, this reduces log spew and avoids unnecessary warnings
+# before importing grpc
+if "GRPC_VERBOSITY" not in os.environ:
+    os.environ["GRPC_VERBOSITY"] = "ERROR"
+    os.environ["GRPC_CPP_MIN_LOG_LEVEL"] = "ERROR"
+    # Disable fork support (stops "skipping fork() handlers")
+    os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "0"
+    # Reduce absl/glog verbosity
+    os.environ["GLOG_minloglevel"] = "2"
+    os.environ["ABSL_LOG"] = "0"
+#### Has to be before 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
+import grpc
+from flyteidl.service import admin_pb2_grpc, dataproxy_pb2_grpc, identity_pb2_grpc
+from flyteidl2.app import app_service_pb2_grpc
+from flyteidl2.secret import secret_pb2_grpc
+from flyteidl2.task import task_service_pb2_grpc
+from flyteidl2.trigger import trigger_service_pb2_grpc
+from flyteidl2.workflow import run_logs_service_pb2_grpc, run_service_pb2_grpc
 
 from ._protocols import (
+    AppService,
     DataProxyService,
+    IdentityService,
     MetadataServiceProtocol,
     ProjectDomainService,
     RunLogsService,
     RunService,
     SecretService,
     TaskService,
+    TriggerService,
 )
 from .auth import create_channel
 
 
+class Console:
+    """
+    Console URL builder for Flyte resources.
+
+    Constructs console URLs for various Flyte resources (tasks, runs, apps, triggers)
+    based on the configured endpoint and security settings.
+
+    Args:
+        endpoint: The Flyte endpoint (e.g., "dns:///localhost:8090", "https://example.com")
+        insecure: Whether to use HTTP (True) or HTTPS (False)
+
+    Example:
+        >>> console = Console("dns:///example.com", insecure=False)
+        >>> url = console.task_url(project="myproject", domain="development", task_name="mytask")
+    """
+
+    def __init__(self, endpoint: str, insecure: bool = False):
+        """
+        Initialize Console with endpoint and security configuration.
+
+        Args:
+            endpoint: The Flyte endpoint URL
+            insecure: Whether to use HTTP (True) or HTTPS (False)
+        """
+        self._endpoint = endpoint
+        self._insecure = insecure
+        self._http_domain = self._compute_http_domain()
+
+    def _compute_http_domain(self) -> str:
+        """
+        Compute the HTTP domain from the endpoint.
+
+        Internal method that extracts and normalizes the domain from various
+        endpoint formats (dns://, http://, https://).
+
+        Returns:
+            The normalized HTTP(S) domain URL
+        """
+        scheme = "http" if self._insecure else "https"
+        parsed = urlparse(self._endpoint)
+        if parsed.scheme == "dns":
+            domain = parsed.path.lstrip("/")
+        else:
+            domain = parsed.netloc or parsed.path
+
+        # TODO: make console url configurable
+        domain_split = domain.split(":")
+        if domain_split[0] == "localhost":
+            # Always use port 8080 for localhost, until the to do is done.
+            domain = "localhost:8080"
+
+        return f"{scheme}://{domain}"
+
+    def _resource_url(self, project: str, domain: str, resource: str, resource_name: str) -> str:
+        """
+        Internal helper to build a resource URL.
+
+        Args:
+            project: Project name
+            domain: Domain name
+            resource: Resource type (e.g., "tasks", "runs", "apps", "triggers")
+            resource_name: Resource identifier
+
+        Returns:
+            The full console URL for the resource
+        """
+        return f"{self._http_domain}/v2/domain/{domain}/project/{project}/{resource}/{resource_name}"
+
+    def run_url(self, project: str, domain: str, run_name: str) -> str:
+        """
+        Build console URL for a run.
+
+        Args:
+            project: Project name
+            domain: Domain name
+            run_name: Run identifier
+
+        Returns:
+            Console URL for the run
+        """
+        return self._resource_url(project, domain, "runs", run_name)
+
+    def app_url(self, project: str, domain: str, app_name: str) -> str:
+        """
+        Build console URL for an app.
+
+        Args:
+            project: Project name
+            domain: Domain name
+            app_name: App identifier
+
+        Returns:
+            Console URL for the app
+        """
+        return self._resource_url(project, domain, "apps", app_name)
+
+    def task_url(self, project: str, domain: str, task_name: str) -> str:
+        """
+        Build console URL for a task.
+
+        Args:
+            project: Project name
+            domain: Domain name
+            task_name: Task identifier
+
+        Returns:
+            Console URL for the task
+        """
+        return self._resource_url(project, domain, "tasks", task_name)
+
+    def trigger_url(self, project: str, domain: str, task_name: str, trigger_name: str) -> str:
+        """
+        Build console URL for a trigger.
+
+        Args:
+            project: Project name
+            domain: Domain name
+            task_name: Task identifier
+            trigger_name: Trigger identifier
+
+        Returns:
+            Console URL for the trigger
+        """
+        return self._resource_url(project, domain, "triggers", f"{task_name}/{trigger_name}")
+
+    @property
+    def endpoint(self) -> str:
+        """The configured endpoint."""
+        return self._endpoint
+
+    @property
+    def insecure(self) -> bool:
+        """Whether insecure (HTTP) mode is enabled."""
+        return self._insecure
+
+
 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._console = Console(self.endpoint, self.insecure)
         self._admin_client = admin_pb2_grpc.AdminServiceStub(channel=channel)
         self._task_service = task_service_pb2_grpc.TaskServiceStub(channel=channel)
+        self._app_service = app_service_pb2_grpc.AppServiceStub(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)
+        self._identity_service = identity_pb2_grpc.IdentityServiceStub(channel=channel)
+        self._trigger_service = trigger_service_pb2_grpc.TriggerServiceStub(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)
+        return cls(
+            await create_channel(endpoint, None, insecure=insecure, **kwargs), endpoint, insecure=insecure, **kwargs
+        )
 
     @classmethod
-    async def for_api_key(cls, api_key: str, **kwargs) -> ClientSet:
-        raise NotImplementedError
+    async def for_api_key(cls, api_key: str, *, insecure: bool = False, **kwargs) -> ClientSet:
+        from flyte.remote._client.auth._auth_utils import decode_api_key
+
+        # Parsing the API key is done in create_channel, but cleaner to redo it here rather than getting create_channel
+        # to return the endpoint
+        endpoint, _, _, _ = decode_api_key(api_key)
+
+        return cls(
+            await create_channel(None, api_key, insecure=insecure, **kwargs), endpoint, insecure=insecure, **kwargs
+        )
 
     @classmethod
     async def for_serverless(cls) -> ClientSet:
@@ -75,6 +234,10 @@ class ClientSet:
     def task_service(self) -> TaskService:
         return self._task_service
 
+    @property
+    def app_service(self) -> AppService:
+        return self._app_service
+
     @property
     def run_service(self) -> RunService:
         return self._run_service
@@ -91,5 +254,30 @@ class ClientSet:
     def secrets_service(self) -> SecretService:
         return self._secrets_service
 
+    @property
+    def identity_service(self) -> IdentityService:
+        return self._identity_service
+
+    @property
+    def trigger_service(self) -> TriggerService:
+        return self._trigger_service
+
+    @property
+    def console(self) -> Console:
+        """
+        Get the Console instance for this client.
+
+        Returns a Console configured with this client's endpoint and security settings.
+        Use this to build console URLs for Flyte resources.
+
+        Returns:
+            Console instance
+
+        Example:
+            >>> client = get_client()
+            >>> url = client.console.task_url(project="myproj", domain="dev", task_name="mytask")
+        """
+        return self._console
+
     async def close(self, grace: float | None = None):
         return await self._channel.close(grace=grace)

flyte/remote/_common.py

@@ -0,0 +1,66 @@
+import json
+from typing import Literal, Tuple
+
+from flyteidl2.common import list_pb2
+from google.protobuf.json_format import MessageToDict, MessageToJson
+
+
+class ToJSONMixin:
+    """
+    A mixin class that provides a method to convert an object to a JSON-serializable dictionary.
+    """
+
+    def to_dict(self) -> dict:
+        """
+        Convert the object to a JSON-serializable dictionary.
+
+        Returns:
+            dict: A dictionary representation of the object.
+        """
+        if hasattr(self, "pb2"):
+            return MessageToDict(self.pb2)
+        else:
+            return {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
+
+    def to_json(self) -> str:
+        """
+        Convert the object to a JSON string.
+
+        Returns:
+            str: A JSON string representation of the object.
+        """
+        return MessageToJson(self.pb2) if hasattr(self, "pb2") else json.dumps(self.to_dict())
+
+
+def sorting(sort_by: Tuple[str, Literal["asc", "desc"]] | None = None) -> list_pb2.Sort:
+    """
+    Create a protobuf Sort object from a sorting tuple.
+
+    :param sort_by: Tuple of (field_name, direction) for sorting, defaults to ("created_at", "asc").
+    :return: A protobuf Sort object.
+    """
+    sort_by = sort_by or ("created_at", "asc")
+    return list_pb2.Sort(
+        key=sort_by[0],
+        direction=(list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING),
+    )
+
+
+def filtering(created_by_subject: str | None = None, *filters: list_pb2.Filter) -> list[list_pb2.Filter]:
+    """
+    Create a list of filter objects, optionally including a filter by creator subject.
+
+    :param created_by_subject: Optional subject to filter by creator.
+    :param filters: Additional filters to include.
+    :return: A list of protobuf Filter objects.
+    """
+    filter_list = list(filters) if filters else []
+    if created_by_subject:
+        filter_list.append(
+            list_pb2.Filter(
+                function=list_pb2.Filter.Function.EQUAL,
+                field="created_by",
+                values=[created_by_subject],
+            ),
+        )
+    return filter_list