From a8b9e03318694530010da84446d66e8771552237 Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Sun, 14 Jan 2024 07:13:40 -0500 Subject: [PATCH] Add v2_client and http_client modules --- neon_client/__init__.py | 2 +- neon_client/__version__.py | 1 + neon_client/client.py | 86 -------- neon_client/http_client.py | 73 ++++++ neon_client/{model.py => openapi_models.py} | 0 neon_client/resources.py | 232 ++++++++++++++++++++ neon_client/utils.py | 41 ++++ neon_client/v2_client.py | 12 + 8 files changed, 360 insertions(+), 87 deletions(-) create mode 100644 neon_client/__version__.py delete mode 100644 neon_client/client.py create mode 100644 neon_client/http_client.py rename neon_client/{model.py => openapi_models.py} (100%) create mode 100644 neon_client/resources.py create mode 100644 neon_client/utils.py create mode 100644 neon_client/v2_client.py diff --git a/neon_client/__init__.py b/neon_client/__init__.py index c87c53f..fcfd5f6 100644 --- a/neon_client/__init__.py +++ b/neon_client/__init__.py @@ -1 +1 @@ -from .client import NeonClient +from .v2_client import NeonClient diff --git a/neon_client/__version__.py b/neon_client/__version__.py new file mode 100644 index 0000000..3dc1f76 --- /dev/null +++ b/neon_client/__version__.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/neon_client/client.py b/neon_client/client.py deleted file mode 100644 index 0d15107..0000000 --- a/neon_client/client.py +++ /dev/null @@ -1,86 +0,0 @@ -from typing import List -import requests - -from .model import * - - -class NeonAPI: - def __init__( - self, api_key: str, *, base_url: str = "https://console.neon.tech/api/v2/" - ): - self.api_key = api_key - self.base_url = base_url - self.session = requests.Session() - - def _request(self, method: str, path: str, **kwargs): - # Set HTTP headers for outgoing requests. - headers = kwargs.pop("headers", {}) - headers["Authorization"] = f"Bearer {self.api_key}" - headers["Accept"] = "application/json" - headers["Content-Type"] = "application/json" - - return self.session.request( - method, self.base_url + path, headers=headers, **kwargs - ) - - -class NeonClient(NeonAPI): - def __init__(self, api_key: str, **kwargs): - super().__init__(api_key, **kwargs) - # self.api_keys = APIKeyResource(self) - - def me(self) -> CurrentUserInfoResponse: - """Get information about the user.""" - - r = self._request("GET", "users/me") - r.raise_for_status() - - return CurrentUserInfoResponse(**r.json()) - - def get_api_keys(self) -> List[ApiKeysListResponseItem]: - """Get a list of API keys.""" - - r = self._request("GET", "api_keys") - r.raise_for_status() - - return [ApiKeysListResponseItem(**item) for item in r.json()] - - def create_api_key(self, name: str) -> ApiKeyCreateResponse: - """Create a new API key.""" - - r = self._request("POST", "api_keys", json={"key_name": name}) - r.raise_for_status() - - return ApiKeyCreateResponse(**r.json()) - - def delete_api_key(self, key_id: str) -> ApiKeyRevokeResponse: - """Delete an API key.""" - - r = self._request("DELETE", f"api_keys/{key_id}") - r.raise_for_status() - - return ApiKeyRevokeResponse(**r.json()) - - def get_projects(self) -> List[ProjectListItem]: - """Get a list of projects.""" - - r = self._request("GET", "projects") - r.raise_for_status() - - return [ProjectListItem(**item) for item in r.json()["projects"]] - - def create_project(self, project_id: str) -> ProjectCreateRequest: - """Create a new project.""" - - r = self._request("POST", "projects", json={"project": {"name": project_id}}) - r.raise_for_status() - - return ProjectCreateRequest(**r.json()) - - def delete_project(self, project_id: str) -> ProjectUpdateRequest: - """Delete a project.""" - - r = self._request("DELETE", f"projects/{project_id}") - r.raise_for_status() - - return ProjectUpdateRequest(**r.json()) diff --git a/neon_client/http_client.py b/neon_client/http_client.py new file mode 100644 index 0000000..10ddce1 --- /dev/null +++ b/neon_client/http_client.py @@ -0,0 +1,73 @@ +from typing import List +import requests +from pydantic import BaseModel + +from .openapi_models import * +from .__version__ import __version__ + + +class Neon_API_V2: + def __init__( + self, api_key: str, *, base_url: str = "https://console.neon.tech/api/v2/" + ): + self.api_key = api_key + self.base_url = base_url + self.session = requests.Session() + self.user_agent = f"neon-client/{__version__}" + + def request( + self, + method: str, + path: str, + *, + response_model: BaseModel = None, + response_is_array=False, + check_status_code=True, + **kwargs, + ): + """ + Sends an HTTP request to the specified path using the specified method. + + Args: + method (str): The HTTP method to use for the request. + path (str): The path to send the request to. + response_model (BaseModel, optional): The model to deserialize the response into. Defaults to None. + response_is_array (bool or str, optional): Indicates whether the response is a list of items. If a string is provided, it is used as the key to access the list in the response JSON. Defaults to False. + check_status_code (bool, optional): Indicates whether to check the status code of the response and raise an exception if it is not successful. Defaults to True. + **kwargs: Additional keyword arguments to be passed to the request. + + Returns: + The deserialized response if a response model is provided, otherwise the raw response object. + """ + # Set HTTP headers for outgoing requests. + headers = kwargs.pop("headers", {}) + headers["Authorization"] = f"Bearer {self.api_key}" + headers["Accept"] = "application/json" + headers["Content-Type"] = "application/json" + headers["User-Agent"] = self.user_agent + + r = self.session.request( + method, self.base_url + path, headers=headers, **kwargs + ) + + if check_status_code: + # TODO: add custom exception classes here. + r.raise_for_status() + + if response_model: + if response_is_array: + # Shortcut for when the response is a list of items. + if type(response_is_array) == "str": + return [ + response_model(**item) for item in r.json()[response_is_array] + ] + elif response_is_array == True: + return [response_model(**item) for item in r.json()] + else: + return response_model(**r.json()) + else: + return r + + def url_join(self, *args): + """Join the specified path components into a URL.""" + return "/".join(args) diff --git a/neon_client/model.py b/neon_client/openapi_models.py similarity index 100% rename from neon_client/model.py rename to neon_client/openapi_models.py diff --git a/neon_client/resources.py b/neon_client/resources.py new file mode 100644 index 0000000..84af2fc --- /dev/null +++ b/neon_client/resources.py @@ -0,0 +1,232 @@ +from typing import List + +from .http_client import Neon_API_V2 +from .openapi_models import ( + ApiKeyCreateResponse, + ApiKeyRevokeResponse, + ApiKeysListResponseItem, + CurrentUserInfoResponse, + Project, + ProjectResponse, + ProjectUpdateRequest, + ProjectsResponse, + DatabaseResponse, + DatabasesResponse, + DatabaseUpdateRequest, + Database, + Database1, + Database2, + DatabaseCreateRequest, +) +from .utils import validate_with_model + + +class Resource: + base_path = None + + def __init__(self, api: Neon_API_V2): + self.api = api + + +class UserResource(Resource): + base_path = "users" + response_model = CurrentUserInfoResponse + + def get_current_user_info(self): + """Get information about the user.""" + + return self.api.request( + method="GET", + path=self.api.url_join(self.path, "me"), + response_model=self.response_model, + ) + + +class APIKeyResource(Resource): + path = "api_keys" + response_model = ApiKeysListResponseItem + + def get_keys(self): + """Get a list of API keys.""" + + keys = self.api.request( + method="GET", + path=self.path, + response_model=self.response_model, + response_is_array=True, + ) + + def create_key(self, name: str): + """Create a new API key.""" + + return self.api.request( + method="POST", + path=self.path, + json={"key_name": name}, + response_model=ApiKeyCreateResponse, + ) + + def revoke_key(self, key_id: str): + """Revoke an API key.""" + + return self.api.request( + method="DELETE", + path=self.api.url_join(self.path, key_id), + response_model=ApiKeyRevokeResponse, + ) + + def get_key(self, key_id: str): + """Get an API key.""" + + return self.api.request( + "GET", + f"api_keys/{key_id}", + response_model=ApiKeyRevokeResponse, + ) + + +class ProjectResource(Resource): + path = "projects" + response_model = ProjectsResponse + response_model_single = ProjectResponse + + def get_projects(self, *, shared=False): + """Get a list of projects.""" + + project_response = self.api.request( + method="GET", + path=(self.api.url_join(self.path, "shared") if shared else self.path), + response_model=self.response_model, + ) + return project_response.projects + + def get_project(self, project_id: str): + """Get a project.""" + + project_response = self.api.request( + method="GET", + path=self.api.url_join(self.path, project_id), + response_model=self.response_model_single, + ) + return project_response.project + + def create_project(self, name: str, **kwargs): + """Create a new project.""" + + project_create_response = self.api.request( + method="POST", + path=self.path, + json={"project": {"name": name, **kwargs}}, + response_model=self.response_model_single, + ) + return project_create_response.project + + def update_project(self, project: Project): + """Update a project.""" + + payload = ProjectUpdateRequest(project=project.model_dump()) + + return self.api.request( + method="PATCH", + path=self.api.url_join(self.path, project.id), + json={"project": payload.model_dump()}, + response_model=self.response_model_single, + ) + + # def new(self, project_id: str): + # """Create a new project.""" + + # project_create_response = self.api.request( + # method="POST", + # path=self.path, + # json={"project": {"name": project_id}}, + # response_model=self.response_model, + # ) + # return project_create_response.project + + # def get_project(self, project_id: str) -> ProjectResponse: + # """Get a project.""" + + # return self.api.request( + # "GET", f"projects/{project_id}", response_model=ProjectResponse + # ) + + # def update_project(self, project: Project) -> ProjectResponse: + # """Update a project.""" + + # payload = ProjectUpdateRequest(project=project.model_dump()) + + # return self.api.request( + # "PATCH", + # f"projects/{project.id}", + # json={"project": payload.model_dump()}, + # response_model=ProjectResponse, + # ) + + # def delete_project(self, project_id: str) -> ProjectResponse: + # """Delete a project.""" + + # return self.api.request( + # "DELETE", f"projects/{project_id}", response_model=ProjectResponse + # ) + + +class DatabaseResource(Resource): + def get_databases(self, project_id: str, branch_id: str): + """Get a list of databases.""" + + databases_response = self.api.request( + method="GET", + path=self.api.url_join( + "projects", project_id, "branches", branch_id, "databases" + ), + response_model=DatabasesResponse, + ) + + return databases_response.databases + + def get_database(self, project_id: str, database_id: str): + """Get a database.""" + + database_response = self.api.request( + method="GET", + path=self.api.url_join("projects", project_id, "databases", database_id), + response_model=DatabaseResponse, + ) + return database_response.database + + @validate_with_model(DatabaseCreateRequest) + def create_database( + self, project_id: str, branch_id: str, *, obj: DatabaseCreateRequest, **kwargs + ): + """Create a new database.""" + # TODO: untested. + + database_create_response = self.api.request( + method="POST", + path=self.api.url_join( + "projects", project_id, "branches", branch_id, "databases" + ), + json=obj.model_dump(), + response_model=DatabaseResponse, + ) + return database_create_response.database + + def update_database(self, project_id: str, database: Database2): + """Update a database.""" + # TODO: This is not working yet. + + payload = DatabaseUpdateRequest(database=database.model_dump()) + + return self.api.request( + method="PATCH", + path=self.api.url_join("projects", project_id, "databases", database.id), + json=payload.model_dump(), + response_model=DatabaseResponse, + ) + + def new(self, name, **kwargs): + """Create a new database.""" + + db = Database1.model_construct(name=name, **kwargs) + return db diff --git a/neon_client/utils.py b/neon_client/utils.py new file mode 100644 index 0000000..d7d790e --- /dev/null +++ b/neon_client/utils.py @@ -0,0 +1,41 @@ +import inspect +from typing import List, Dict, Any, Union, Optional +from pydantic import BaseModel + + +def validate_with_model(*models): + """a decorator that will use the Pydantic model to parse and validate the input""" + + def decorator(func): + def wrapper(*args, **kwargs): + # Merge args and kwargs into a single dict + sig = inspect.signature(func) + bound_args = sig.bind_partial(*args, **kwargs) + bound_args.apply_defaults() + keyword_args = bound_args.arguments.pop("kwargs") + + # creating Pydantic classes + for name, param in sig.parameters.items(): + if issubclass(param.annotation, BaseModel): + for model in models: + if model == param.annotation: + model_dict = { + key: value + for key, value in keyword_args.items() + if key in model.model_fields + } + model_object = model(**model_dict) + keyword_args[name] = model_object + keyword_args = { + key: value + for key, value in keyword_args.items() + if key not in model_dict.keys() + } + break + + # Pass the model instance(s) to the wrapped function + return func(*bound_args.args, **keyword_args) + + return wrapper + + return decorator diff --git a/neon_client/v2_client.py b/neon_client/v2_client.py new file mode 100644 index 0000000..19bba32 --- /dev/null +++ b/neon_client/v2_client.py @@ -0,0 +1,12 @@ +from .http_client import Neon_API_V2 +from .resources import APIKeyResource, UserResource, ProjectResource, DatabaseResource + + +class NeonClient: + def __init__(self, api_key: str, **kwargs): + self.api = Neon_API_V2(api_key, **kwargs) + + self.api_keys = APIKeyResource(self.api) + self.users = UserResource(self.api) + self.projects = ProjectResource(self.api) + self.databases = DatabaseResource(self.api)