Run session notebook against fresh db

FlorentinD · FlorentinD · commit d0bc72acc73e · 2024-07-10T14:07:33.000+02:00
diff --git a/scripts/ci/aura_api_ci.py b/scripts/ci/aura_api_ci.py
@@ -0,0 +1,155 @@
+import logging
+import os
+import time
+from time import sleep
+from typing import Any, Dict, Optional
+
+import requests
+
+CLIENT_ID = os.environ["AURA_API_CLIENT_ID"]
+CLIENT_SECRET = os.environ["AURA_API_CLIENT_SECRET"]
+
+
+class AuraApiCI:
+    class AuraAuthToken:
+        access_token: str
+        expires_in: int
+        token_type: str
+
+        def __init__(self, json: Dict[str, Any]) -> None:
+            self.access_token = json["access_token"]
+            expires_in: int = json["expires_in"]
+            self.expires_at = int(time.time()) + expires_in
+            self.token_type = json["token_type"]
+
+        def is_expired(self) -> bool:
+            return self.expires_at >= int(time.time())
+
+    def __init__(self) -> None:
+        self._token: Optional[AuraApiCI.AuraAuthToken] = None
+        self._logger = logging.getLogger()
+
+    def _build_header(self) -> Dict[str, str]:
+        return {"Authorization": f"Bearer {self._auth_token()}", "User-agent": "neo4j-graphdatascience-ci"}
+
+    def _auth_token(self) -> str:
+        if self._token is None or self._token.is_expired():
+            self._token = self._update_token()
+        return self._token.access_token
+
+    def _update_token(self) -> AuraAuthToken:
+        data = {
+            "grant_type": "client_credentials",
+        }
+
+        self._logger.debug("Updating oauth token")
+
+        response = requests.post("https://api-staging.neo4j.io/oauth/token", data=data, auth=(CLIENT_ID, CLIENT_SECRET))
+
+        response.raise_for_status()
+
+        return AuraApiCI.AuraAuthToken(response.json())
+
+    def create_ds_instance(self, name: str) -> Dict[str, Any]:
+        return self.create_instance(name, memory="8GB", type="gds")
+
+    def create_instance(self, name: str, memory: str, type: str) -> Dict[str, Any]:
+        CREATE_OK_MAX_WAIT_TIME = 10
+
+        data = {
+            "name": name,
+            "memory": memory,
+            "version": "5",
+            "region": "europe-west1",
+            "type": type,
+            "cloud_provider": "gcp",
+            "tenant_id": self.get_tenant_id(),
+        }
+
+        should_retry = True
+        wait_time = 1
+
+        while should_retry:
+            sleep(wait_time)
+            wait_time *= 2
+
+            response = requests.post(
+                "https://api-staging.neo4j.io/v1/instances",
+                json=data,
+                headers=self._build_header(),
+            )
+            should_retry = response.status_code in [500, 502, 503, 504, 405] and CREATE_OK_MAX_WAIT_TIME > wait_time
+
+            if should_retry:
+                logging.debug(f"Error code: {response.status_code} - Retrying in {wait_time} s")
+
+        response.raise_for_status()
+
+        return response.json()["data"]  # type: ignore
+
+    def check_running(self, db_id: str) -> None:
+        RUNNING_MAX_WAIT_TIME = 60 * 5
+
+        should_retry = True
+        wait_time = 1
+
+        while should_retry:
+            sleep(wait_time)
+            wait_time *= 2
+
+            response = requests.get(
+                f"https://api-staging.neo4j.io/v1/instances/{db_id}",
+                headers=self._build_header(),
+            )
+
+            instance_status = "?"
+            if response.status_code == 200:
+                instance_status = response.json()["data"]["status"]
+
+            should_retry = (
+                response.status_code in [500, 502, 503, 504] or instance_status == "creating"
+            ) and RUNNING_MAX_WAIT_TIME > wait_time
+
+            if should_retry:
+                logging.debug(
+                    f"Status code: {response.status_code}, Status: {instance_status} - Retrying in {wait_time} s"
+                )
+
+        response.raise_for_status()
+
+    def teardown_instance(self, db_id: str) -> None:
+        TEARDOWN_MAX_WAIT_TIME = 10
+
+        should_retry = True
+        wait_time = 1
+
+        while should_retry:
+            sleep(wait_time)
+            wait_time *= 2
+
+            response = requests.delete(
+                f"https://api-staging.neo4j.io/v1/instances/{db_id}",
+                headers=self._build_header(),
+            )
+
+            if response.status_code == 202:
+                should_retry = False
+
+            should_retry = (response.status_code in [500, 502, 503, 504]) and TEARDOWN_MAX_WAIT_TIME > wait_time
+
+            if should_retry:
+                logging.debug(f"Status code: {response.status_code} - Retrying in {wait_time} s")
+
+        response.raise_for_status()
+
+    def get_tenant_id(self) -> str:
+        response = requests.get(
+            "https://api-staging.neo4j.io/v1/tenants",
+            headers=self._build_header(),
+        )
+        response.raise_for_status()
+
+        raw_data = response.json()["data"]
+        assert len(raw_data) == 1
+
+        return raw_data[0]["id"]  # type: ignore
diff --git a/scripts/ci/run_targeting_aura.py b/scripts/ci/run_targeting_aura.py
@@ -2,106 +2,11 @@
 import os
 import random as rd
 import sys
-from time import sleep
-from typing import Any, Dict
 
-import requests as req
+from aura_api_ci import AuraApiCI
 
 logging.basicConfig(level=logging.INFO)
 
-CLIENT_ID = os.environ["AURA_API_CLIENT_ID"]
-CLIENT_SECRET = os.environ["AURA_API_CLIENT_SECRET"]
-
-
-def get_access_token() -> str:
-    data = {
-        "grant_type": "client_credentials",
-    }
-
-    # getting a token like {'access_token':'X','expires_in':3600,'token_type':'bearer'}
-    response = req.post("https://api-staging.neo4j.io/oauth/token", data=data, auth=(CLIENT_ID, CLIENT_SECRET))
-
-    response.raise_for_status()
-
-    return response.json()["access_token"]  # type: ignore
-
-
-def create_instance(name: str, access_token: str) -> Dict[str, Any]:
-    CREATE_OK_MAX_WAIT_TIME = 10
-
-    data = {
-        "name": name,
-        "memory": "8GB",
-        "version": "5",
-        "region": "europe-west1",
-        "type": "gds",
-        "cloud_provider": "gcp",
-        "tenant_id": get_tenant_id(access_token),
-    }
-
-    should_retry = True
-    wait_time = 1
-
-    while should_retry:
-        sleep(wait_time)
-        wait_time *= 2
-
-        response = req.post(
-            "https://api-staging.neo4j.io/v1/instances",
-            json=data,
-            headers={"Authorization": f"Bearer {access_token}"},
-        )
-        should_retry = response.status_code in [500, 502, 503, 504, 405] and CREATE_OK_MAX_WAIT_TIME > wait_time
-
-        if should_retry:
-            logging.debug(f"Error code: {response.status_code} - Retrying in {wait_time} s")
-
-    response.raise_for_status()
-
-    return response.json()["data"]  # type: ignore
-
-
-def check_running(access_token: str, db_id: str) -> None:
-    RUNNING_MAX_WAIT_TIME = 60 * 5
-
-    should_retry = True
-    wait_time = 1
-
-    while should_retry:
-        sleep(wait_time)
-        wait_time *= 2
-
-        response = req.get(
-            f"https://api-staging.neo4j.io/v1/instances/{db_id}",
-            headers={"Authorization": f"Bearer {access_token}"},
-        )
-
-        instance_status = "?"
-        if response.status_code == 200:
-            instance_status = response.json()["data"]["status"]
-
-        should_retry = (
-            response.status_code in [500, 502, 503, 504] or instance_status == "creating"
-        ) and RUNNING_MAX_WAIT_TIME > wait_time
-
-        if should_retry:
-            logging.debug(f"Status code: {response.status_code}, Status: {instance_status} - Retrying in {wait_time} s")
-
-    response.raise_for_status()
-
-
-def get_tenant_id(access_token: str) -> str:
-    response = req.get(
-        "https://api-staging.neo4j.io/v1/tenants",
-        headers={"Authorization": f"Bearer {access_token}"},
-    )
-    response.raise_for_status()
-
-    raw_data = response.json()["data"]
-    assert len(raw_data) == 1
-
-    return raw_data[0]["id"]  # type: ignore
-
 
 def run_tests(uri: str, username: str, password: str) -> None:
     cmd = (
@@ -120,45 +25,18 @@ def run_notebooks(uri: str, username: str, password: str) -> None:
         raise Exception("Failed to run notebooks")
 
 
-def teardown_instance(access_token: str, db_id: str) -> None:
-    TEARDOWN_MAX_WAIT_TIME = 10
-
-    should_retry = True
-    wait_time = 1
-
-    while should_retry:
-        sleep(wait_time)
-        wait_time *= 2
-
-        response = req.delete(
-            f"https://api-staging.neo4j.io/v1/instances/{db_id}",
-            headers={"Authorization": f"Bearer {access_token}"},
-        )
-
-        if response.status_code == 202:
-            should_retry = False
-
-        should_retry = (response.status_code in [500, 502, 503, 504]) and TEARDOWN_MAX_WAIT_TIME > wait_time
-
-        if should_retry:
-            logging.debug(f"Status code: {response.status_code} - Retrying in {wait_time} s")
-
-    response.raise_for_status()
-
-
 def main() -> None:
-    access_token = get_access_token()
-    logging.info("Access token for creation acquired")
+    aura_api = AuraApiCI()
 
     MAX_INT = 1000000
     instance_name = f"ci-build-{sys.argv[2]}" if len(sys.argv) > 1 else "ci-instance-" + str(rd.randint(0, MAX_INT))
 
-    create_result = create_instance(instance_name, access_token)
+    create_result = aura_api.create_ds_instance(instance_name)
     instance_id = create_result["id"]
     logging.info("Creation of database accepted")
 
     try:
-        check_running(access_token, instance_id)
+        aura_api.check_running(instance_id)
         logging.info("Database %s up and running", instance_id)
 
         if sys.argv[1] == "tests":
@@ -178,10 +56,7 @@ def main() -> None:
         else:
             logging.error(f"Invalid target: {sys.argv[1]}")
     finally:
-        access_token = get_access_token()
-        logging.info("Access token for teardown acquired")
-
-        teardown_instance(access_token, instance_id)
+        aura_api.teardown_instance(instance_id)
         logging.info("Teardown of instance %s successful", instance_id)
 
 
diff --git a/scripts/ci/run_targeting_aura_sessions.py b/scripts/ci/run_targeting_aura_sessions.py
@@ -0,0 +1,42 @@
+# run `tox -e jupyter-notebook-session-ci`
+
+import logging
+import os
+import random as rd
+import sys
+
+from aura_api_ci import AuraApiCI
+
+logging.basicConfig(level=logging.INFO)
+
+
+def main() -> None:
+    aura_api = AuraApiCI()
+
+    MAX_INT = 1000000
+    instance_name = f"ci-build-{sys.argv[1]}" if len(sys.argv) > 0 else "ci-instance-" + str(rd.randint(0, MAX_INT))
+
+    create_result = aura_api.create_instance(instance_name, memory="1GB", type="professional-db")
+    instance_id = create_result["id"]
+    logging.info("Creation of database accepted")
+
+    try:
+        aura_api.check_running(instance_id)
+        logging.info("Database %s up and running", instance_id)
+
+        uri = (create_result["connection_url"],)
+        username = (create_result["username"],)
+        password = (create_result["password"],)
+
+        cmd = f"AURA_DB_ADDRESS={uri} AURA_DB_USER={username} AURA_DB_PW={password} tox -e jupyter-notebook-session-ci"
+
+        if os.system(cmd) != 0:
+            raise Exception("Failed to run notebooks")
+
+    finally:
+        aura_api.teardown_instance(instance_id)
+        logging.info("Teardown of instance %s successful", instance_id)
+
+
+if __name__ == "__main__":
+    main()
