fix: emit regular pipeline heartbeats (HEXA-1406) (#329)

Author: yolanfery

openhexa/graphql/graphql_client/__init__.py

@@ -388,6 +388,10 @@
     UpdateDatasetUpdateDataset,
     UpdateDatasetUpdateDatasetDataset,
 )
+from .update_pipeline_heartbeat import (
+    UpdatePipelineHeartbeat,
+    UpdatePipelineHeartbeatUpdatePipelineHeartbeat,
+)
 from .update_webapp import UpdateWebapp, UpdateWebappUpdateWebapp
 from .update_workspace import (
     UpdateWorkspace,
@@ -737,6 +741,8 @@
     "UpdateOrganizationMemberError",
     "UpdateOrganizationMemberInput",
     "UpdatePipelineError",
+    "UpdatePipelineHeartbeat",
+    "UpdatePipelineHeartbeatUpdatePipelineHeartbeat",
     "UpdatePipelineInput",
     "UpdatePipelineProgressInput",
     "UpdatePipelineRecipientInput",

openhexa/graphql/graphql_client/client.py

@@ -84,6 +84,10 @@
 from .stop_pipeline import StopPipeline, StopPipelineStopPipeline
 from .update_connection import UpdateConnection, UpdateConnectionUpdateConnection
 from .update_dataset import UpdateDataset, UpdateDatasetUpdateDataset
+from .update_pipeline_heartbeat import (
+    UpdatePipelineHeartbeat,
+    UpdatePipelineHeartbeatUpdatePipelineHeartbeat,
+)
 from .update_webapp import UpdateWebapp, UpdateWebappUpdateWebapp
 from .update_workspace import UpdateWorkspace, UpdateWorkspaceUpdateWorkspace
 from .upgrade_pipeline_version_from_template import (
@@ -310,6 +314,29 @@ def stop_pipeline(
         data = self.get_data(response)
         return StopPipeline.model_validate(data).stop_pipeline
 
+    def update_pipeline_heartbeat(
+        self, **kwargs: Any
+    ) -> UpdatePipelineHeartbeatUpdatePipelineHeartbeat:
+        query = gql(
+            """
+            mutation UpdatePipelineHeartbeat {
+              updatePipelineHeartbeat {
+                success
+                errors
+              }
+            }
+            """
+        )
+        variables: Dict[str, object] = {}
+        response = self.execute(
+            query=query,
+            operation_name="UpdatePipelineHeartbeat",
+            variables=variables,
+            **kwargs
+        )
+        data = self.get_data(response)
+        return UpdatePipelineHeartbeat.model_validate(data).update_pipeline_heartbeat
+
     def add_pipeline_recipient(
         self, input: CreatePipelineRecipientInput, **kwargs: Any
     ) -> AddPipelineRecipientAddPipelineRecipient:

openhexa/graphql/graphql_client/create_pipeline_from_template_version.py

@@ -10,9 +10,9 @@
 
 
 class CreatePipelineFromTemplateVersion(BaseModel):
-    create_pipeline_from_template_version: "CreatePipelineFromTemplateVersionCreatePipelineFromTemplateVersion" = Field(
-        alias="createPipelineFromTemplateVersion"
-    )
+    create_pipeline_from_template_version: (
+        "CreatePipelineFromTemplateVersionCreatePipelineFromTemplateVersion"
+    ) = Field(alias="createPipelineFromTemplateVersion")
 
 
 class CreatePipelineFromTemplateVersionCreatePipelineFromTemplateVersion(BaseModel):

openhexa/graphql/graphql_client/create_pipeline_template_version.py

@@ -10,9 +10,9 @@
 
 
 class CreatePipelineTemplateVersion(BaseModel):
-    create_pipeline_template_version: "CreatePipelineTemplateVersionCreatePipelineTemplateVersion" = Field(
-        alias="createPipelineTemplateVersion"
-    )
+    create_pipeline_template_version: (
+        "CreatePipelineTemplateVersionCreatePipelineTemplateVersion"
+    ) = Field(alias="createPipelineTemplateVersion")
 
 
 class CreatePipelineTemplateVersionCreatePipelineTemplateVersion(BaseModel):

openhexa/graphql/graphql_client/enums.py

@@ -750,6 +750,7 @@ class UpdateTeamError(str, Enum):
 
 
 class UpdateTemplateError(str, Enum):
+    INVALID_CONFIG = "INVALID_CONFIG"
     NOT_FOUND = "NOT_FOUND"
     PERMISSION_DENIED = "PERMISSION_DENIED"
 

openhexa/graphql/graphql_client/input_types.py

@@ -147,6 +147,7 @@ class CreatePipelineInput(BaseModel):
     )
     name: str
     notebook_path: Optional[str] = Field(alias="notebookPath", default=None)
+    tags: Optional[List[str]] = None
     workspace_slug: str = Field(alias="workspaceSlug")
 
 
@@ -638,8 +639,12 @@ class UpdateTeamInput(BaseModel):
 class UpdateTemplateInput(BaseModel):
     config: Optional[Any] = None
     description: Optional[str] = None
+    functional_type: Optional[PipelineFunctionalType] = Field(
+        alias="functionalType", default=None
+    )
     id: Any
     name: Optional[str] = None
+    tags: Optional[List[str]] = None
 
 
 class UpdateTemplateVersionInput(BaseModel):

openhexa/graphql/graphql_client/update_pipeline_heartbeat.py

@@ -0,0 +1,23 @@
+# Generated by ariadne-codegen
+# Source: openhexa/graphql/queries.graphql
+
+from typing import List
+
+from pydantic import Field
+
+from .base_model import BaseModel
+from .enums import PipelineError
+
+
+class UpdatePipelineHeartbeat(BaseModel):
+    update_pipeline_heartbeat: "UpdatePipelineHeartbeatUpdatePipelineHeartbeat" = Field(
+        alias="updatePipelineHeartbeat"
+    )
+
+
+class UpdatePipelineHeartbeatUpdatePipelineHeartbeat(BaseModel):
+    success: bool
+    errors: List[PipelineError]
+
+
+UpdatePipelineHeartbeat.model_rebuild()

openhexa/graphql/graphql_client/upgrade_pipeline_version_from_template.py

@@ -10,9 +10,9 @@
 
 
 class UpgradePipelineVersionFromTemplate(BaseModel):
-    upgrade_pipeline_version_from_template: "UpgradePipelineVersionFromTemplateUpgradePipelineVersionFromTemplate" = Field(
-        alias="upgradePipelineVersionFromTemplate"
-    )
+    upgrade_pipeline_version_from_template: (
+        "UpgradePipelineVersionFromTemplateUpgradePipelineVersionFromTemplate"
+    ) = Field(alias="upgradePipelineVersionFromTemplate")
 
 
 class UpgradePipelineVersionFromTemplateUpgradePipelineVersionFromTemplate(BaseModel):

openhexa/graphql/queries.graphql

@@ -112,6 +112,13 @@ mutation StopPipeline($input: StopPipelineInput!) {
   }
 }
 
+mutation UpdatePipelineHeartbeat {
+  updatePipelineHeartbeat {
+    success
+    errors
+  }
+}
+
 mutation AddPipelineRecipient($input: CreatePipelineRecipientInput!) {
   addPipelineRecipient(input: $input) {
     success

openhexa/graphql/schema.generated.graphql

@@ -797,6 +797,7 @@ input CreatePipelineInput {
   functionalType: PipelineFunctionalType
   name: String!
   notebookPath: String
+  tags: [String!]
   workspaceSlug: String!
 }
 
@@ -1863,7 +1864,8 @@ type File {
   path: String!
   size: BigInt
   type: FileType!
-  updated: DateTime
+  updated: DateTime @deprecated(reason: "Use updatedAt instead")
+  updatedAt: DateTime
 }
 
 """Represents a file or directory node in a flattened structure."""
@@ -2574,6 +2576,12 @@ type Mutation {
   """Updates an existing pipeline."""
   updatePipeline(input: UpdatePipelineInput!): UpdatePipelineResult!
 
+  """
+  Updates the heartbeat timestamp for the current pipeline run.
+  Must be called by an authenticated pipeline run.
+  """
+  updatePipelineHeartbeat: UpdatePipelineHeartbeatResult!
+
   """Updates the progress of a pipeline."""
   updatePipelineProgress(input: UpdatePipelineProgressInput!): UpdatePipelineProgressResult!
 
@@ -2625,6 +2633,9 @@ type Organization {
   """The contact information of the organization."""
   contactInfo: String!
 
+  """Dataset links available in the organization"""
+  datasetLinks(page: Int = 1, perPage: Int = 15, query: String): DatasetLinkPage!
+
   """Datasets available in the organization"""
   datasets(page: Int = 1, perPage: Int = 15, query: String): DatasetPage!
 
@@ -2640,9 +2651,20 @@ type Organization {
   """The name of the organization."""
   name: String!
 
+  """
+  The direct invitations sent to join a specific workspace in the organization.
+  """
+  pendingWorkspaceInvitations(page: Int, perPage: Int): WorkspaceInvitationPage!
+
   """The permissions the current user has in the organization."""
   permissions: OrganizationPermissions!
 
+  """Pipeline tags used within this organization."""
+  pipelineTags: [String!]!
+
+  """Pipeline template tags used within this organization."""
+  pipelineTemplateTags: [String!]!
+
   """The short name of the organization."""
   shortName: String
 
@@ -2899,6 +2921,7 @@ type PipelineParameter {
   code: String!
   connection: String
   default: Generic
+  directory: String
   help: String
   multiple: Boolean!
   name: String!
@@ -3018,10 +3041,12 @@ type PipelineTemplate {
   config: String
   currentVersion: PipelineTemplateVersion
   description: String
+  functionalType: PipelineFunctionalType
   id: UUID!
   name: String!
   permissions: PipelineTemplatePermissions!
   sourcePipeline: Pipeline
+  tags: [Tag!]!
   updatedAt: DateTime!
   versions(page: Int, perPage: Int): TemplateVersionPage!
   workspace: Workspace
@@ -3229,7 +3254,7 @@ enum PrepareObjectUploadError {
 
 """
 Input for preparing to upload an object to a workspace's bucket.
-The `contentType`
+The `contentType` 
 """
 input PrepareObjectUploadInput {
   contentType: String
@@ -3332,8 +3357,11 @@ type Query {
   """Retrieves a pipeline run by ID."""
   pipelineRun(id: UUID!): PipelineRun
 
+  """Retrieves a pipeline template version by ID."""
+  pipelineTemplateVersion(id: UUID!): PipelineTemplateVersion
+
   """Retrieves a page of pipeline templates."""
-  pipelineTemplates(page: Int = 1, perPage: Int = 15, search: String, workspaceSlug: String): PipelineTemplatePage!
+  pipelineTemplates(functionalType: PipelineFunctionalType, page: Int = 1, perPage: Int = 15, search: String, tags: [String!], workspaceSlug: String): PipelineTemplatePage!
 
   """Retrieves a pipeline version by ID."""
   pipelineVersion(id: UUID!): PipelineVersion
@@ -3342,7 +3370,7 @@ type Query {
   pipelines(functionalType: PipelineFunctionalType, name: String, page: Int, perPage: Int, search: String, tags: [String!], workspaceSlug: String): PipelinesPage!
   searchDatabaseTables(organizationId: UUID, page: Int = 1, perPage: Int = 15, query: String!, workspaceSlugs: [String]): DatabaseTableResultPage!
   searchDatasets(organizationId: UUID, page: Int = 1, perPage: Int = 15, query: String!, workspaceSlugs: [String]): DatasetResultPage!
-  searchFiles(organizationId: UUID, page: Int = 1, perPage: Int = 15, query: String!, workspaceSlugs: [String]): FileResultPage!
+  searchFiles(organizationId: UUID, page: Int = 1, perPage: Int = 15, prefix: String, query: String!, workspaceSlugs: [String]): FileResultPage!
   searchPipelineTemplates(organizationId: UUID, page: Int = 1, perPage: Int = 15, query: String!, workspaceSlugs: [String]): PipelineTemplateResultPage!
   searchPipelines(functionalType: PipelineFunctionalType, organizationId: UUID, page: Int = 1, perPage: Int = 15, query: String!, workspaceSlugs: [String]): PipelineResultPage!
   team(id: UUID!): Team
@@ -4052,6 +4080,12 @@ enum UpdatePipelineError {
   PERMISSION_DENIED
 }
 
+"""Represents the result of updating a pipeline heartbeat."""
+type UpdatePipelineHeartbeatResult {
+  errors: [PipelineError!]!
+  success: Boolean!
+}
+
 """Represents the input for updating a pipeline."""
 input UpdatePipelineInput {
   autoUpdateFromTemplate: Boolean
@@ -4162,6 +4196,7 @@ type UpdateTeamResult {
 Enum representing the possible errors that can occur when updating a template.
 """
 enum UpdateTemplateError {
+  INVALID_CONFIG
   NOT_FOUND
   PERMISSION_DENIED
 }
@@ -4170,8 +4205,10 @@ enum UpdateTemplateError {
 input UpdateTemplateInput {
   config: JSON
   description: String
+  functionalType: PipelineFunctionalType
   id: UUID!
   name: String
+  tags: [String!]
 }
 
 """Represents the result of updating a template."""
@@ -4482,6 +4519,8 @@ type Workspace {
   name: String!
   organization: Organization
   permissions: WorkspacePermissions!
+  pipelineTags: [String!]!
+  pipelineTemplateTags: [String!]!
   slug: String!
   updatedAt: DateTime
 }
@@ -4531,6 +4570,7 @@ enum WorkspaceInvitationStatus {
 type WorkspaceMembership {
   createdAt: DateTime!
   id: UUID!
+  organizationMembership: OrganizationMembership
   role: WorkspaceMembershipRole!
   updatedAt: DateTime
   user: User!
@@ -4595,4 +4635,4 @@ type WorkspacePermissions {
   launchNotebookServer: Boolean!
   manageMembers: Boolean!
   update: Boolean!
-}
+}