mass simplification

This commit is contained in:
2024-01-19 16:28:17 -05:00
parent a861c24882
commit 5860a94d2b
+162 -342
View File
@@ -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 }"
)