diff --git a/openhexa/cli/graphql/graphql_client/__init__.py b/openhexa/cli/graphql/graphql_client/__init__.py index a82fb51c..f4320cb0 100644 --- a/openhexa/cli/graphql/graphql_client/__init__.py +++ b/openhexa/cli/graphql/graphql_client/__init__.py @@ -134,6 +134,11 @@ GraphQLClientHttpError, GraphQLClientInvalidResponseError, ) +from .get_countries import ( + GetCountries, + GetCountriesWorkspace, + GetCountriesWorkspaceCountries, +) from .get_workspace_pipelines import ( GetWorkspacePipelines, GetWorkspacePipelinesPipelines, @@ -377,6 +382,9 @@ "GeneratePipelineWebhookUrlInput", "GenerateWorkspaceTokenError", "GenerateWorkspaceTokenInput", + "GetCountries", + "GetCountriesWorkspace", + "GetCountriesWorkspaceCountries", "GetWorkspacePipelines", "GetWorkspacePipelinesPipelines", "GetWorkspacePipelinesPipelinesItems", diff --git a/openhexa/cli/graphql/graphql_client/client.py b/openhexa/cli/graphql/graphql_client/client.py index c5f97d2d..1bb62ae1 100644 --- a/openhexa/cli/graphql/graphql_client/client.py +++ b/openhexa/cli/graphql/graphql_client/client.py @@ -5,6 +5,7 @@ from .base_client import BaseClient from .base_model import UNSET, UnsetType +from .get_countries import GetCountries from .get_workspace_pipelines import GetWorkspacePipelines @@ -60,3 +61,25 @@ def get_workspace_pipelines( ) data = self.get_data(response) return GetWorkspacePipelines.model_validate(data) + + def get_countries(self, workspace_slug: str, **kwargs: Any) -> GetCountries: + query = gql( + """ + query getCountries($workspaceSlug: String!) { + workspace(slug: $workspaceSlug) { + countries { + code + name + flag + alpha3 + } + } + } + """ + ) + variables: Dict[str, object] = {"workspaceSlug": workspace_slug} + response = self.execute( + query=query, operation_name="getCountries", variables=variables, **kwargs + ) + data = self.get_data(response) + return GetCountries.model_validate(data) diff --git a/openhexa/cli/graphql/graphql_client/get_countries.py b/openhexa/cli/graphql/graphql_client/get_countries.py new file mode 100644 index 00000000..2292c773 --- /dev/null +++ b/openhexa/cli/graphql/graphql_client/get_countries.py @@ -0,0 +1,27 @@ +# Generated by ariadne-codegen +# Source: openhexa/cli/graphql/queries.graphql + +from typing import List, Optional + +from pydantic import Field + +from .base_model import BaseModel + + +class GetCountries(BaseModel): + workspace: Optional["GetCountriesWorkspace"] + + +class GetCountriesWorkspace(BaseModel): + countries: List["GetCountriesWorkspaceCountries"] + + +class GetCountriesWorkspaceCountries(BaseModel): + code: str + name: str + flag: str + alpha_3: str = Field(alias="alpha3") + + +GetCountries.model_rebuild() +GetCountriesWorkspace.model_rebuild() diff --git a/openhexa/cli/graphql/queries.graphql b/openhexa/cli/graphql/queries.graphql index e49b0cea..03be1b88 100644 --- a/openhexa/cli/graphql/queries.graphql +++ b/openhexa/cli/graphql/queries.graphql @@ -13,4 +13,15 @@ query getWorkspacePipelines($workspaceSlug: String!, $name: String, $page: Int = } } } +} + +query getCountries($workspaceSlug:String!){ + workspace(slug: $workspaceSlug) { + countries { + code + name + flag + alpha3 + } + } } \ No newline at end of file diff --git a/openhexa/sdk/workspaces/current_workspace.py b/openhexa/sdk/workspaces/current_workspace.py index 0dc6e590..189edbab 100644 --- a/openhexa/sdk/workspaces/current_workspace.py +++ b/openhexa/sdk/workspaces/current_workspace.py @@ -9,6 +9,7 @@ from openhexa.utils import stringcase +from ...cli.graphql.graphql_client import GetCountriesWorkspaceCountries from ..datasets import Dataset from ..utils import graphql from .connection import ( @@ -59,6 +60,16 @@ def slug(self) -> str: except KeyError: raise WorkspaceConfigError("The workspace slug is not available in this environment.") + @property + def countries(self) -> list[GetCountriesWorkspaceCountries]: + """The countries of the workspace.""" + from openhexa.cli.api import OpenHexaClient + + try: + return OpenHexaClient().get_countries(workspace_slug=self.slug).workspace.countries + except KeyError: + raise WorkspaceConfigError("The workspace countries are not available in this environment.") + @property def database_host(self) -> str: """The workspace database host.""" diff --git a/pyproject.toml b/pyproject.toml index baffc7da..4763a6ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,7 +97,7 @@ ignore = [ "ANN202", "ANN204", "ANN205", - "ANN401", + "ANN401" ] [tool.ruff.pycodestyle]