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

flyte/remote/_secret.py

@@ -1,31 +1,46 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import AsyncGenerator, Iterator, Literal, Union
+from typing import AsyncIterator, Literal, Union
 
 import rich.repr
+from flyteidl2.secret import definition_pb2, payload_pb2
 
-from flyte._api_commons import syncer
-from flyte._initialize import get_client, get_common_config, requires_client
-from flyte._protos.secret import definition_pb2, payload_pb2
+from flyte._initialize import ensure_client, get_client, get_init_config
+from flyte.remote._common import ToJSONMixin
+from flyte.syncify import syncify
 
 SecretTypes = Literal["regular", "image_pull"]
 
 
 @dataclass
-class Secret:
+class Secret(ToJSONMixin):
     pb2: definition_pb2.Secret
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def create(cls, name: str, value: Union[str, bytes], type: SecretTypes = "regular"):
-        cfg = get_common_config()
-        secret_type = (
-            definition_pb2.SecretType.SECRET_TYPE_GENERIC
-            if type == "regular"
-            else definition_pb2.SecretType.SECRET_TYPE_IMAGE_PULL_SECRET
-        )
+        """
+        Create a new secret.
+
+        :param name: The name of the secret.
+        :param value: The secret value as a string or bytes.
+        :param type: Type of secret - either "regular" or "image_pull".
+        """
+        ensure_client()
+        cfg = get_init_config()
+        project = cfg.project
+        domain = cfg.domain
+
+        if type == "regular":
+            secret_type = definition_pb2.SecretType.SECRET_TYPE_GENERIC
+
+        else:
+            secret_type = definition_pb2.SecretType.SECRET_TYPE_IMAGE_PULL_SECRET
+            if project or domain:
+                raise ValueError(
+                    f"Project `{project}` or domain `{domain}` should not be set when creating the image pull secret."
+                )
 
         if isinstance(value, str):
             secret = definition_pb2.SecretSpec(
@@ -41,19 +56,25 @@ class Secret:
             request=payload_pb2.CreateSecretRequest(
                 id=definition_pb2.SecretIdentifier(
                     organization=cfg.org,
-                    project=cfg.project,
-                    domain=cfg.domain,
+                    project=project,
+                    domain=domain,
                     name=name,
                 ),
                 secret_spec=secret,
             ),
         )
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def get(cls, name: str) -> Secret:
-        cfg = get_common_config()
+        """
+        Retrieve a secret by name.
+
+        :param name: The name of the secret to retrieve.
+        :return: A Secret object.
+        """
+        ensure_client()
+        cfg = get_init_config()
         resp = await get_client().secrets_service.GetSecret(
             request=payload_pb2.GetSecretRequest(
                 id=definition_pb2.SecretIdentifier(
@@ -66,33 +87,46 @@ class Secret:
         )
         return Secret(pb2=resp.secret)
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
-    async def listall(cls, limit: int = 100) -> Union[Iterator[Secret], AsyncGenerator[Secret, None]]:
-        cfg = get_common_config()
-        token = None
+    async def listall(cls, limit: int = 10) -> AsyncIterator[Secret]:
+        """
+        List all secrets in the current project and domain.
+
+        :param limit: Maximum number of secrets to return per page.
+        :return: An async iterator of Secret objects.
+        """
+        ensure_client()
+        cfg = get_init_config()
+        per_cluster_tokens = None
         while True:
             resp = await get_client().secrets_service.ListSecrets(  # type: ignore
                 request=payload_pb2.ListSecretsRequest(
                     organization=cfg.org,
                     project=cfg.project,
                     domain=cfg.domain,
-                    token=token,
+                    per_cluster_tokens=per_cluster_tokens,
                     limit=limit,
                 ),
             )
-            token = resp.token
+            per_cluster_tokens = resp.per_cluster_tokens
+            round_items = [v for _, v in per_cluster_tokens.items() if v]
+            has_next = any(round_items)
             for r in resp.secrets:
                 yield cls(r)
-            if not token:
+            if not has_next:
                 break
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def delete(cls, name):
-        cfg = get_common_config()
+        """
+        Delete a secret by name.
+
+        :param name: The name of the secret to delete.
+        """
+        ensure_client()
+        cfg = get_init_config()
         await get_client().secrets_service.DeleteSecret(  # type: ignore
             request=payload_pb2.DeleteSecretRequest(
                 id=definition_pb2.SecretIdentifier(
@@ -106,10 +140,16 @@ class Secret:
 
     @property
     def name(self) -> str:
+        """
+        Get the name of the secret.
+        """
         return self.pb2.id.name
 
     @property
     def type(self) -> str:
+        """
+        Get the type of the secret as a string ("regular" or "image_pull").
+        """
         if self.pb2.secret_metadata.type == definition_pb2.SecretType.SECRET_TYPE_GENERIC:
             return "regular"
         elif self.pb2.secret_metadata.type == definition_pb2.SecretType.SECRET_TYPE_IMAGE_PULL_SECRET:
@@ -117,6 +157,9 @@ class Secret:
         raise ValueError("unknown type")
 
     def __rich_repr__(self) -> rich.repr.Result:
+        """
+        Rich representation of the Secret object for pretty printing.
+        """
         yield "project", self.pb2.id.project or "-"
         yield "domain", self.pb2.id.domain or "-"
         yield "name", self.name