Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 156 additions & 0 deletions scripts/ci/aura_api_ci.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,156 @@
import logging
import time
from time import sleep
from typing import Any, Dict, Optional

import requests


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, client_id: str, client_secret: str, tenant_id: Optional[str] = None) -> None:
self._token: Optional[AuraApiCI.AuraAuthToken] = None
self._logger = logging.getLogger()
self._auth = (client_id, client_secret)
self._tenant_id = tenant_id

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=self._auth)

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:
if self._tenant_id:
return self._tenant_id

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
138 changes: 8 additions & 130 deletions scripts/ci/run_targeting_aura.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
Expand All @@ -120,45 +25,21 @@ 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")
client_id = os.environ["AURA_API_CLIENT_ID"]
client_secret = os.environ["AURA_API_CLIENT_SECRET"]
tenant_id = os.environ.get("TENANT_ID")
aura_api = AuraApiCI(client_id=client_id, client_secret=client_secret, tenant_id=tenant_id)

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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So this file runs tests and notebooks using just an AuraDS instance?
I suggest renaming the file run_targeting_aurads

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I realise this will need another TC change

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":
Expand All @@ -178,10 +59,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)


Expand Down
45 changes: 45 additions & 0 deletions scripts/ci/run_targeting_aura_sessions.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
# 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:
client_id = os.environ["AURA_API_CLIENT_ID"]
client_secret = os.environ["AURA_API_CLIENT_SECRET"]
tenant_id = os.environ.get("TENANT_ID")
aura_api = AuraApiCI(client_id=client_id, client_secret=client_secret, tenant_id=tenant_id)

MAX_INT = 1000000
instance_name = f"ci-build-{sys.argv[1]}" if len(sys.argv) > 1 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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Meanwhile, this test runs with an AuraDB instance and dedicated or AuraDS-based sessions.
I suggest renaming it run_sessions_targeting_auradb


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