import os
import typing as t
from datetime import datetime
from functools import wraps
import requests
from . import schema
from .utils import compact_mapping, to_iso8601
from .exceptions import NeonAPIError
__VERSION__ = "0.1.0"
NEON_API_KEY_ENVIRON = "NEON_API_KEY"
NEON_API_BASE_URL = "https://console.neon.tech/api/v2/"
ENABLE_PYDANTIC = True
def returns_model(model, is_array=False):
"""Decorator that returns a Pydantic dataclass.
:param model: The Pydantic dataclass to return.
:param is_array: Whether the return value is an array (default is False).
:return: A Pydantic dataclass.
If Pydantic is not enabled, the original return value is returned.
"""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
if not ENABLE_PYDANTIC:
return func(*args, **kwargs)
if is_array:
return [model(**item) for item in func(*args, **kwargs)]
else:
return model(**func(*args, **kwargs))
return wrapper
return decorator
def returns_subkey(key):
"""Decorator that returns a subkey.
:param key: The key to return.
:return: The value of the key in the return value.
"""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
try:
return getattr(func(*args, **kwargs), key)
except AttributeError:
return func(*args, **kwargs)[key]
return wrapper
return decorator
[docs]
class NeonAPI:
def __init__(self, api_key: str, *, base_url: str = None):
"""A Neon API client.
:param api_key: The API key to use for authentication.
:param base_url: The base URL of the Neon API (default is https://console.neon.tech/api/v2/).
"""
# Set the base URL.
if not base_url:
base_url = NEON_API_BASE_URL
# Private attributes.
self._api_key = api_key
self._session = requests.Session()
# Public attributes.
self.base_url = base_url
self.user_agent = f"neon-client/python version=({__VERSION__})"
def __repr__(self):
return f"<NeonAPI base_url={self.base_url!r}>"
def _request(
self,
method: str,
path: str,
**kwargs,
):
"""Send an HTTP request to the specified API path using the specified method.
:param method: The HTTP method to use (e.g., "GET", "POST", "PUT", "DELETE").
:param path: The API path to send the request to.
:param kwargs: Additional keyword arguments to pass to the requests.Session.request method.
:return: The JSON response from the server.
"""
# Set HTTP headers for outgoing requests.
headers = kwargs.pop("headers", {})
headers["Authorization"] = f"Bearer {self._api_key}"
headers["Accept"] = "application/json"
headers["Content-Type"] = "application/json"
headers["User-Agent"] = self.user_agent
# Send the request.
r = self._session.request(
method, self.base_url + path, headers=headers, **kwargs
)
# Check the response status code.
try:
r.raise_for_status()
except requests.exceptions.HTTPError:
raise NeonAPIError(r.text)
return r.json()
def _url_join(self, *args):
"""Join a list of URL components into a single URL."""
return "/".join(args)
[docs]
@classmethod
def from_environ(cls):
"""Create a new Neon API client from the `NEON_API_KEY` environment variable."""
return cls(os.environ[NEON_API_KEY_ENVIRON])
[docs]
@returns_model(schema.CurrentUserInfoResponse)
def me(self) -> t.Dict[str, t.Any]:
"""Get the current user.
More info: https://api-docs.neon.tech/reference/getcurrentuserinfo
"""
return self._request("GET", "users/me")
[docs]
@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: A dataclass representing the API key.
More info: https://api-docs.neon.tech/reference/listapikeys
"""
return self._request("GET", "api_keys")
[docs]
@returns_model(schema.ApiKeyCreateResponse)
def api_key_create(self, **json: dict) -> t.Dict[str, t.Any]:
"""Create a new API key.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the API key.
Example usage:
>>> neon.api_key_create(name="My API Key")
More info: https://api-docs.neon.tech/reference/createapikey
"""
return self._request("POST", "api_keys", json=json)
[docs]
@returns_model(schema.ApiKeyRevokeResponse)
def api_key_revoke(self, api_key_id: str) -> t.Dict[str, t.Any]:
"""Revoke an API key.
:param api_key_id: The ID of the API key to revoke.
:return: A dataclass representing the API key.
More info: https://api-docs.neon.tech/reference/revokeapikey
"""
return self._request("DELETE", f"api_keys/{ api_key_id }")
[docs]
@returns_model(schema.ProjectsResponse)
def projects(
self,
*,
shared: bool = False,
cursor: str = None,
limit: int = None,
) -> t.List[t.Dict[str, t.Any]]:
"""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 dataclasses representing the projects.
More info: https://api-docs.neon.tech/reference/listprojects
"""
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)
[docs]
@returns_model(schema.ProjectResponse)
def project(self, project_id: str) -> t.Dict[str, t.Any]:
"""Get a project.
:param project_id: The ID of the project.
:return: A dataclass representing the project.
More info: https://api-docs.neon.tech/reference/getproject
"""
r_path = f"projects/{project_id}"
return self._request("GET", r_path)
[docs]
@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.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the project.
More info: https://api-docs.neon.tech/reference/createproject
"""
return self._request("POST", "projects", json=json)
[docs]
@returns_model(schema.ProjectResponse)
def project_update(self, project_id: str, **json: dict) -> t.Dict[str, t.Any]:
"""Updates a project. Accepts all keyword arguments for json body.
:param project_id: The ID of the project.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the project.
More info: https://api-docs.neon.tech/reference/updateproject"""
return self._request("PATCH", f"projects/{ project_id }", json=json)
[docs]
@returns_model(schema.ProjectResponse)
def project_delete(self, project_id: str) -> t.Dict[str, t.Any]:
"""Delete a project.
:param project_id: The ID of the project.
:return: A dataclass representing the project.
More info: https://api-docs.neon.tech/reference/deleteproject
"""
return self._request("DELETE", f"projects/{ project_id }")
[docs]
@returns_model(schema.ProjectPermissions)
def project_permissions(self, project_id: str) -> t.Dict[str, t.Any]:
"""Get a project permissions.
:param project_id: The ID of the project.
:return: A dataclass representing the project permissions.
More info: https://api-docs.neon.tech/reference/listprojectpermissions
"""
return self._request("GET", f"projects/{ project_id }/permissions")
[docs]
@returns_model(schema.ProjectPermission)
def project_permissions_grant(
self, project_id: str, **json: dict
) -> t.Dict[str, t.Any]:
"""Update a project permissions. Accepts all keyword arguments for json body.
:param project_id: The ID of the project.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the project permissions.
More info: https://api-docs.neon.tech/reference/grantpermissiontoproject
"""
return self._request("POST", f"projects/{ project_id }/permissions", json=json)
[docs]
@returns_model(schema.ProjectPermission)
def project_permissions_revoke(
self, project_id: str, **json: dict
) -> t.Dict[str, t.Any]:
"""Update a project permissions. Accepts all keyword arguments for json body.
:param project_id: The ID of the project.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the project permissions.
More info: https://api-docs.neon.tech/reference/revokepermissionfromproject
"""
return self._request(
"DELETE", f"projects/{ project_id }/permissions", json=json
)
[docs]
@returns_model(schema.BranchesResponse)
def branches(
self,
project_id: str,
*,
cursor: str = None,
limit: int = None,
) -> t.Dict[str, t.Any]:
"""Get a list of branches.
:param project_id: The ID of the project.
: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 dataclasses representing the projects.
More info: https://api-docs.neon.tech/reference/listprojectbranches
"""
# 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)
[docs]
@returns_model(schema.BranchResponse)
def branch(self, project_id: str, branch_id: str) -> t.Dict[str, t.Any]:
"""Get a branch.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:return: A dataclass representing the branch.
More info: https://api-docs.neon.tech/reference/getprojectbranch
"""
# Construct the request path.
r_path = self._url_join("projects", project_id, "branches", branch_id)
# Make the request.
return self._request("GET", r_path)
[docs]
@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.
:param project_id: The ID of the project.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the branch.
More info: https://api-docs.neon.tech/reference/createprojectbranch
"""
return self._request("POST", f"projects/{ project_id }/branches", json=json)
[docs]
@returns_model(schema.BranchOperations)
def branch_update(
self, project_id: str, branch_id: str, **json: dict
) -> t.Dict[str, t.Any]:
"""Update a branch by branch_id. Accepts all keyword arguments for json body.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the branch.
More info: https://api-docs.neon.tech/reference/updateprojectbranch
"""
return self._request(
"PATCH", f"projects/{ project_id }/branches/{ branch_id }", json=json
)
[docs]
@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.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:return: A dataclass representing the branch.
More info: https://api-docs.neon.tech/reference/deleteprojectbranch
"""
return self._request(
"DELETE", f"projects/{ project_id }/branches/{ branch_id }"
)
[docs]
@returns_model(schema.BranchOperations)
def branch_set_as_primary(
self, project_id: str, branch_id: str
) -> t.Dict[str, t.Any]:
"""Set a branch as primary by branch_id.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:return: A dataclass representing the branch.
More info: https://api-docs.neon.tech/reference/setprimaryprojectbranch"""
return self._request(
"POST", f"projects/{ project_id }/branches/{ branch_id }/set_as_primary"
)
[docs]
@returns_model(schema.DatabasesResponse)
def databases(
self,
project_id: str,
branch_id: str,
*,
cursor: str = None,
limit: int = None,
) -> t.List[t.Dict[str, t.Any]]:
"""Get a list of databases.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
: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 dataclasses representing the database.
More info: https://api-docs.neon.tech/reference/listprojectbranchdatabases
"""
# 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)
[docs]
@returns_model(schema.DatabaseResponse)
def database(
self, project_id: str, branch_id: str, database_id: str
) -> t.Dict[str, t.Any]:
"""Get a database.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param database_id: The ID of the database.
:return: A dataclass representing the database.
More info: https://api-docs.neon.tech/reference/getprojectbranchdatabase
"""
# 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)
[docs]
@returns_model(schema.DatabaseResponse)
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.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the database.
More info: https://api-docs.neon.tech/reference/createprojectbranchdatabase
"""
return self._request(
"POST",
f"projects/{ project_id }/branches/{ branch_id }/databases",
json=json,
)
[docs]
@returns_model(schema.DatabaseResponse)
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.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param database_id: The ID of the database.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the database.
More info: https://api-docs.neon.tech/reference/updateprojectbranchdatabase
"""
return self._request(
"PUT",
f"projects/{ project_id }/branches/{ branch_id }/databases/{ database_id }",
json=json,
)
[docs]
@returns_model(schema.DatabaseResponse)
def database_delete(
self, project_id: str, branch_id: str, database_id: str
) -> t.Dict[str, t.Any]:
"""Delete a database by database_id.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param database_id: The ID of the database.
:return: A dataclass representing the database.
More info: https://api-docs.neon.tech/reference/deleteprojectbranchdatabase
"""
return self._request(
"DELETE",
f"projects/{ project_id }/branches/{ branch_id }/databases/{ database_id }",
)
[docs]
@returns_model(schema.EndpointsResponse)
def endpoints(self, project_id: str) -> t.Dict[str, t.Any]:
"""Get a list of endpoints for a given branch
:param project_id: The ID of the project.
:return: A list of dataclasses representing the endpoints.
More info: https://api-docs.neon.tech/reference/listprojectendpoints
"""
return self._request("GET", f"projects/{ project_id }/endpoints")
[docs]
@returns_model(schema.EndpointResponse)
def endpoint(self, project_id: str, endpoint_id: str) -> t.Dict[str, t.Any]:
"""Get an endpoint for a given branch.
:param project_id: The ID of the project.
:param endpoint_id: The ID of the endpoint.
:return: A dataclass representing the endpoint.
More info: https://api-docs.neon.tech/reference/getprojectendpoint
"""
return self._request(
"GET",
f"projects/{ project_id }/endpoints/{ endpoint_id }",
)
[docs]
@returns_model(schema.EndpointOperations)
def endpoint_create(
self,
project_id: str,
**json: dict,
) -> t.Dict[str, t.Any]:
"""Create a new endpoint. Accepts all keyword arguments for json body.
:param project_id: The ID of the project.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the endpoint.
More info: https://api-docs.neon.tech/reference/createprojectendpoint
"""
return self._request("POST", f"projects/{ project_id }/endpoints", json=json)
[docs]
@returns_model(schema.EndpointOperations)
def endpoint_delete(self, project_id: str, endpoint_id: str) -> t.Dict[str, t.Any]:
"""Delete an endpoint by endpoint_id.
:param project_id: The ID of the project.
:param endpoint_id: The ID of the endpoint.
:return: A dataclass representing the endpoint.
More info: https://api-docs.neon.tech/reference/deleteprojectendpoint
"""
return self._request(
"DELETE",
f"projects/{ project_id }/endpoints/{ endpoint_id }",
)
[docs]
@returns_model(schema.EndpointOperations)
def endpoint_update(
self, project_id: str, endpoint_id: str, **json: dict
) -> t.Dict[str, t.Any]:
"""Update an endpoint. Accepts all keyword arguments for json body.
:param project_id: The ID of the project.
:param endpoint_id: The ID of the endpoint.
:param json: The JSON paypload to send to the server.
:return: A dataclass representing the endpoint.
More info: https://api-docs.neon.tech/reference/updateprojectendpoint
"""
return self._request(
"PATCH",
f"projects/{ project_id }/endpoints/{ endpoint_id }",
json=json,
)
[docs]
@returns_model(schema.EndpointOperations)
def endpoint_start(self, project_id: str, endpoint_id: str):
"""Start an endpoint by endpoint_id.
:param project_id: The ID of the project.
:param endpoint_id: The ID of the endpoint.
:return: A dataclass representing the endpoint.
More info: https://api-docs.neon.tech/reference/startprojectendpoint"""
return self._request(
"POST",
f"projects/{ project_id }/endpoints/{ endpoint_id }/start",
)
[docs]
@returns_model(schema.EndpointOperations)
def endpoint_suspend(self, project_id: str, endpoint_id: str):
"""Suspend an endpoint by endpoint_id.
:param project_id: The ID of the project.
:param endpoint_id: The ID of the endpoint.
:return: A dataclass representing the endpoint.
More info: https://api-docs.neon.tech/reference/suspendprojectendpoint
"""
return self._request(
"POST",
f"projects/{ project_id }/endpoints/{ endpoint_id }/suspend",
)
[docs]
@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.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:return: A list of dataclasses representing the roles.
More info: https://api-docs.neon.tech/reference/listprojectbranchroles
"""
return self._request(
"GET", f"projects/{ project_id }/branches/{ branch_id }/roles"
)
[docs]
@returns_model(schema.RoleResponse)
def role(
self, project_id: str, branch_id: str, role_name: str
) -> t.Dict[str, t.Any]:
"""Get a role for a given branch.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param role_name: The name of the role.
:return: A dataclass representing the role.
More info: https://api-docs.neon.tech/reference/getprojectbranchrole
"""
return self._request(
"GET", f"projects/{ project_id }/branches/{ branch_id }/roles/{ role_name }"
)
[docs]
@returns_model(schema.RoleOperations)
def role_create(
self,
project_id: str,
branch_id: str,
role_name: str,
) -> t.Dict[str, t.Any]:
"""Create a new role. Accepts all keyword arguments for json body.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param role_name: The name of the role.
:return: A dataclass representing the role.
More info: https://api-docs.neon.tech/reference/createprojectbranchrole
"""
return self._request(
"POST",
f"projects/{ project_id }/branches/{ branch_id }/roles",
json={"role": {"name": role_name}},
)
[docs]
@returns_model(schema.RoleOperations)
def role_delete(
self,
project_id: str,
branch_id: str,
role_name: str,
) -> t.Dict[str, t.Any]:
"""Delete a role by given role name.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param role_name: The name of the role.
:return: A dataclass representing the role.
More info: https://api-docs.neon.tech/reference/deleteprojectbranchrole
"""
return self._request(
"DELETE",
f"projects/{ project_id }/branches/{ branch_id }/roles/{ role_name }",
)
[docs]
@returns_model(schema.RolePasswordResponse)
def role_password_reveal(
self,
project_id: str,
branch_id: str,
role_name: str,
) -> t.Dict[str, t.Any]:
"""Get a role password.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param role_name: The name of the role.
:return: A dataclass representing the role password.
More info: https://api-docs.neon.tech/reference/getprojectbranchrolepassword"""
return self._request(
"POST",
f"projects/{ project_id }/branches/{ branch_id }/roles/{ role_name }/reveal_password",
)
[docs]
@returns_model(schema.RoleOperations)
def role_password_reset(
self,
project_id: str,
branch_id: str,
role_name: str,
) -> t.Dict[str, t.Any]:
"""Reset a role password.
:param project_id: The ID of the project.
:param branch_id: The ID of the branch.
:param role_name: The name of the role.
:return: A dataclass representing the role.
More info: https://api-docs.neon.tech/reference/resetprojectbranchrolepassword
"""
return self._request(
"POST",
f"projects/{ project_id }/branches/{ branch_id }/roles/{ role_name }/reset_password",
)
[docs]
@returns_model(schema.OperationsResponse)
def operations(
self,
project_id: str,
*,
cursor: str = None,
limit: int = None,
) -> t.Dict[str, t.Any]:
"""Get a list of operations.
:param project_id: The ID of the project.
: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 dataclasses representing the operations.
More info: https://api-docs.neon.tech/reference/listprojectoperations
"""
r_params = compact_mapping({"cursor": cursor, "limit": limit})
return self._request(
"GET", f"projects/{ project_id }/operations", params=r_params
)
[docs]
@returns_model(schema.OperationResponse)
def operation(self, project_id: str, operation_id: str) -> t.Dict[str, t.Any]:
"""Get an operation.
:param project_id: The ID of the project.
:param operation_id: The ID of the operation.
:return: A dataclass representing the operation.
More info: https://api-docs.neon.tech/reference/getprojectoperation
"""
return self._request(
"GET", f"projects/{ project_id }/operations/{ operation_id }"
)
[docs]
@returns_model(schema.ProjectsConsumptionResponse)
def consumption(
self,
*,
cursor: str = None,
limit: int = None,
from_date: datetime = None,
to_date: datetime = None,
) -> t.Dict[str, t.Any]:
"""Experimental — get a list of consumption metrics for all projects.
:param cursor: The cursor for pagination (default is None).
:param limit: The maximum number of projects to retrieve (default is None).
:param from_date: The start date for the consumption metrics (default is None).
:param to_date: The end date for the consumption metrics (default is None).
:return: A dataclass representing the consumption metrics.
More info: https://api-docs.neon.tech/reference/listprojectsconsumption
"""
# Convert datetime objects to ISO 8601 strings.
from_date = (
to_iso8601(from_date) if isinstance(from_date, datetime) else from_date
)
to_date = to_iso8601(to_date) if isinstance(to_date, datetime) else to_date
# Construct the request parameters.
r_params = compact_mapping(
{"cursor": cursor, "limit": limit, "from": from_date, "to": to_date}
)
# Make the request.
return self._request("GET", "consumption/projects", params=r_params)