feat: add a "country" property to the current_workspace

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",

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)

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()

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
+        }
+    }
 }

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."""

pyproject.toml

@@ -97,7 +97,7 @@ ignore = [
     "ANN202",
     "ANN204",
     "ANN205",
-    "ANN401",
+    "ANN401"
 ]
 
 [tool.ruff.pycodestyle]