From 62237c24860d1993802bc10e77be3da19f1e7830 Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Mon, 22 Jan 2024 15:29:50 -0500 Subject: [PATCH] Update provisioner field to be optional --- neon_client/client.py | 58 ++++++++++++++++++++++++++++++++++++++++++- neon_client/schema.py | 2 +- 2 files changed, 58 insertions(+), 2 deletions(-) diff --git a/neon_client/client.py b/neon_client/client.py index 164a53b..2112475 100644 --- a/neon_client/client.py +++ b/neon_client/client.py @@ -3,6 +3,7 @@ import typing as t import requests +from . import schema from .utils import compact_mapping from .exceptions import NeonClientException @@ -13,6 +14,21 @@ NEON_API_KEY_ENVIRON = "NEON_API_KEY" NEON_API_BASE_URL = "https://console.neon.tech/api/v2/" +def returns_model(model, is_array=False): + """Decorator that returns a model instance.""" + + def decorator(func): + def wrapper(*args, **kwargs): + if is_array: + return [model.model_construct(**item) for item in func(*args, **kwargs)] + else: + return model.model_construct(**func(*args, **kwargs)) + + return wrapper + + return decorator + + class NeonAPI: def __init__(self, api_key: str, *, base_url: str = None): """A Neon API client.""" @@ -70,22 +86,27 @@ class NeonAPI: return cls(os.environ[NEON_API_KEY_ENVIRON]) + @returns_model(schema.CurrentUserInfoResponse) def me(self) -> t.Dict[str, t.Any]: """Get the current user.""" return self.request("GET", "users/me") + @returns_model(schema.ApiKeysListResponseItem, is_array=True) def api_keys(self) -> t.List[t.Dict[str, t.Any]]: """Get a list of API keys.""" return self.request("GET", "api_keys") + @returns_model(schema.ApiKeyCreateResponse) def api_key_create(self, key_name: str) -> t.Dict[str, t.Any]: """Create a new API key.""" return self.request("POST", "api_keys", json={"name": key_name}) + @returns_model(schema.ApiKeyRevokeResponse) def api_key_revoke(self, api_key_id: str) -> t.Dict[str, t.Any]: """Revoke an API key.""" return self.request("DELETE", f"api_keys/{ api_key_id }") + @returns_model(schema.ProjectsResponse) def projects( self, *, @@ -93,13 +114,20 @@ class NeonAPI: cursor: str | None = None, limit: int | None = None, ) -> t.List[t.Dict[str, t.Any]]: - """Get a list of projects.""" + """Get a list of projects. If shared is True, get a list of shared projects. + + :param shared: Whether to retrieve shared projects (default is False). + :param cursor: The cursor for pagination (default is None). + :param limit: The maximum number of projects to retrieve (default is None). + :return: A list of dictionaries representing the projects. + """ r_path = "projects" if not shared else "projects/shared" r_params = compact_mapping({"cursor": cursor, "limit": limit}) return self.request("GET", r_path, params=r_params) + @returns_model(schema.ProjectResponse) def project(self, project_id: str) -> t.Dict[str, t.Any]: """Get a project.""" @@ -110,16 +138,19 @@ class NeonAPI: r_path, ) + @returns_model(schema.ProjectResponse) def project_create(self, **json: dict) -> t.Dict[str, t.Any]: """Create a new project. Accepts all keyword arguments for json body.""" return self.request("POST", "projects", json=json) + @returns_model(schema.ProjectResponse) def project_delete(self, project_id: str) -> t.Dict[str, t.Any]: """Delete a project.""" return self.request("DELETE", f"projects/{ project_id }") + @returns_model(schema.BranchesResponse) def branches( self, project_id: str, @@ -136,6 +167,7 @@ class NeonAPI: # Make the request. return self.request("GET", r_path, params=r_params) + @returns_model(schema.BranchResponse) def branch(self, project_id: str, branch_id: str) -> t.Dict[str, t.Any]: """Get a branch.""" @@ -145,10 +177,12 @@ class NeonAPI: # Make the request. return self.request("GET", r_path) + @returns_model(schema.BranchOperations) def branch_create(self, project_id: str, **json: dict) -> t.Dict[str, t.Any]: """Create a new branch. Accepts all keyword arguments for json body.""" return self.request("POST", f"projects/{ project_id }/branches", json=json) + @returns_model(schema.BranchOperations) def branch_update( self, project_id: str, branch_id: str, **json: dict ) -> t.Dict[str, t.Any]: @@ -158,10 +192,12 @@ class NeonAPI: "PATCH", f"projects/{ project_id }/branches/{ branch_id }", json=json ) + @returns_model(schema.BranchOperations) def branch_delete(self, project_id: str, branch_id: str) -> t.Dict[str, t.Any]: """Delete a branch by branch_id.""" return self.request("DELETE", f"projects/{ project_id }/branches/{ branch_id }") + @returns_model(schema.BranchOperations) def branch_set_as_primary( self, project_id: str, branch_id: str ) -> t.Dict[str, t.Any]: @@ -171,6 +207,7 @@ class NeonAPI: "POST", f"projects/{ project_id }/branches/{ branch_id }/set_as_primary" ) + @returns_model(schema.DatabasesResponse) def databases( self, project_id: str, @@ -190,6 +227,7 @@ class NeonAPI: # Make the request. return self.request("GET", r_path, params=r_params) + @returns_model(schema.DatabaseResponse) def database( self, project_id: str, branch_id: str, database_id: str ) -> t.Dict[str, t.Any]: @@ -203,6 +241,7 @@ class NeonAPI: # Make the request. return self.request("GET", r_path) + @returns_model(schema.DatabaseResponse) def database_create( self, project_id: str, branch_id: str, **json: dict ) -> t.Dict[str, t.Any]: @@ -214,6 +253,7 @@ class NeonAPI: json=json, ) + @returns_model(schema.DatabaseResponse) def database_update( self, project_id: str, branch_id: str, database_id: str, **json: dict ) -> t.Dict[str, t.Any]: @@ -225,6 +265,7 @@ class NeonAPI: json=json, ) + @returns_model(schema.DatabaseResponse) def database_delete( self, project_id: str, branch_id: str, database_id: str ) -> t.Dict[str, t.Any]: @@ -235,12 +276,14 @@ class NeonAPI: f"projects/{ project_id }/branches/{ branch_id }/databases/{ database_id }", ) + @returns_model(schema.EndpointsResponse) def endpoints(self, project_id: str, branch_id: str) -> t.Dict[str, t.Any]: """Get a list of endpoints for a given branch.""" return self.request( "GET", f"projects/{ project_id }/branches/{ branch_id }/endpoints" ) + @returns_model(schema.EndpointResponse) def endpoint( self, project_id: str, branch_id: str, endpoint_id: str ) -> t.Dict[str, t.Any]: @@ -250,6 +293,7 @@ class NeonAPI: f"projects/{ project_id }/branches/{ branch_id }/endpoints/{ endpoint_id }", ) + @returns_model(schema.EndpointOperations) def endpoint_create( self, project_id: str, @@ -264,6 +308,7 @@ class NeonAPI: json=json, ) + @returns_model(schema.EndpointOperations) def endpoint_delete( self, project_id: str, branch_id: str, endpoint_id: str ) -> t.Dict[str, t.Any]: @@ -274,6 +319,7 @@ class NeonAPI: f"projects/{ project_id }/branches/{ branch_id }/endpoints/{ endpoint_id }", ) + @returns_model(schema.EndpointOperations) def endpoint_update( self, project_id: str, branch_id: str, endpoint_id: str, **json: dict ) -> t.Dict[str, t.Any]: @@ -285,6 +331,7 @@ class NeonAPI: json=json, ) + @returns_model(schema.EndpointOperations) def endpoint_start(self, project_id: str, branch_id: str, endpoint_id: str): """Start an endpoint by endpoint_id.""" @@ -293,6 +340,7 @@ class NeonAPI: f"projects/{ project_id }/branches/{ branch_id }/endpoints/{ endpoint_id }/start", ) + @returns_model(schema.EndpointOperations) def endpoint_suspend(self, project_id: str, branch_id: str, endpoint_id: str): """Suspend an endpoint by endpoint_id.""" @@ -301,12 +349,14 @@ class NeonAPI: f"projects/{ project_id }/branches/{ branch_id }/endpoints/{ endpoint_id }/suspend", ) + @returns_model(schema.RolesResponse) def roles(self, project_id: str, branch_id: str) -> t.Dict[str, t.Any]: """Get a list of roles for a given branch.""" return self.request( "GET", f"projects/{ project_id }/branches/{ branch_id }/roles" ) + @returns_model(schema.RoleResponse) def role( self, project_id: str, branch_id: str, role_name: str ) -> t.Dict[str, t.Any]: @@ -315,6 +365,7 @@ class NeonAPI: "GET", f"projects/{ project_id }/branches/{ branch_id }/roles/{ role_name }" ) + @returns_model(schema.RoleOperations) def role_create( self, project_id: str, @@ -329,6 +380,7 @@ class NeonAPI: json={"role": {"name": role_name}}, ) + @returns_model(schema.RoleOperations) def role_delete( self, project_id: str, @@ -342,6 +394,7 @@ class NeonAPI: f"projects/{ project_id }/branches/{ branch_id }/roles/{ role_name }", ) + @returns_model(schema.RolePasswordResponse) def role_password_reveal( self, project_id: str, @@ -355,6 +408,7 @@ class NeonAPI: f"projects/{ project_id }/branches/{ branch_id }/roles/{ role_name }/reveal_password", ) + @returns_model(schema.RoleOperations) def role_password_reset( self, project_id: str, @@ -368,10 +422,12 @@ class NeonAPI: f"projects/{ project_id }/branches/{ branch_id }/roles/{ role_name }/reset_password", ) + @returns_model(schema.OperationsResponse) def operations(self, project_id: str) -> t.Dict[str, t.Any]: """Get a list of operations.""" return self.request("GET", f"projects/{ project_id }/operations") + @returns_model(schema.OperationResponse) def operation(self, project_id: str, operation_id: str) -> t.Dict[str, t.Any]: """Get an operation.""" return self.request( diff --git a/neon_client/schema.py b/neon_client/schema.py index 6eaf557..a5829c9 100644 --- a/neon_client/schema.py +++ b/neon_client/schema.py @@ -716,7 +716,7 @@ class ProjectListItem(BaseModel): ) region_id: str = Field(..., description="The region identifier\n") name: str = Field(..., description="The project name\n") - provisioner: Provisioner + provisioner: Optional[Provisioner] default_endpoint_settings: Optional[DefaultEndpointSettings] = None settings: Optional[ProjectSettingsData] = None pg_version: conint(ge=14, le=16) = Field(