From 19e6e71733222fe605202b7b1256bc8adb9d7c72 Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Wed, 17 Jan 2024 11:15:18 -0500 Subject: [PATCH] snapshop --- neon_client/client.py | 81 ++++++++++++++++++++++++++++++++------ neon_client/http_client.py | 16 +++++--- neon_client/resources.py | 22 ++++++++++- 3 files changed, 99 insertions(+), 20 deletions(-) diff --git a/neon_client/client.py b/neon_client/client.py index b4746ed..29284ed 100644 --- a/neon_client/client.py +++ b/neon_client/client.py @@ -7,8 +7,7 @@ from . import models class ItemView: """A view into a single item.""" - def __init__(self, item, key_id=None, neon=None): - self.neon = neon + def __init__(self, item, key_id=None): self._item = item self._key_id = key_id @@ -44,14 +43,7 @@ class ItemView: class CollectionView: """A view into a collection of items.""" - def __init__( - self, - collection, - key_ids=None, - collection_id=None, - neon=None, - ): - self.neon = neon + def __init__(self, collection, key_ids=None, collection_id=None): self.pagination = None if not key_ids: @@ -107,16 +99,26 @@ class NeonClient: return self.resources.users.get_current_user_info() def api_keys(self): + """Get a list of API keys.""" + return CollectionView(self.resources.api_keys.get_list(), key_ids=["id"]) - def projects(self, **kwargs): + def projects(self, shared=False, **kwargs): + """Get a list of projects.""" + + a = self.resources.projects.get_list(shared=shared, **kwargs) + print(a) + exit() + return CollectionView( - self.resources.projects.get_list(**kwargs), + self.resources.projects.get_list(shared=shared, **kwargs), key_ids=["id", "name"], collection_id="projects", ) def project(self, project_id: str, **kwargs): + """Get a single project.""" + return ItemView( self.resources.projects.get(project_id, **kwargs), key_id="project" ) @@ -155,6 +157,61 @@ class NeonClient: key_id="branch", ) + def branch_create(self, project_id: str, **kwargs): + return ItemView( + self.resources.branches.create(project_id, **kwargs), + key_id="branch", + ) + + def branch_delete(self, project_id: str, branch_id: str, **kwargs): + return ItemView( + self.resources.branches.delete(project_id, branch_id, **kwargs), + key_id="branch", + ) + + def branch_update(self, project_id: str, branch_id: str, **kwargs): + # TODO: untested. + return ItemView( + self.resources.branches.update(project_id, branch_id, **kwargs), + key_id="branch", + ) + + def branch_rename(self, project_id: str, branch_id: str, **kwargs): + # TODO: untested. + return ItemView( + self.resources.branches.rename(project_id, branch_id, **kwargs), + key_id="branch", + ) + + def branch_add_compute(self, project_id: str, branch_id: str, **kwargs): + # TODO: untested. + return ItemView( + self.resources.branches.add_compute(project_id, branch_id, **kwargs), + key_id="branch", + ) + + def branch_remove_compute(self, project_id: str, branch_id: str, **kwargs): + # TODO: untested. + return ItemView( + self.resources.branches.remove_compute(project_id, branch_id, **kwargs), + key_id="branch", + ) + + def branch_set_primary(self, project_id: str, branch_id: str, **kwargs): + # TODO: untested. + return ItemView( + self.resources.branches.set_primary(project_id, branch_id, **kwargs), + key_id="branch", + ) + + def get_connection_string(self, project_id: str, branch_id: str, database_id: str): + # TODO: implement this. + return self.resources.databases.get_connection_string( + project_id, + branch_id, + database_id, + ) + # def branch_create(self, project_id: str, **kwargs): # return ItemView( # self.resources.branches.create(project_id, **kwargs), diff --git a/neon_client/http_client.py b/neon_client/http_client.py index ccfbea4..a051b19 100644 --- a/neon_client/http_client.py +++ b/neon_client/http_client.py @@ -1,4 +1,5 @@ from typing import List + import requests from pydantic import BaseModel @@ -24,7 +25,7 @@ class Neon_API_V2: response_model: BaseModel = None, response_is_array=False, check_status_code=True, - _debug_pagination=False, + _debug=False, **kwargs, ): """ @@ -56,7 +57,7 @@ class Neon_API_V2: # TODO: add custom exception classes here. try: r.raise_for_status() - except: + except requests.exceptions.HTTPError: raise NeonClientException(r.text) if response_model: @@ -64,14 +65,17 @@ class Neon_API_V2: # Shortcut for when the response is a list of items. if type(response_is_array) == "str": response_parsed = [ - response_model(**item) for item in r.json()[response_is_array] + response_model.model_construct(**item) + for item in r.json()[response_is_array] ] elif response_is_array == True: - response_parsed = [response_model(**item) for item in r.json()] + response_parsed = [ + response_model.model_construct(**item) for item in r.json() + ] else: - response_parsed = response_model(**r.json()) + response_parsed = response_model.model_construct(**r.json()) - if _debug_pagination: + if _debug: print(r.json()) return response_parsed diff --git a/neon_client/resources.py b/neon_client/resources.py index 5cdd5a6..fbd08c4 100644 --- a/neon_client/resources.py +++ b/neon_client/resources.py @@ -121,9 +121,9 @@ class ProjectResource(Resource): return self.api.request( method="POST", path=self.base_path, - json={"project": kwargs}, + json=kwargs, response_model=models.ProjectResponse, - ).model_dump() + ) def update(self, project: models.Project): """Update a project.""" @@ -252,6 +252,24 @@ class DatabaseResource(Resource): response_model=models.DatabaseResponse, ) + def get_connection_string(self, project_id: str, branch_id: str, database_id: str): + """Get a database connection string.""" + + # TODO: this isn't correct. + return self.api.request( + method="GET", + path=self.api.url_join( + "projects", + project_id, + "branches", + branch_id, + "databases", + database_id, + "connection_string", + ), + response_model=models.DatabaseConnectionStringResponse, + ) + class BranchResource(Resource): """A resource for interacting with branches."""