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,
)
from .utils import validate_obj_model


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."""

    path = "api_keys"
    response_model = ApiKeysListResponseItem
    # response_model_single =

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

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

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

        return self.api.request(
            method="POST",
            path=self.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.path, key_id),
            response_model=ApiKeyRevokeResponse,
        )

    def get(self, key_id: str):
        """Get an API key."""

        return self.api.request(
            "GET",
            f"api_keys/{key_id}",
            response_model=ApiKeyRevokeResponse,
        )


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

    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.path, "shared") if shared else self.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.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.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.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.path, project_id),
            response_model=self.response_model_single,
        )


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

    CreateRequest = DatabaseCreateRequest
    UpdateRequest = DatabaseUpdateRequest
    Database = Database

    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
        """

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

        return databases_response.databases

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

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

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

        db = self._extract_database(db)

        database_create_response = 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,
        )
        return database_create_response.database

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

        db = self._extract_database(db)

        database_update_response = 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,
        )
        return database_update_response.database


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."""

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

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

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

        return branch_response.branch


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)