fix: use codegen

Author: nazarfil

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

@@ -20,6 +20,8 @@
     PostgreSQLConnection,
     S3Connection,
 )
+from ...cli.api import OpenHexaClient
+
 
 class Country:
     """Represents a country with its code, name, alpha3 code and flag."""
@@ -32,6 +34,8 @@ def __init__(self, code: str, name: str, alpha3: str, flag: str):
 
     def __repr__(self):
         return f"Country(code={self.code}, name={self.name}, alpha3={self.alpha3}, flag={self.flag})"
+
+
 class WorkspaceConfigError(Exception):
     """Raised whenever the system cannot find an environment variable required to configure the current workspace."""
 
@@ -70,46 +74,13 @@ def slug(self) -> str:
             raise WorkspaceConfigError("The workspace slug is not available in this environment.")
 
     @property
-    def country(self) -> Country:
-        """The country of the workspace."""
+    def countries(self) -> list[Country]:
+        """The countries of the workspace."""
         try:
-            if self._connected:
-                response = graphql(
-                    """
-                    query getWorkspaceCountry($slug: String!) {
-                        workspace(slug: $slug) {
-                            countries {
-                                code
-                                name
-                                alpha3
-                                flag
-                            }
-                        }
-                    }
-                    """,
-                    {"slug": self.slug},
-                )
-                countries = response["workspace"]["countries"]
-                if not countries:
-                    raise WorkspaceConfigError("The workspace does not have a country configured.")
-                if len(response["workspace"]["countries"]) > 1:
-                    warn(
-                        "The workspace has multiple countries configured. The first one will be used.",
-                        UserWarning,
-                        stacklevel=2,
-                    )
-                # return the first country
-                first_country = response["workspace"]["countries"][0]
-                return Country(
-                    code=first_country["code"],
-                    name=first_country["name"],
-                    alpha3=first_country["alpha3"],
-                    flag=first_country["flag"],
-                )
-            else:
-                return os.environ["WORKSPACE_COUNTRY"]
+            response = OpenHexaClient().get_countries(workspace_slug=self.slug)
+            return response["workspace"]["countries"]
         except KeyError:
-            raise WorkspaceConfigError("The workspace country is not available in this environment.")
+            raise WorkspaceConfigError("The workspace countries are not available in this environment.")
 
     @property
     def database_host(self) -> str: