From 5860a94d2b6221d988b1bd57d0c153bdefe8184e Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Fri, 19 Jan 2024 16:28:17 -0500 Subject: [PATCH] mass simplification --- neon_client/client.py | 504 ++++++++++++++---------------------------- 1 file changed, 162 insertions(+), 342 deletions(-) diff --git a/neon_client/client.py b/neon_client/client.py index a3e7754..0303f4a 100644 --- a/neon_client/client.py +++ b/neon_client/client.py @@ -17,319 +17,6 @@ NEON_API_KEY_ENVIRON = "NEON_API_KEY" NEON_API_BASE_URL = "https://console.neon.tech/api/v2/" -class NeonResource: - def __init__(self, client, obj, data_model: BaseModel): - """A Neon Resource. Used to wrap API responses, and deserialize them. - - Args: - client (NeonAPI): The Neon API client. - obj (dict): The API key data. - data_model (BaseModel, optional): The data model to use for deserialization. Defaults to None. - """ - self._client = client - self._data = obj - self._data_model = data_model - self.__cached_obj = None - - @property - def obj(self): - """The object in question.""" - - if self._data_model: - # Cache the object. - if not self.__cached_obj: - # Deserialize the object. Fallback to the model constructor if the - # data model is not a Pydantic model. - self.__cached_obj = self._data_model(**self._data) - # try: - # self.__cached_obj = self._data_model(**self._data) - # except ValidationError: - # self.__cached_obj = self._data_model.model_construct(**self._data) - - # Return the cached object. - return self.__cached_obj - else: - # Return the raw data. - return self._data - - def __getattribute__(self, name): - """Get an attribute from the object.""" - - try: - # Try to get the attribute from this object. - return super().__getattribute__(name) - except AttributeError: - # Try to get the attribute from the object data. - return getattr(self.obj, name) - - def __getitem__(self, key): - """Get an item from the object.""" - - return getattr(self.obj, key, None) or self.obj[key] - - def __repr__(self): - """Return a string representation.""" - - return repr(self.obj) - - -class APIKey(NeonResource): - """A Neon API key.""" - - @classmethod - def create(cls, client, key_name: str): - """Create a new API key.""" - - obj = schema.ApiKeyCreateRequest(key_name=key_name) - r = client.request("POST", "api_keys", json=jsonable_encoder(obj)) - - return cls(client=client, obj=r, data_model=schema.ApiKeyCreateResponse) - - @classmethod - def get_list(cls, client): - """Get a list of API keys.""" - - obj = client.request("GET", "api_keys") - return [ - cls(client=client, obj=obj, data_model=schema.ApiKeysListResponseItem) - for obj in obj - ] - - @classmethod - def revoke_with_id(cls, client, api_key_id): - """Revoke an API key, via object instance.""" - - obj = client.request("DELETE", f"api_keys/{ api_key_id }") - - return cls(client=client, obj=obj, data_model=schema.ApiKeyRevokeResponse) - - def revoke(self): - """Revoke the API key.""" - - return self.revoke_with(self._client, self.id) - - -class User(NeonResource): - """A Neon user.""" - - @classmethod - def get_current_user_info(cls, client): - """Get the current user.""" - - # Make the Request. - r = client.request("GET", "users/me") - - # Deserialize the response. - return cls( - client=client, - obj=r, - data_model=schema.CurrentUserInfoResponse, - ) - - -class Project(NeonResource): - @classmethod - def get_list( - cls, - client, - *, - shared: bool = False, - cursor: str | None = None, - limit: int | None = None, - ): - """Get a list of projects.""" - - r_path = "projects" if not shared else "projects/shared" - r_params = compact_mapping({"cursor": cursor, "limit": limit}) - - r = client.request("GET", r_path, params=r_params) - - return [ - cls(client=client, obj=obj, data_model=schema.ProjectListItem) - for obj in r["projects"] - ] - - @classmethod - def get(cls, client, project_id: str): - """Get a project.""" - - r_path = client.url_join("projects", project_id) - r = client.request( - "GET", - r_path, - response_model=schema.ProjectResponse, - ) - - return cls( - client=client, - obj=r.project, - data_model=schema.ProjectListItem, - ) - - @classmethod - def create(cls, client, **json: dict): - """Create a new project. Accepts all keyword arguments from ProjectCreateRequest.""" - - r = client.request("POST", "projects", json=jsonable_encoder(json)) - - return cls( - client=client, - obj=r, - data_model=schema.ProjectResponse, - ) - - @classmethod - def delete_with_id(cls, client, project_id: str): - """Delete a project.""" - - r = client.request("DELETE", f"projects/{ project_id }") - - return cls( - client=client, - obj=r, - data_model=schema.ProjectResponse, - ) - - def delete(self): - """Delete the project.""" - - return self.delete_with_id(self._client, self.id) - - -class Branch(NeonResource): - @classmethod - def list( - cls, - client, - project_id: str, - *, - cursor: str | None = None, - limit: int | None = None, - ): - """Get a list of branches.""" - - # Construct the request path and parameters. - r_path = "/".join(["projects", project_id, "branches"]) - r_params = compact_mapping({"cursor": cursor, "limit": limit}) - - # Make the request. - r = client.request("GET", r_path, params=r_params) - - # Deserialize the response. - return [ - cls(client=client, obj=obj, data_model=schema.ProjectListItem) - for obj in r["branches"] - ] - - @classmethod - def get(cls, client, project_id: str, branch_id: str): - """Get a list of branches.""" - - # Construct the request path. - r_path = client.url_join("projects", project_id, "branches", branch_id) - - # Make the request. - obj = client.request("GET", r_path, response_model=schema.ProjectResponse) - - # Deserialize the response. - return cls( - client=client, obj=jsonable_encoder(obj.branch), data_model=schema.Project - ) - - @classmethod - def create( - cls, - client, - project_id: str, - *, - endpoints: t.List[schema.BranchCreateRequestEndpointOptions] | None = None, - branch: schema.Branch2 | None = None, - **kwargs, - ): - """Create a new branch.""" - - # Construct the request object. - kwargs.setdefault("endpoints", endpoints) - kwargs.setdefault("branch", branch) - - # Validate and prepare the request body. - obj = schema.BranchCreateRequest(**kwargs) - - # Make the request. - r_path = client.url_join("projects", project_id, "branches") - - # Skip the request body if there are no endpoints or branch. - if obj.endpoints or obj.branch: - r = client.request("POST", r_path, json=jsonable_encoder(obj)) - else: - r = client.request("POST", r_path) - - return cls(client=client, obj=r, data_model=schema.BranchResponse) - - @classmethod - def get_connection_string(cls, client, project_id: str, branch_id: str): - """Get a connection string for a branch.""" - - # Construct the request path. - r_path = client.url_join( - "projects", project_id, "branches", branch_id, "connection_string" - ) - - # Make the request. - obj = client.request( - "GET", r_path, response_model=schema.BranchConnectionStringResponse - ) - - # Deserialize the response. - return obj.connection_string - - -class Operation(NeonResource): - @classmethod - def list( - cls, - client, - project_id: str, - *, - cursor: str | None = None, - limit: int | None = None, - ): - """Get a list of operations.""" - - # Construct the request path and parameters. - r_path = client.url_join("projects", project_id, "operations") - r_params = compact_mapping({"cursor": cursor, "limit": limit}) - - # Make the request. - obj = client.request( - "GET", r_path, params=r_params, response_model=schema.OperationsResponse - ) - - # Deserialize the response. - return [ - cls(client=client, obj=o, data_model=schema.Operation) - for o in obj.operations - ] - - @classmethod - def get( - cls, - client, - project_id: str, - operation_id: str, - ): - """Get a list of operations.""" - - # Construct the request path and parameters. - r_path = client.url_join("projects", project_id, "operations", operation_id) - - # Make the request. - obj = client.request("GET", r_path, response_model=schema.OperationResponse) - - # Deserialize the response. - return cls(client=client, obj=obj.operation, data_model=schema.Operation) - - class NeonAPI: def __init__(self, api_key: str, *, base_url=None): """A Neon API client.""" @@ -353,6 +40,7 @@ class NeonAPI: method: str, path: str, response_model: BaseModel | None = None, + response_is_array=False, **kwargs, ): """Send an HTTP request to the specified path using the specified method.""" @@ -375,11 +63,7 @@ class NeonAPI: except requests.exceptions.HTTPError: raise NeonClientException(r.text) - # Deserialize the response. - if response_model: - return response_model.model_construct(**r.json()) - else: - return r.json() + return r.json() def url_join(self, *args): """Join multiple URL path components.""" @@ -392,32 +76,168 @@ class NeonAPI: return cls(os.environ[NEON_API_KEY_ENVIRON]) - def me(self): + def me(self) -> t.Dict[str, t.Any]: """Get the current user.""" - return User.get_current_user_info(client=self) + return self.request("GET", "users/me") - def api_keys(self): + def api_keys(self) -> t.List[t.Dict[str, t.Any]]: """Get a list of API keys.""" - return APIKey.get_list(client=self) + return self.request("GET", "api_keys", response_is_array=True) - def projects(self, *, shared: bool = False, **kwargs): - """Get a list of projects.""" - return Project.get_list(client=self, shared=shared, **kwargs) - - def project(self, project_id: str): - """Get a project.""" - return Project.get(client=self, project_id=project_id) - - def branches(self, project_id: str): - """Get a list of branches.""" - return Branch.list(client=self, project_id=project_id) - - def operations(self, project_id: str): - """Get a list of operations.""" - return Operation.list(client=self, project_id=project_id) - - def operation(self, project_id: str, operation_id: str): - """Get an operation.""" - return Operation.get( - client=self, project_id=project_id, operation_id=operation_id + def api_key_create(self, key_name: str): + """Create a new API key.""" + return self.request( + "POST", "api_keys", json=jsonable_encoder({"name": key_name}) + ) + + 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 }") + + def projects( + self, + *, + shared: bool = False, + cursor: str | None = None, + limit: int | None = None, + ) -> t.List[t.Dict[str, t.Any]]: + """Get a list of 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, response_is_array=True) + + def project(self, project_id: str) -> t.Dict[str, t.Any]: + """Get a project.""" + + r_path = self.url_join("projects", project_id) + + return self.request( + "GET", + r_path, + response_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=jsonable_encoder(json)) + + def project_delete(self, project_id: str) -> t.Dict[str, t.Any]: + """Delete a project.""" + + return self.request("DELETE", f"projects/{ project_id }") + + def branches( + self, + project_id: str, + *, + cursor: str | None = None, + limit: int | None = None, + ): + """Get a list of branches.""" + + # Construct the request path and parameters. + r_path = self.url_join("projects", project_id, "branches") + r_params = compact_mapping({"cursor": cursor, "limit": limit}) + + # Make the request. + return self.request("GET", r_path, params=r_params) + + def branch(self, project_id: str, branch_id: str) -> t.Dict[str, t.Any]: + """Get a branch.""" + + # Construct the request path. + r_path = self.url_join("projects", project_id, "branches", branch_id) + + # Make the request. + return self.request("GET", r_path) + + 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=jsonable_encoder(json) + ) + + def branch_delete(self, project_id: str, branch_id: str) -> t.Dict[str, t.Any]: + """Delete a branch.""" + + return self.request("DELETE", f"projects/{ project_id }/branches/{ branch_id }") + + def databases( + self, + project_id: str, + branch_id: str, + *, + cursor: str | None = None, + limit: int | None = None, + ) -> t.List[t.Dict[str, t.Any]]: + """Get a list of databases.""" + + # Construct the request path and parameters. + r_path = self.url_join( + "projects", project_id, "branches", branch_id, "databases" + ) + r_params = compact_mapping({"cursor": cursor, "limit": limit}) + + # Make the request. + return self.request("GET", r_path, params=r_params) + + def database( + self, project_id: str, branch_id: str, database_id: str + ) -> t.Dict[str, t.Any]: + """Get a database.""" + + # Construct the request path. + r_path = self.url_join( + "projects", project_id, "branches", branch_id, "databases", database_id + ) + + # Make the request. + return self.request("GET", r_path) + + def database_create( + self, project_id: str, branch_id: str, **json: dict + ) -> t.Dict[str, t.Any]: + """Create a new database. Accepts all keyword arguments for json body.""" + + return self.request( + "POST", + f"projects/{ project_id }/branches/{ branch_id }/databases", + json=jsonable_encoder(json), + ) + + def database_update( + self, project_id: str, branch_id: str, database_id: str, **json: dict + ) -> t.Dict[str, t.Any]: + """Update a database. Accepts all keyword arguments for json body.""" + + return self.request( + "PUT", + f"projects/{ project_id }/branches/{ branch_id }/databases/{ database_id }", + json=jsonable_encoder(json), + ) + + def database_delete( + self, project_id: str, branch_id: str, database_id: str + ) -> t.Dict[str, t.Any]: + """Delete a database.""" + + return self.request( + "DELETE", + f"projects/{ project_id }/branches/{ branch_id }/databases/{ database_id }", + ) + + def operations(self, project_id: str) -> t.List[t.Dict[str, t.Any]]: + """Get a list of operations.""" + + return self.request("GET", f"projects/{ project_id }/operations") + + def operation(self, project_id: str, operation_id: str) -> t.Dict[str, t.Any]: + """Get an operation.""" + + return self.request( + "GET", f"projects/{ project_id }/operations/{ operation_id }" )