Use new regional endpoints (#51)

* Use new regional endpoints

* lint

.github/workflows/stage.yml

@@ -32,4 +32,4 @@ jobs:
         run: make lint
 
       - name: Test
-        run: USE_NEW_REGIONAL_ENDPOINTS="TRUE" BASE_NUCLIA_DOMAIN="stashify.cloud" GA_TESTING_SERVICE_TOKEN="${{ secrets.STAGE_TESTING_SERVICE_TOKEN }}" GA_TESTING_TOKEN="${{ secrets.STAGE_TESTING_TOKEN }}" GA_TESTING_NUA="${{ secrets.STAGE_TESTING_NUA }}" make test
+        run: BASE_NUCLIA_DOMAIN="stashify.cloud" GA_TESTING_SERVICE_TOKEN="${{ secrets.STAGE_TESTING_SERVICE_TOKEN }}" GA_TESTING_TOKEN="${{ secrets.STAGE_TESTING_TOKEN }}" GA_TESTING_NUA="${{ secrets.STAGE_TESTING_NUA }}" make test

CHANGELOG.md

@@ -1,9 +1,10 @@
 # Changelog
 
-## 1.1.23 (unreleased)
+## 1.2.0 (unreleased)
 
+** BREAKING CHANGE **
 
-- Nothing changed yet.
+- Use the new regional endpoints
 
 
 ## 1.1.22 (2023-12-15)

nuclia/__init__.py

@@ -6,8 +6,6 @@
 REGIONAL = "https://{region}." + BASE_DOMAIN
 CLOUD_ID = BASE.split("/")[-1]
 
-USE_NEW_REGIONAL_ENDPOINTS = os.environ.get("USE_NEW_REGIONAL_ENDPOINTS", "") == "TRUE"
-
 
 def get_global_url(path: str):
     return BASE + path

nuclia/decorators.py

@@ -2,7 +2,7 @@
 
 import yaml
 
-from nuclia import BASE_DOMAIN, USE_NEW_REGIONAL_ENDPOINTS
+from nuclia import BASE_DOMAIN
 from nuclia.data import get_auth
 from nuclia.exceptions import NotDefinedDefault
 from nuclia.lib.kb import Environment, NucliaDBClient
@@ -24,11 +24,8 @@ def kbs(func):
     def wrapper_checkout_kbs(*args, **kwargs):
         if "account" in kwargs:
             auth = get_auth()
-            if not USE_NEW_REGIONAL_ENDPOINTS:
-                auth.kbs(kwargs["account"])
-            else:
-                account_id = auth.get_account_id(kwargs["account"])
-                auth.kbs(account_id)
+            account_id = auth.get_account_id(kwargs["account"])
+            auth.kbs(account_id)
         return func(*args, **kwargs)
 
     return wrapper_checkout_kbs

nuclia/sdk/auth.py

@@ -6,7 +6,7 @@
 import requests
 from prompt_toolkit import prompt
 
-from nuclia import USE_NEW_REGIONAL_ENDPOINTS, get_global_url, get_regional_url
+from nuclia import get_global_url, get_regional_url
 from nuclia.cli.utils import yes_no
 from nuclia.config import Account, Config, KnowledgeBox, Zone, retrieve_account
 from nuclia.exceptions import NeedUserToken, UserTokenExpired
@@ -287,61 +287,35 @@ def zones(self) -> List[Zone]:
     def kbs(self, account: str):
         result = []
         zones = self.zones()
-        if not USE_NEW_REGIONAL_ENDPOINTS:
-            path = get_global_url(LIST_KBS.format(account=account))
+        for zoneObj in zones:
+            zoneSlug = zoneObj.slug
+            if not zoneSlug:
+                continue
+            path = get_regional_url(zoneSlug, LIST_KBS.format(account=account))
             try:
                 kbs = self._request("GET", path)
             except UserTokenExpired:
                 return []
-            region = {zone.id: zone.slug for zone in zones}
+            except requests.exceptions.ConnectionError:
+                print(
+                    f"Connection error to {get_regional_url(zoneSlug, '')}, skipping zone"
+                )
+                continue
             for kb in kbs:
-                zone = region[kb["zone"]]
-                if not zone:
-                    continue
-                url = get_regional_url(zone, f"/api/v1/kb/{kb['id']}")
+                url = get_regional_url(zoneSlug, f"/api/v1/kb/{kb['id']}")
                 kb_obj = KnowledgeBox(
                     url=url,
                     id=kb["id"],
                     slug=kb["slug"],
                     title=kb["title"],
                     account=account,
-                    region=zone,
+                    region=zoneSlug,
                 )
                 result.append(kb_obj)
-        else:
-            for zoneObj in zones:
-                zoneSlug = zoneObj.slug
-                if not zoneSlug:
-                    continue
-                path = get_regional_url(zoneSlug, LIST_KBS.format(account=account))
-                try:
-                    kbs = self._request("GET", path)
-                except UserTokenExpired:
-                    return []
-                except requests.exceptions.ConnectionError:
-                    print(
-                        f"Connection error to {get_regional_url(zoneSlug, '')}, skipping zone"
-                    )
-                    continue
-                for kb in kbs:
-                    url = get_regional_url(zoneSlug, f"/api/v1/kb/{kb['id']}")
-                    kb_obj = KnowledgeBox(
-                        url=url,
-                        id=kb["id"],
-                        slug=kb["slug"],
-                        title=kb["title"],
-                        account=account,
-                        region=zoneSlug,
-                    )
-                    result.append(kb_obj)
         return result
 
     def get_account_id(self, account_slug: str) -> str:
-        if not USE_NEW_REGIONAL_ENDPOINTS:
-            account_id = account_slug
-        else:
-            account_obj = retrieve_account(self._config.accounts or [], account_slug)
-            if not account_obj:
-                raise ValueError(f"Account {account_slug} not found")
-            account_id = account_obj.id
-        return account_id
+        account_obj = retrieve_account(self._config.accounts or [], account_slug)
+        if not account_obj:
+            raise ValueError(f"Account {account_slug} not found")
+        return account_obj.id

nuclia/sdk/kbs.py

@@ -1,6 +1,6 @@
 from typing import Dict, Optional
 
-from nuclia import USE_NEW_REGIONAL_ENDPOINTS, get_global_url, get_regional_url
+from nuclia import get_regional_url
 from nuclia.config import retrieve, retrieve_account
 from nuclia.data import get_auth
 from nuclia.decorators import account, accounts, zone
@@ -27,10 +27,7 @@ def list(self, account: Optional[str] = None):
             )
             for account_obj in accounts:
                 if account_obj.slug is not None:
-                    account_id = (
-                        USE_NEW_REGIONAL_ENDPOINTS and account_obj.id
-                    ) or account_obj.slug
-                    result.extend(self._auth.kbs(account_id))
+                    result.extend(self._auth.kbs(account_obj.id))
             self._auth._config.kbs = result
             self._auth._config.save()
 
@@ -43,13 +40,12 @@ def list(self, account: Optional[str] = None):
 
             return result
         else:
-            if not USE_NEW_REGIONAL_ENDPOINTS:
-                return self._auth.kbs(account)
-            else:
-                matching_account = retrieve_account(self._auth._config.accounts or [], account)
-                if not matching_account:
-                    raise ValueError("Account not found")
-                return self._auth.kbs(matching_account.id)
+            matching_account = retrieve_account(
+                self._auth._config.accounts or [], account
+            )
+            if not matching_account:
+                raise ValueError("Account not found")
+            return self._auth.kbs(matching_account.id)
 
     @accounts
     @account
@@ -67,14 +63,9 @@ def add(
     ):
         if not slug:
             raise ValueError("slug is required.")
-        if USE_NEW_REGIONAL_ENDPOINTS:
-            if not zone:
-                raise ValueError("zone is required")
-            path = get_regional_url(
-                zone, KBS_ENDPOINT.format(account=kwargs["account_id"])
-            )
-        else:
-            path = get_global_url(KBS_ENDPOINT.format(account=kwargs["account"]))
+        if not zone:
+            raise ValueError("zone is required")
+        path = get_regional_url(zone, KBS_ENDPOINT.format(account=kwargs["account_id"]))
         data = {
             "slug": slug,
             "anonymization": anonymization,
@@ -92,59 +83,49 @@ def add(
     @zone
     def get(
         self,
-        slug: Optional[str]=None,
-        id: Optional[str]=None,
+        slug: Optional[str] = None,
+        id: Optional[str] = None,
         **kwargs,
     ):
-        if USE_NEW_REGIONAL_ENDPOINTS:
-            zone = kwargs.get("zone")
-            if not zone:
-                raise ValueError("zone is required")
-            if not id and not slug:
-                raise ValueError("id or slug is required")
-            if slug and not id:
-                kbs = self._auth.kbs(kwargs["account_id"])
-                kb_obj = retrieve(kbs, slug)
-                if not kb_obj:
-                    raise ValueError("Knowledge Box not found")
-                id = kb_obj.id
-            path = get_regional_url(
-                zone, KB_ENDPOINT.format(account=kwargs["account_id"], kb=id)
-            )
-        else:
-            path = get_global_url(
-                KB_ENDPOINT.format(account=kwargs["account"], kb=slug)
-            )
+        zone = kwargs.get("zone")
+        if not zone:
+            raise ValueError("zone is required")
+        if not id and not slug:
+            raise ValueError("id or slug is required")
+        if slug and not id:
+            kbs = self._auth.kbs(kwargs["account_id"])
+            kb_obj = retrieve(kbs, slug)
+            if not kb_obj:
+                raise ValueError("Knowledge Box not found")
+            id = kb_obj.id
+        path = get_regional_url(
+            zone, KB_ENDPOINT.format(account=kwargs["account_id"], kb=id)
+        )
         return self._auth._request("GET", path)
 
     @accounts
     @account
     @zone
     def delete(
         self,
-        slug: Optional[str]=None,
-        id: Optional[str]=None,
+        slug: Optional[str] = None,
+        id: Optional[str] = None,
         **kwargs,
     ):
-        if USE_NEW_REGIONAL_ENDPOINTS:
-            zone = kwargs.get("zone")
-            if not zone:
-                raise ValueError("zone is required")
-            if not id and not slug:
-                raise ValueError("id or slug is required")
-            if slug and not id:
-                kbs = self._auth.kbs(kwargs["account_id"])
-                kb_obj = retrieve(kbs, slug)
-                if not kb_obj:
-                    raise ValueError("Knowledge Box not found")
-                id = kb_obj.id
-            path = get_regional_url(
-                zone, KB_ENDPOINT.format(account=kwargs["account_id"], kb=id)
-            )
-        else:
-            path = get_global_url(
-                KB_ENDPOINT.format(account=kwargs["account"], kb=slug)
-            )
+        zone = kwargs.get("zone")
+        if not zone:
+            raise ValueError("zone is required")
+        if not id and not slug:
+            raise ValueError("id or slug is required")
+        if slug and not id:
+            kbs = self._auth.kbs(kwargs["account_id"])
+            kb_obj = retrieve(kbs, slug)
+            if not kb_obj:
+                raise ValueError("Knowledge Box not found")
+            id = kb_obj.id
+        path = get_regional_url(
+            zone, KB_ENDPOINT.format(account=kwargs["account_id"], kb=id)
+        )
         return self._auth._request("DELETE", path)
 
     def default(self, kb: str):

nuclia/sdk/process.py

@@ -1,11 +1,12 @@
 from typing import Optional
 
+from nucliadb_protos.writer_pb2 import BrokerMessage
+
 from nuclia.data import get_auth
 from nuclia.decorators import nua
 from nuclia.lib.nua import NuaClient
 from nuclia.lib.nua_responses import LearningConfig, ProcessingStatus
 from nuclia.sdk.auth import NucliaAuth
-from nucliadb_protos.writer_pb2 import BrokerMessage
 
 
 class NucliaProcessing:

nuclia/tests/test_manage/test_kb.py

@@ -20,11 +20,16 @@ def test_add_and_delete_kb(testing_config):
         assert True
         return
     kbs = NucliaKBS()
-    kb = kbs.add(account=TESTING_ACCOUNT_SLUG, slug=NEW_KB_SLUG, title="Test KB", zone="europe-1")
+    kb = kbs.add(
+        account=TESTING_ACCOUNT_SLUG, slug=NEW_KB_SLUG, title="Test KB", zone="europe-1"
+    )
     assert kb["id"] is not None
     assert kb["slug"] == NEW_KB_SLUG
     assert kb["title"] == "Test KB"
-    assert kbs.get(account=TESTING_ACCOUNT_SLUG, slug=NEW_KB_SLUG, zone="europe-1") is not None
+    assert (
+        kbs.get(account=TESTING_ACCOUNT_SLUG, slug=NEW_KB_SLUG, zone="europe-1")
+        is not None
+    )
 
     kbs.delete(account=TESTING_ACCOUNT_SLUG, id=kb["id"], zone="europe-1")
     with pytest.raises(ValueError):