snapshop

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

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

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."""