neon-api-python/neon_client/resources.py

from typing import List

from .http_client import Neon_API_V2
from .openapi_models import (
    ApiKeyCreateRequest,
    ApiKeyCreateResponse,
    ApiKeyRevokeResponse,
    ApiKeysListResponseItem,
    CurrentUserInfoResponse,
    Project,
    ProjectResponse,
    ProjectUpdateRequest,
    ProjectsResponse,
    DatabaseResponse,
    DatabasesResponse,
    DatabaseUpdateRequest,
    Database,
    Database1,
    Database2,
    DatabaseCreateRequest,
    BranchesResponse,
    BranchResponse,
    OperationResponse,
    OperationsResponse,
    PaginationResponse,
)
from .utils import validate_obj_model


class PagedOperationsResponse(OperationsResponse, PaginationResponse):
    """A response containing a list of operations and pagination information."""

    pass


class Resource:
    base_path = None

    def __init__(self, api: Neon_API_V2):
        self.api = api


class UserResource(Resource):
    """A resource for interacting with users."""

    base_path = "users"
    response_model = CurrentUserInfoResponse

    def get_current_user_info(self):
        """Get information about the user."""

        return self.api.request(
            method="GET",
            path=self.api.url_join(self.path, "me"),
            response_model=self.response_model,
        )


class APIKeyResource(Resource):
    """A resource for interacting with API keys."""

    base_path = "api_keys"

    def get_list(self):
        """Get a list of API keys."""

        return self.api.request(
            method="GET",
            path=self.base_path,
            response_model=ApiKeysListResponseItem,
            response_is_array=True,
        )

    def create(self, key_name: str):
        """Create a new API key."""

        return self.api.request(
            method="POST",
            path=self.base_path,
            json=ApiKeyCreateRequest(key_name=key_name).model_dump(),
            response_model=ApiKeyCreateResponse,
        )

    def revoke(self, key_id: str):
        """Revoke an API key."""

        return self.api.request(
            method="DELETE",
            path=self.api.url_join(self.base_path, str(key_id)),
            response_model=ApiKeyRevokeResponse,
        )


class ProjectResource(Resource):
    """A resource for interacting with projects."""

    base_path = "projects"
    response_model = ProjectsResponse
    response_model_single = ProjectResponse

    def get_list(self, *, shared=False):
        """Get a list of projects."""

        project_response = self.api.request(
            method="GET",
            path=(
                self.api.url_join(self.base_path, "shared")
                if shared
                else self.base_path
            ),
            response_model=self.response_model,
        )
        return project_response.projects

    def get(self, project_id: str):
        """Get a project."""

        project_response = self.api.request(
            method="GET",
            path=self.api.url_join(self.base_path, project_id),
            response_model=self.response_model_single,
        )
        return project_response.project

    def create(self, name: str, **kwargs):
        """Create a new project."""

        project_create_response = self.api.request(
            method="POST",
            path=self.base_path,
            json={"project": {"name": name, **kwargs}},
            response_model=self.response_model_single,
        )
        return project_create_response.project

    def update(self, project: Project):
        """Update a project."""

        payload = ProjectUpdateRequest(project=project.model_dump())

        return self.api.request(
            method="PATCH",
            path=self.api.url_join(self.base_path, project.id),
            json={"project": payload.model_dump()},
            response_model=self.response_model_single,
        )

    def delete(self, project_id: str):
        """Delete a project."""

        return self.api.request(
            method="DELETE",
            path=self.api.url_join(self.base_path, project_id),
            response_model=self.response_model_single,
        )


class DatabaseResource(Resource):
    base_path = "databases"
    response_model = DatabasesResponse
    response_model_single = DatabaseResponse

    def _extract_database(self, obj):
        """Extract a database from the specified object."""

        assert isinstance(obj, (DatabaseCreateRequest, Database1, Database2, Database))

        # Object mappings.
        if isinstance(obj, DatabaseCreateRequest):
            obj = obj.database.model_dump()
        if isinstance(obj, Database1):
            obj = obj.database.model_dump()
        if isinstance(obj, Database2):
            obj = obj.database.model_dump()
        if isinstance(obj, Database):
            obj = obj.model_dump()

        return obj

    def get_list(self, project_id: str, branch_id: str):
        """Get a list of databases.

        See also: https://api-docs.neon.tech/reference/listprojectbranchdatabases
        """

        return self.api.request(
            method="GET",
            path=self.api.url_join(
                "projects", project_id, "branches", branch_id, "databases"
            ),
            response_model=DatabasesResponse,
        ).model_dump()

    def get(
        self,
        project_id: str,
        branch_id: str,
        database_name: str,
    ):
        """Get a database."""

        return self.api.request(
            method="GET",
            path=self.api.url_join(
                "projects",
                project_id,
                "branches",
                branch_id,
                "databases",
                database_name,
            ),
            response_model=DatabaseResponse,
        ).model_dump()

    def create(
        self,
        project_id: str,
        branch_id: str,
        db: DatabaseCreateRequest | Database1 | Database2 | Database,
    ):
        """Create a new database."""

        db = self._extract_database(db)

        return self.api.request(
            method="POST",
            path=self.api.url_join(
                "projects", project_id, "branches", branch_id, "databases"
            ),
            json=DatabaseCreateRequest(database=db).model_dump(),
            response_model=DatabaseResponse,
        ).model_dump()

    def update(
        self,
        project_id: str,
        branch_id: str,
        database_id: str,
        db: DatabaseUpdateRequest | Database2,
    ):
        """Update a database."""

        db = self._extract_database(db)

        return self.api.request(
            method="PATCH",
            path=self.api.url_join(
                "projects", project_id, "branches", branch_id, "databases", database_id
            ),
            json=DatabaseUpdateRequest(database=db).model_dump(),
            response_model=DatabaseResponse,
        ).model_dump()


class BranchResource(Resource):
    """A resource for interacting with branches."""

    path = "branches"
    response_model = BranchesResponse
    response_model_single = BranchResponse

    def get_list(self, project_id: str):
        """Get a list of branches."""

        return self.api.request(
            method="GET",
            path=self.api.url_join("projects", project_id, "branches"),
            response_model=self.response_model,
        ).model_dump()

    def get(self, project_id: str, branch_id: str):
        """Get a branch."""

        return self.api.request(
            method="GET",
            path=self.api.url_join("projects", project_id, "branches", branch_id),
            response_model=self.response_model_single,
        ).model_dump()


class OperationResource(Resource):
    """A resource for interacting with operations."""

    base_path = "operations"

    def get_list(
        self,
        project_id: str,
        cursor: int | None = None,
        limit: int | None = None,
    ):
        """Get a list of operations."""

        operations_params = {}
        if cursor is not None:
            operations_params["cursor"] = cursor
        if limit is not None:
            operations_params["limit"] = limit

        return self.api.request(
            method="GET",
            path=self.api.url_join("projects", project_id, "operations"),
            params=operations_params,
            response_model=PagedOperationsResponse,
        ).model_dump()

    def get(self, project_id: str, operation_id: str):
        """Get an operation."""

        return self.api.request(
            method="GET",
            path=self.api.url_join("projects", project_id, "operations", operation_id),
            response_model=OperationResponse,
        ).model_dump()


class ResourceCollection:
    """A collection of resources."""

    def __init__(self, api: Neon_API_V2):
        # Initialize resources.
        self.api_keys = APIKeyResource(api)
        self.users = UserResource(api)
        self.projects = ProjectResource(api)
        self.databases = DatabaseResource(api)
        self.branches = BranchResource(api)
        self.operations = OperationResource(api)