pipekit-sdk 2.0.0__tar.gz → 2.1.0__tar.gz

{pipekit_sdk-2.0.0 → pipekit_sdk-2.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pipekit-sdk
-Version: 2.0.0
+Version: 2.1.0
 Summary: Pipekit Python SDK
 Home-page: https://pipekit.io
 License: BSD-3-Clause
@@ -16,8 +16,10 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: hera (>=5.5.2)
 Requires-Dist: pydantic (<3)
+Requires-Dist: pyjwt (>=2.9.0,<3.0.0)
 Requires-Dist: responses (>=0.25.3,<0.26.0)
 Project-URL: Documentation, https://docs.pipekit.io/
 Description-Content-Type: text/markdown

{pipekit_sdk-2.0.0 → pipekit_sdk-2.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pipekit-sdk"
-version = "2.0.0"
+version = "2.1.0"
 description = "Pipekit Python SDK"
 authors = ["Pipekit <ci@pipekit.io>"]
 license = "BSD-3-Clause"
@@ -27,6 +27,7 @@ python = ">=3.8,<4"
 hera = ">=5.5.2"
 pydantic = "<3"
 responses = "^0.25.3"
+pyjwt = "^2.9.0"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "*"

{pipekit_sdk-2.0.0 → pipekit_sdk-2.1.0}/src/pipekit_sdk/models/__init__.py

@@ -4,8 +4,8 @@ from __future__ import annotations
 
 from typing import Any
 
-from pipekit_sdk.models._config import BaseModel
+from pydantic import RootModel
 
 
-class Model(BaseModel):
-    __root__: Any
+class Model(RootModel[Any]):
+    root: Any

pipekit_sdk-2.1.0/src/pipekit_sdk/models/_config.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel as PydanticBaseModel
+from pydantic import ConfigDict
+
+
+class BaseModel(PydanticBaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,  # support populating object by Field alias
+        use_enum_values=True,  # supports using enums, which are then unpacked to obtain the actual `.value`
+    )

{pipekit_sdk-2.0.0 → pipekit_sdk-2.1.0}/src/pipekit_sdk/models/model.py

@@ -2,10 +2,9 @@
 
 from __future__ import annotations
 
-from datetime import datetime
 from typing import Dict, List, Optional
 
-from pydantic import Field
+from pydantic import AwareDatetime, Field
 from typing_extensions import Annotated
 
 from pipekit_sdk.models._config import BaseModel
@@ -13,7 +12,7 @@ from pipekit_sdk.models._config import BaseModel
 
 class Cluster(BaseModel):
     api_key: Annotated[Optional[str], Field(alias="apiKey")] = None
-    created_at: Annotated[datetime, Field(alias="createdAt")]
+    created_at: Annotated[AwareDatetime, Field(alias="createdAt")]
     description: str
     identity_uuid: Annotated[str, Field(alias="identityUUID")]
     is_active: Annotated[bool, Field(alias="isActive")]
@@ -23,7 +22,7 @@ class Cluster(BaseModel):
     pipekit_agent_version: Annotated[str, Field(alias="pipekitAgentVersion")]
     port: Optional[int] = None
     type: List[str]
-    updated_at: Annotated[datetime, Field(alias="updatedAt")]
+    updated_at: Annotated[AwareDatetime, Field(alias="updatedAt")]
     uuid: str
     workflow_group_priorities: Annotated[Optional[Dict[str, int]], Field(alias="workflowGroupPriorities")] = None
 
@@ -43,7 +42,7 @@ class GitInfo(BaseModel):
 
 class Logs(BaseModel):
     container_name: Annotated[str, Field(alias="containerName")]
-    created_at: Annotated[datetime, Field(alias="createdAt")]
+    created_at: Annotated[AwareDatetime, Field(alias="createdAt")]
     identity_uuid: Annotated[str, Field(alias="identityUUID")]
     log_type: Annotated[str, Field(alias="logType")]
     node_id: Annotated[Optional[str], Field(alias="nodeId")] = None
@@ -52,7 +51,7 @@ class Logs(BaseModel):
     pod_name: Annotated[str, Field(alias="podName")]
     run_uuid: Annotated[str, Field(alias="runUUID")]
     stream_cursor: Annotated[str, Field(alias="streamCursor")]
-    updated_at: Annotated[datetime, Field(alias="updatedAt")]
+    updated_at: Annotated[AwareDatetime, Field(alias="updatedAt")]
 
 
 class Pipe(BaseModel):
@@ -77,7 +76,7 @@ class PipeRun(BaseModel):
     check_id: Annotated[Optional[int], Field(alias="checkId")] = None
     cluster_name: Annotated[Optional[str], Field(alias="clusterName")] = None
     cluster_uuid: Annotated[Optional[str], Field(alias="clusterUUID")] = None
-    created_at: Annotated[Optional[datetime], Field(alias="createdAt")] = None
+    created_at: Annotated[Optional[AwareDatetime], Field(alias="createdAt")] = None
     git_info: Annotated[Optional[GitInfo], Field(alias="gitInfo")] = None
     identity_uuid: Annotated[Optional[str], Field(alias="identityUUID")] = None
     message: Optional[str] = None
@@ -92,6 +91,6 @@ class PipeRun(BaseModel):
     source_name: Annotated[Optional[str], Field(alias="sourceName")] = None
     source_type: Annotated[Optional[str], Field(alias="sourceType")] = None
     status: Optional[str] = None
-    updated_at: Annotated[Optional[datetime], Field(alias="updatedAt")] = None
+    updated_at: Annotated[Optional[AwareDatetime], Field(alias="updatedAt")] = None
     user_name: Annotated[Optional[str], Field(alias="userName")] = None
     uuid: Optional[str] = None

{pipekit_sdk-2.0.0 → pipekit_sdk-2.1.0}/src/pipekit_sdk/models/v1.py

@@ -2,10 +2,9 @@
 
 from __future__ import annotations
 
-from datetime import datetime
 from typing import Any, Dict, List, Optional
 
-from pydantic import Field
+from pydantic import AwareDatetime, Field, RootModel
 from typing_extensions import Annotated
 
 from pipekit_sdk.models._config import BaseModel
@@ -276,8 +275,8 @@ class FCVolumeSource(BaseModel):
     ] = None
 
 
-class FieldsV1(BaseModel):
-    __root__: Annotated[
+class FieldsV1(RootModel[Any]):
+    root: Annotated[
         Any,
         Field(
             description=(
@@ -1317,7 +1316,7 @@ class TCPSocketAction(BaseModel):
 
 
 class Time(BaseModel):
-    time: Annotated[datetime, Field(alias="Time")]
+    time: Annotated[AwareDatetime, Field(alias="Time")]
 
 
 class Toleration(BaseModel):

{pipekit_sdk-2.0.0 → pipekit_sdk-2.1.0}/src/pipekit_sdk/models/v1alpha1.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from typing import Any, Dict, List, Optional
 
-from pydantic import Field
+from pydantic import Field, RootModel
 from typing_extensions import Annotated
 
 from pipekit_sdk.models._config import BaseModel
@@ -12,8 +12,8 @@ from pipekit_sdk.models._config import BaseModel
 from . import intstr, v1
 
 
-class Amount(BaseModel):
-    __root__: Any
+class Amount(RootModel[Any]):
+    root: Any
 
 
 class ArtGCStatus(BaseModel):
@@ -190,8 +190,8 @@ class Histogram(BaseModel):
     value: str
 
 
-class Item(BaseModel):
-    __root__: Any
+class Item(RootModel[Any]):
+    root: Any
 
 
 class LabelValueFrom(BaseModel):
@@ -244,8 +244,8 @@ class NodeSynchronizationStatus(BaseModel):
     waiting: Optional[str] = None
 
 
-class NoneStrategy(BaseModel):
-    __root__: Any
+class NoneStrategy(RootModel[Any]):
+    root: Any
 
 
 class OAuth2EndpointParam(BaseModel):
@@ -258,12 +258,12 @@ class OSSLifecycleRule(BaseModel):
     mark_infrequent_access_after_days: Annotated[Optional[int], Field(alias="markInfrequentAccessAfterDays")] = None
 
 
-class ParallelSteps(BaseModel):
-    __root__: Any
+class ParallelSteps(RootModel[Any]):
+    root: Any
 
 
-class Plugin(BaseModel):
-    __root__: Any
+class Plugin(RootModel[Any]):
+    root: Any
 
 
 class Prometheus(BaseModel):
@@ -280,8 +280,8 @@ class RawArtifact(BaseModel):
     data: str
 
 
-class RetryNodeAntiAffinity(BaseModel):
-    __root__: Any
+class RetryNodeAntiAffinity(RootModel[Any]):
+    root: Any
 
 
 class S3EncryptionOptions(BaseModel):
@@ -315,8 +315,8 @@ class Sequence(BaseModel):
     start: Optional[intstr.IntOrString] = None
 
 
-class SuppliedValueFrom(BaseModel):
-    __root__: Any
+class SuppliedValueFrom(RootModel[Any]):
+    root: Any
 
 
 class SuspendTemplate(BaseModel):
@@ -388,8 +388,8 @@ class WorkflowTemplateRef(BaseModel):
     name: Optional[str] = None
 
 
-class ZipStrategy(BaseModel):
-    __root__: Any
+class ZipStrategy(RootModel[Any]):
+    root: Any
 
 
 class ArchiveStrategy(BaseModel):
@@ -1689,7 +1689,7 @@ class WorkflowStatus(BaseModel):
     ] = None
 
 
-CronWorkflow.update_forward_refs()
-CronWorkflowSpec.update_forward_refs()
-DAGTask.update_forward_refs()
-Workflow.update_forward_refs()
+CronWorkflow.model_rebuild()
+CronWorkflowSpec.model_rebuild()
+DAGTask.model_rebuild()
+Workflow.model_rebuild()

{pipekit_sdk-2.0.0 → pipekit_sdk-2.1.0}/src/pipekit_sdk/service.py

@@ -2,19 +2,71 @@
 
 import json
 import os
-from typing import Any, Generator, Optional, cast
+from datetime import datetime
+from enum import Enum
+from pathlib import Path
+from typing import Any, Generator, List, Optional, cast
 
+import jwt
 import requests
 from hera.workflows import CronWorkflow, Workflow
 from hera.workflows.models import WorkflowCreateRequest
+from pydantic import BaseModel, ValidationError
 
 from pipekit_sdk._helpers import is_valid_uuid
 from pipekit_sdk.models.model import Cluster, Logs, Pipe, PipeRun
 
 
+class _LoginResponse(BaseModel):
+    """LoginResponse is a model for Login responses."""
+
+    token_type: str
+    access_token: str
+    expires_in: int
+    refresh_token: str
+
+
+class _AccessToken(BaseModel):
+    """AccessToken is a model for decoded access tokens."""
+
+    identityUUID: str
+    ClusterUUID: str
+    OrgUUID: str
+    isActive: bool
+    products: List[str]
+    iss: str
+    sub: str
+    exp: int
+
+
+class _ValidationResult(Enum):
+    OK = 0
+    INVALID_TOKEN = 1
+    EXPIRED_TOKEN = 2
+
+
 class PipekitService:
     """PipekitService is a wrapper around the Pipekit API. It provides a simple interface to interact with Pipekit."""
 
+    @staticmethod
+    def __from_home_dir() -> Optional[_LoginResponse]:
+        home = Path.home()
+        pipekit_token_file_path = os.path.join(home, ".pipekit", "token")
+        try:
+            with open(pipekit_token_file_path, "r") as f:
+                raw_data = f.read()
+                return _LoginResponse.model_validate_json(raw_data)
+        except ValidationError:
+            return None
+
+    @staticmethod
+    def __write_home_dir(login_resp: _LoginResponse):
+        home = Path.home()
+        pipekit_token_file_path = os.path.join(home, ".pipekit", "token")
+        with open(pipekit_token_file_path, "w") as f:
+            payload = login_resp.model_dump_json(indent=2)
+            f.write(payload)
+
     def __init__(
         self,
         username: str = "",
@@ -38,8 +90,60 @@ class PipekitService:
             self.access_token = token
             return
 
+        maybe_resp = self.__from_home_dir()
+        self.token_data = maybe_resp
+
+        if self.token_data is not None:
+            self.__token_login_flow()
+            return
+
         self.__login()
 
+    def __update_token(self) -> None:
+        login_url = f"{self.id_url}/api/id/v1/oauth/token"
+        assert self.token_data is not None, "Token data was None, this implies a bug"
+        login_request = {"grant_type": "refresh_token", "refresh_token": self.token_data.refresh_token}
+        login_response = requests.post(
+            login_url,
+            data=json.dumps(login_request),
+            headers={"Content-Type": "application/json"},
+            timeout=10,
+        )
+        if login_response.status_code // 100 != 2:
+            raise Exception(
+                "Couldn't refresh access token, obtained non 2xx status code, try logging in again via the cli"
+            )
+
+        login_response_model = _LoginResponse.model_validate(login_response.json())
+        self.__write_home_dir(login_response_model)
+        self.access_token = login_response_model.access_token
+
+    def __validate_token(self) -> _ValidationResult:
+        assert self.token_data is not None
+        try:
+            decoded_payload = jwt.decode(self.token_data.access_token, options={"verify_signature": False})
+            acc_tk = _AccessToken.model_validate(decoded_payload)
+            exp_time = datetime.fromtimestamp(acc_tk.exp)
+            now_time = datetime.now()
+
+            if exp_time <= now_time:
+                return _ValidationResult.EXPIRED_TOKEN
+        except jwt.PyJWTError:
+            return _ValidationResult.INVALID_TOKEN
+        return _ValidationResult.OK
+
+    def __token_login_flow(self) -> None:
+        assert self.token_data is not None, "token was None when it was expected to be valid"
+        res = self.__validate_token()
+        if res == _ValidationResult.EXPIRED_TOKEN:
+            self.__update_token()
+        elif res == _ValidationResult.INVALID_TOKEN:
+            raise RuntimeError(
+                "Token was invalid, if the `.pipekit/token` file is corrupted, please delete the file and login again"
+            )
+        else:
+            self.access_token = self.token_data.access_token
+
     def __login(self) -> None:
         login_request = {
             "username": self.username,
@@ -80,7 +184,7 @@ class PipekitService:
         if clusters_response.status_code // 100 != 2:
             raise Exception(f"List clusters failed: {clusters_response.text}")
 
-        clusters = [Cluster.parse_obj(cluster) for cluster in clusters_response.json()]
+        clusters = [Cluster.model_validate(cluster) for cluster in clusters_response.json()]
         return clusters
 
     def list_pipes(
@@ -107,7 +211,7 @@ class PipekitService:
             if pipes_response.status_code // 100 != 2:
                 raise Exception(f"List pipes failed: {pipes_response.text}")
 
-            pipes_page = [Pipe.parse_obj(item) for item in pipes_response.json()["items"]]
+            pipes_page = [Pipe.model_validate(item) for item in pipes_response.json()["items"]]
             if len(pipes_page) == 0:
                 break
 
@@ -135,7 +239,7 @@ class PipekitService:
         if pipe_response.status_code // 100 != 2:
             raise Exception(f"Get pipe failed: {pipe_response.text}")
 
-        return Pipe.parse_obj(pipe_response.json())
+        return Pipe.model_validate(pipe_response.json())
 
     def list_runs_by_cluster(
         self,
@@ -160,7 +264,7 @@ class PipekitService:
             if runs_response.status_code // 100 != 2:
                 raise Exception(f"List runs failed: {runs_response.text}")
 
-            runs_page = [PipeRun.parse_obj(item) for item in runs_response.json()["items"]]
+            runs_page = [PipeRun.model_validate(item) for item in runs_response.json()["items"]]
             if len(runs_page) == 0:
                 break
 
@@ -196,7 +300,7 @@ class PipekitService:
             if runs_response.status_code // 100 != 2:
                 raise Exception(f"List runs failed: {runs_response.text}")
 
-            runs_page = [PipeRun.parse_obj(item) for item in runs_response.json()["items"]]
+            runs_page = [PipeRun.model_validate(item) for item in runs_response.json()["items"]]
             if len(runs_page) == 0:
                 break
 
@@ -224,7 +328,7 @@ class PipekitService:
         if run_response.status_code // 100 != 2:
             raise Exception(f"Get run failed: {run_response.text}")
 
-        return PipeRun.parse_obj(run_response.json())
+        return PipeRun.model_validate(run_response.json())
 
     def get_cluster(
         self,
@@ -241,7 +345,7 @@ class PipekitService:
         if cluster_response.status_code // 100 != 2:
             raise Exception(f"Get cluster failed: {cluster_response.text}")
 
-        return Cluster.parse_obj(cluster_response.json())
+        return Cluster.model_validate(cluster_response.json())
 
     def get_cluster_by_name(
         self,
@@ -293,7 +397,7 @@ class PipekitService:
         if submit_response.status_code // 100 != 2:
             raise Exception(f"Submit failed: {submit_response.text}")
 
-        return PipeRun.parse_obj(submit_response.json())
+        return PipeRun.model_validate(submit_response.json())
 
     def create(self, cron_workflow: CronWorkflow, cluster: str, pipe_name: str = "") -> PipeRun:
         """Create a CronWorkflow on the cluster, under an existing or new pipe with the given name."""
@@ -328,7 +432,7 @@ class PipekitService:
         if create_response.status_code // 100 != 2:
             raise Exception(f"Create failed: {create_response.text}")
 
-        return PipeRun.parse_obj(create_response.json())
+        return PipeRun.model_validate(create_response.json())
 
     def __apply_run_action(
         self,
@@ -348,7 +452,7 @@ class PipekitService:
         if action_response.status_code // 100 != 2:
             raise Exception(f"Action failed: {action_response.text}")
 
-        return PipeRun.parse_obj(action_response.json())
+        return PipeRun.model_validate(action_response.json())
 
     def restart(
         self,
@@ -388,7 +492,7 @@ class PipekitService:
         if logs_response.status_code // 100 != 2:
             raise Exception(f"Get logs failed: {logs_response.text}")
 
-        return [Logs.parse_obj(log) for log in list(logs_response.json())]
+        return [Logs.model_validate(log) for log in list(logs_response.json())]
 
     def __get_logs_stream(
         self,
@@ -464,7 +568,7 @@ class PipekitService:
                         # bye
                         return
                     else:
-                        yield Logs.parse_obj(json.loads(line))
+                        yield Logs.model_validate(json.loads(line))
 
             except requests.exceptions.Timeout:
                 pass  # we'll ignore timeout errors and reconnect
@@ -486,3 +590,7 @@ class PipekitService:
 
         for log in self.get_logs(run_uuid, pod_name=pod_name, container_name=container_name):
             self.__print_log_message(log)
+
+
+if __name__ == "__main__":
+    p = PipekitService()

pipekit_sdk-2.0.0/src/pipekit_sdk/models/_config.py

@@ -1,15 +0,0 @@
-from pydantic.v1 import BaseModel as PydanticBaseModel
-
-
-class BaseModel(PydanticBaseModel):
-    class Config:
-        """General config for generated models."""
-
-        allow_population_by_field_name = True
-        """support populating object by Field alias"""
-
-        allow_mutation = True
-        """allow mutation of objects post instantiation"""
-
-        use_enum_values = True
-        """supports using enums, which are then unpacked to obtain the actual `.value`"""