# Generated by ariadne-codegen
# Source: openhexa/graphql/queries.graphql

from typing import Any, Dict, List, Optional, Union

from .add_pipeline_recipient import (
    AddPipelineRecipient,
    AddPipelineRecipientAddPipelineRecipient,
)
from .add_webapp_to_favorites import (
    AddWebappToFavorites,
    AddWebappToFavoritesAddToFavorites,
)
from .archive_workspace import ArchiveWorkspace, ArchiveWorkspaceArchiveWorkspace
from .base_client import BaseClient
from .base_model import UNSET, UnsetType
from .create_connection import CreateConnection, CreateConnectionCreateConnection
from .create_dataset import CreateDataset, CreateDatasetCreateDataset
from .create_pipeline import CreatePipeline, CreatePipelineCreatePipeline
from .create_pipeline_from_template_version import (
    CreatePipelineFromTemplateVersion,
    CreatePipelineFromTemplateVersionCreatePipelineFromTemplateVersion,
)
from .create_pipeline_template_version import (
    CreatePipelineTemplateVersion,
    CreatePipelineTemplateVersionCreatePipelineTemplateVersion,
)
from .create_webapp import CreateWebapp, CreateWebappCreateWebapp
from .create_workspace import CreateWorkspace, CreateWorkspaceCreateWorkspace
from .dataset import Dataset, DatasetDataset
from .datasets import Datasets, DatasetsDatasets
from .delete_connection import DeleteConnection, DeleteConnectionDeleteConnection
from .delete_dataset import DeleteDataset, DeleteDatasetDeleteDataset
from .delete_pipeline import DeletePipeline, DeletePipelineDeletePipeline
from .delete_pipeline_template import (
    DeletePipelineTemplate,
    DeletePipelineTemplateDeletePipelineTemplate,
)
from .delete_pipeline_version import (
    DeletePipelineVersion,
    DeletePipelineVersionDeletePipelineVersion,
)
from .delete_webapp import DeleteWebapp, DeleteWebappDeleteWebapp
from .get_connection import GetConnection, GetConnectionConnectionBySlug
from .get_users import GetUsers, GetUsersUsers
from .input_types import (
    AddToFavoritesInput,
    ArchiveWorkspaceInput,
    CreateConnectionInput,
    CreateDatasetInput,
    CreatePipelineFromTemplateVersionInput,
    CreatePipelineInput,
    CreatePipelineRecipientInput,
    CreatePipelineTemplateVersionInput,
    CreateWebappInput,
    CreateWorkspaceInput,
    DeleteConnectionInput,
    DeleteDatasetInput,
    DeletePipelineInput,
    DeletePipelineTemplateInput,
    DeletePipelineVersionInput,
    DeleteWebappInput,
    InviteWorkspaceMemberInput,
    RemoveFromFavoritesInput,
    StopPipelineInput,
    UpdateConnectionInput,
    UpdateDatasetInput,
    UpdateWebappInput,
    UpdateWorkspaceInput,
    UpgradePipelineVersionFromTemplateInput,
    UploadPipelineInput,
)
from .invite_workspace_member import (
    InviteWorkspaceMember,
    InviteWorkspaceMemberInviteWorkspaceMember,
)
from .pipeline import Pipeline, PipelinePipelineByCode
from .pipelines import Pipelines, PipelinesPipelines
from .remove_webapp_from_favorites import (
    RemoveWebappFromFavorites,
    RemoveWebappFromFavoritesRemoveFromFavorites,
)
from .stop_pipeline import StopPipeline, StopPipelineStopPipeline
from .update_connection import UpdateConnection, UpdateConnectionUpdateConnection
from .update_dataset import UpdateDataset, UpdateDatasetUpdateDataset
from .update_webapp import UpdateWebapp, UpdateWebappUpdateWebapp
from .update_workspace import UpdateWorkspace, UpdateWorkspaceUpdateWorkspace
from .upgrade_pipeline_version_from_template import (
    UpgradePipelineVersionFromTemplate,
    UpgradePipelineVersionFromTemplateUpgradePipelineVersionFromTemplate,
)
from .upload_pipeline import UploadPipeline, UploadPipelineUploadPipeline
from .workspace import Workspace, WorkspaceWorkspace
from .workspaces import Workspaces, WorkspacesWorkspaces


def gql(q: str) -> str:
    return q


class Client(BaseClient):
    def pipelines(
        self,
        workspace_slug: str,
        name: Union[Optional[str], UnsetType] = UNSET,
        page: Union[Optional[int], UnsetType] = UNSET,
        per_page: Union[Optional[int], UnsetType] = UNSET,
        **kwargs: Any
    ) -> PipelinesPipelines:
        query = gql(
            """
            query Pipelines($workspaceSlug: String!, $name: String, $page: Int = 1, $perPage: Int = 10) {
              pipelines(
                workspaceSlug: $workspaceSlug
                name: $name
                page: $page
                perPage: $perPage
              ) {
                totalPages
                items {
                  id
                  code
                  name
                  type
                  currentVersion {
                    id
                    name
                    versionNumber
                  }
                  lastRuns: runs(orderBy: EXECUTION_DATE_DESC, page: 1, perPage: 1) {
                    items {
                      id
                      status
                      executionDate
                    }
                  }
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "workspaceSlug": workspace_slug,
            "name": name,
            "page": page,
            "perPage": per_page,
        }
        response = self.execute(
            query=query, operation_name="Pipelines", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return Pipelines.model_validate(data).pipelines

    def pipeline(
        self, workspace_slug: str, pipeline_code: str, **kwargs: Any
    ) -> Optional[PipelinePipelineByCode]:
        query = gql(
            """
            query Pipeline($workspaceSlug: String!, $pipelineCode: String!) {
              pipelineByCode(workspaceSlug: $workspaceSlug, code: $pipelineCode) {
                id
                code
                name
                description
                type
                currentVersion {
                  id
                  name
                  versionNumber
                  user {
                    id
                    email
                    displayName
                    avatar {
                      initials
                      color
                    }
                  }
                  createdAt
                }
                runs(orderBy: EXECUTION_DATE_DESC, page: 1, perPage: 10) {
                  items {
                    id
                    status
                    executionDate
                    user {
                      id
                      displayName
                      email
                    }
                  }
                }
                webhookUrl
                webhookEnabled
                schedule
                notebookPath
                sourceTemplate {
                  id
                  code
                  name
                }
                hasNewTemplateVersions
                newTemplateVersions {
                  id
                  changelog
                  versionNumber
                  createdAt
                }
                recipients {
                  user {
                    id
                    displayName
                  }
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "workspaceSlug": workspace_slug,
            "pipelineCode": pipeline_code,
        }
        response = self.execute(
            query=query, operation_name="Pipeline", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return Pipeline.model_validate(data).pipeline_by_code

    def create_pipeline(
        self, input: CreatePipelineInput, **kwargs: Any
    ) -> CreatePipelineCreatePipeline:
        query = gql(
            """
            mutation CreatePipeline($input: CreatePipelineInput!) {
              createPipeline(input: $input) {
                success
                errors
                pipeline {
                  code
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="CreatePipeline", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return CreatePipeline.model_validate(data).create_pipeline

    def upload_pipeline(
        self, input: UploadPipelineInput, **kwargs: Any
    ) -> UploadPipelineUploadPipeline:
        query = gql(
            """
            mutation uploadPipeline($input: UploadPipelineInput!) {
              uploadPipeline(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="uploadPipeline", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return UploadPipeline.model_validate(data).upload_pipeline

    def delete_pipeline(
        self, input: DeletePipelineInput, **kwargs: Any
    ) -> DeletePipelineDeletePipeline:
        query = gql(
            """
            mutation DeletePipeline($input: DeletePipelineInput!) {
              deletePipeline(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="DeletePipeline", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return DeletePipeline.model_validate(data).delete_pipeline

    def stop_pipeline(
        self, input: StopPipelineInput, **kwargs: Any
    ) -> StopPipelineStopPipeline:
        query = gql(
            """
            mutation StopPipeline($input: StopPipelineInput!) {
              stopPipeline(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="StopPipeline", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return StopPipeline.model_validate(data).stop_pipeline

    def add_pipeline_recipient(
        self, input: CreatePipelineRecipientInput, **kwargs: Any
    ) -> AddPipelineRecipientAddPipelineRecipient:
        query = gql(
            """
            mutation AddPipelineRecipient($input: CreatePipelineRecipientInput!) {
              addPipelineRecipient(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="AddPipelineRecipient",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return AddPipelineRecipient.model_validate(data).add_pipeline_recipient

    def delete_pipeline_version(
        self, input: DeletePipelineVersionInput, **kwargs: Any
    ) -> DeletePipelineVersionDeletePipelineVersion:
        query = gql(
            """
            mutation DeletePipelineVersion($input: DeletePipelineVersionInput!) {
              deletePipelineVersion(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="DeletePipelineVersion",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return DeletePipelineVersion.model_validate(data).delete_pipeline_version

    def create_pipeline_template_version(
        self, input: CreatePipelineTemplateVersionInput, **kwargs: Any
    ) -> CreatePipelineTemplateVersionCreatePipelineTemplateVersion:
        query = gql(
            """
            mutation CreatePipelineTemplateVersion($input: CreatePipelineTemplateVersionInput!) {
              createPipelineTemplateVersion(input: $input) {
                success
                errors
                pipelineTemplate {
                  name
                  code
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="CreatePipelineTemplateVersion",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return CreatePipelineTemplateVersion.model_validate(
            data
        ).create_pipeline_template_version

    def delete_pipeline_template(
        self, input: DeletePipelineTemplateInput, **kwargs: Any
    ) -> DeletePipelineTemplateDeletePipelineTemplate:
        query = gql(
            """
            mutation DeletePipelineTemplate($input: DeletePipelineTemplateInput!) {
              deletePipelineTemplate(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="DeletePipelineTemplate",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return DeletePipelineTemplate.model_validate(data).delete_pipeline_template

    def create_pipeline_from_template_version(
        self, input: CreatePipelineFromTemplateVersionInput, **kwargs: Any
    ) -> CreatePipelineFromTemplateVersionCreatePipelineFromTemplateVersion:
        query = gql(
            """
            mutation CreatePipelineFromTemplateVersion($input: CreatePipelineFromTemplateVersionInput!) {
              createPipelineFromTemplateVersion(input: $input) {
                success
                errors
                pipeline {
                  id
                  name
                  code
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="CreatePipelineFromTemplateVersion",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return CreatePipelineFromTemplateVersion.model_validate(
            data
        ).create_pipeline_from_template_version

    def upgrade_pipeline_version_from_template(
        self, input: UpgradePipelineVersionFromTemplateInput, **kwargs: Any
    ) -> UpgradePipelineVersionFromTemplateUpgradePipelineVersionFromTemplate:
        query = gql(
            """
            mutation UpgradePipelineVersionFromTemplate($input: UpgradePipelineVersionFromTemplateInput!) {
              upgradePipelineVersionFromTemplate(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="UpgradePipelineVersionFromTemplate",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return UpgradePipelineVersionFromTemplate.model_validate(
            data
        ).upgrade_pipeline_version_from_template

    def update_webapp(
        self, input: UpdateWebappInput, **kwargs: Any
    ) -> UpdateWebappUpdateWebapp:
        query = gql(
            """
            mutation UpdateWebapp($input: UpdateWebappInput!) {
              updateWebapp(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="UpdateWebapp", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return UpdateWebapp.model_validate(data).update_webapp

    def create_webapp(
        self, input: CreateWebappInput, **kwargs: Any
    ) -> CreateWebappCreateWebapp:
        query = gql(
            """
            mutation CreateWebapp($input: CreateWebappInput!) {
              createWebapp(input: $input) {
                success
                errors
                webapp {
                  id
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="CreateWebapp", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return CreateWebapp.model_validate(data).create_webapp

    def delete_webapp(
        self, input: DeleteWebappInput, **kwargs: Any
    ) -> DeleteWebappDeleteWebapp:
        query = gql(
            """
            mutation DeleteWebapp($input: DeleteWebappInput!) {
              deleteWebapp(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="DeleteWebapp", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return DeleteWebapp.model_validate(data).delete_webapp

    def add_webapp_to_favorites(
        self, input: AddToFavoritesInput, **kwargs: Any
    ) -> AddWebappToFavoritesAddToFavorites:
        query = gql(
            """
            mutation AddWebappToFavorites($input: AddToFavoritesInput!) {
              addToFavorites(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="AddWebappToFavorites",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return AddWebappToFavorites.model_validate(data).add_to_favorites

    def remove_webapp_from_favorites(
        self, input: RemoveFromFavoritesInput, **kwargs: Any
    ) -> RemoveWebappFromFavoritesRemoveFromFavorites:
        query = gql(
            """
            mutation RemoveWebappFromFavorites($input: RemoveFromFavoritesInput!) {
              removeFromFavorites(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="RemoveWebappFromFavorites",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return RemoveWebappFromFavorites.model_validate(data).remove_from_favorites

    def workspace(self, slug: str, **kwargs: Any) -> Optional[WorkspaceWorkspace]:
        query = gql(
            """
            query Workspace($slug: String!) {
              workspace(slug: $slug) {
                slug
                name
                description
                dockerImage
                configuration
                countries {
                  code
                  flag
                  name
                  alpha3
                }
                permissions {
                  delete
                  update
                  manageMembers
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"slug": slug}
        response = self.execute(
            query=query, operation_name="Workspace", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return Workspace.model_validate(data).workspace

    def workspaces(
        self,
        page: Union[Optional[int], UnsetType] = UNSET,
        per_page: Union[Optional[int], UnsetType] = UNSET,
        **kwargs: Any
    ) -> WorkspacesWorkspaces:
        query = gql(
            """
            query Workspaces($page: Int = 1, $perPage: Int = 10) {
              workspaces(page: $page, perPage: $perPage) {
                totalPages
                items {
                  slug
                  name
                  description
                  countries {
                    code
                    alpha3
                    name
                  }
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"page": page, "perPage": per_page}
        response = self.execute(
            query=query, operation_name="Workspaces", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return Workspaces.model_validate(data).workspaces

    def create_workspace(
        self, input: CreateWorkspaceInput, **kwargs: Any
    ) -> CreateWorkspaceCreateWorkspace:
        query = gql(
            """
            mutation CreateWorkspace($input: CreateWorkspaceInput!) {
              createWorkspace(input: $input) {
                success
                workspace {
                  slug
                  name
                  description
                  countries {
                    code
                    alpha3
                    name
                  }
                }
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="CreateWorkspace", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return CreateWorkspace.model_validate(data).create_workspace

    def update_workspace(
        self, input: UpdateWorkspaceInput, **kwargs: Any
    ) -> UpdateWorkspaceUpdateWorkspace:
        query = gql(
            """
            mutation UpdateWorkspace($input: UpdateWorkspaceInput!) {
              updateWorkspace(input: $input) {
                success
                workspace {
                  slug
                  name
                  description
                  configuration
                  countries {
                    code
                    alpha3
                    name
                  }
                }
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="UpdateWorkspace", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return UpdateWorkspace.model_validate(data).update_workspace

    def archive_workspace(
        self, input: ArchiveWorkspaceInput, **kwargs: Any
    ) -> ArchiveWorkspaceArchiveWorkspace:
        query = gql(
            """
            mutation ArchiveWorkspace($input: ArchiveWorkspaceInput!) {
              archiveWorkspace(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="ArchiveWorkspace",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return ArchiveWorkspace.model_validate(data).archive_workspace

    def invite_workspace_member(
        self, input: InviteWorkspaceMemberInput, **kwargs: Any
    ) -> InviteWorkspaceMemberInviteWorkspaceMember:
        query = gql(
            """
            mutation InviteWorkspaceMember($input: InviteWorkspaceMemberInput!) {
              inviteWorkspaceMember(input: $input) {
                success
                errors
                workspaceMembership {
                  id
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="InviteWorkspaceMember",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return InviteWorkspaceMember.model_validate(data).invite_workspace_member

    def create_connection(
        self, input: CreateConnectionInput, **kwargs: Any
    ) -> CreateConnectionCreateConnection:
        query = gql(
            """
            mutation CreateConnection($input: CreateConnectionInput!) {
              createConnection(input: $input) {
                success
                connection {
                  __typename
                  id
                  name
                }
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="CreateConnection",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return CreateConnection.model_validate(data).create_connection

    def update_connection(
        self, input: UpdateConnectionInput, **kwargs: Any
    ) -> UpdateConnectionUpdateConnection:
        query = gql(
            """
            mutation UpdateConnection($input: UpdateConnectionInput!) {
              updateConnection(input: $input) {
                success
                errors
                connection {
                  __typename
                  id
                  name
                  slug
                  description
                  fields {
                    code
                    value
                    secret
                  }
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="UpdateConnection",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return UpdateConnection.model_validate(data).update_connection

    def delete_connection(
        self, input: DeleteConnectionInput, **kwargs: Any
    ) -> DeleteConnectionDeleteConnection:
        query = gql(
            """
            mutation DeleteConnection($input: DeleteConnectionInput!) {
              deleteConnection(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query,
            operation_name="DeleteConnection",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return DeleteConnection.model_validate(data).delete_connection

    def get_users(
        self, query: str, workspace_slug: str, **kwargs: Any
    ) -> List[GetUsersUsers]:
        _query = gql(
            """
            query GetUsers($query: String!, $workspaceSlug: String!) {
              users(query: $query, workspaceSlug: $workspaceSlug) {
                id
                email
                displayName
                avatar {
                  initials
                  color
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"query": query, "workspaceSlug": workspace_slug}
        response = self.execute(
            query=_query, operation_name="GetUsers", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return GetUsers.model_validate(data).users

    def datasets(
        self,
        query: Union[Optional[str], UnsetType] = UNSET,
        page: Union[Optional[int], UnsetType] = UNSET,
        per_page: Union[Optional[int], UnsetType] = UNSET,
        **kwargs: Any
    ) -> DatasetsDatasets:
        _query = gql(
            """
            query Datasets($query: String, $page: Int = 1, $perPage: Int = 15) {
              datasets(query: $query, page: $page, perPage: $perPage) {
                totalPages
                items {
                  id
                  slug
                  name
                  description
                  createdAt
                  updatedAt
                  createdBy {
                    id
                    displayName
                  }
                  permissions {
                    update
                    delete
                  }
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "query": query,
            "page": page,
            "perPage": per_page,
        }
        response = self.execute(
            query=_query, operation_name="Datasets", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return Datasets.model_validate(data).datasets

    def dataset(self, id: str, **kwargs: Any) -> Optional[DatasetDataset]:
        query = gql(
            """
            query Dataset($id: ID!) {
              dataset(id: $id) {
                id
                slug
                name
                description
                createdAt
                updatedAt
                createdBy {
                  id
                  displayName
                  email
                }
                permissions {
                  update
                  delete
                  createVersion
                }
                versions {
                  items {
                    id
                    name
                    changelog
                    createdAt
                    createdBy {
                      id
                      displayName
                    }
                  }
                }
                workspace {
                  slug
                  name
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"id": id}
        response = self.execute(
            query=query, operation_name="Dataset", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return Dataset.model_validate(data).dataset

    def create_dataset(
        self, input: CreateDatasetInput, **kwargs: Any
    ) -> CreateDatasetCreateDataset:
        query = gql(
            """
            mutation CreateDataset($input: CreateDatasetInput!) {
              createDataset(input: $input) {
                success
                errors
                dataset {
                  slug
                  name
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="CreateDataset", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return CreateDataset.model_validate(data).create_dataset

    def update_dataset(
        self, input: UpdateDatasetInput, **kwargs: Any
    ) -> UpdateDatasetUpdateDataset:
        query = gql(
            """
            mutation UpdateDataset($input: UpdateDatasetInput!) {
              updateDataset(input: $input) {
                success
                errors
                dataset {
                  slug
                  name
                  description
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="UpdateDataset", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return UpdateDataset.model_validate(data).update_dataset

    def delete_dataset(
        self, input: DeleteDatasetInput, **kwargs: Any
    ) -> DeleteDatasetDeleteDataset:
        query = gql(
            """
            mutation DeleteDataset($input: DeleteDatasetInput!) {
              deleteDataset(input: $input) {
                success
                errors
              }
            }
            """
        )
        variables: Dict[str, object] = {"input": input}
        response = self.execute(
            query=query, operation_name="DeleteDataset", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return DeleteDataset.model_validate(data).delete_dataset

    def get_connection(
        self, workspace_slug: str, connection_slug: str, **kwargs: Any
    ) -> Optional[GetConnectionConnectionBySlug]:
        query = gql(
            """
            query getConnection($workspaceSlug: String!, $connectionSlug: String!) {
              connectionBySlug(workspaceSlug: $workspaceSlug, connectionSlug: $connectionSlug) {
                __typename
                type
                fields {
                  code
                  value
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "workspaceSlug": workspace_slug,
            "connectionSlug": connection_slug,
        }
        response = self.execute(
            query=query, operation_name="getConnection", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return GetConnection.model_validate(data).connection_by_slug
