neon-api-python/neon_client/resources.py

from typing import List

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


class Resource:
    base_path = None

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


class UserResource(Resource):
    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):
    path = "api_keys"
    response_model = ApiKeysListResponseItem

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

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

    def create_key(self, name: str):
        """Create a new API key."""

        return self.api.request(
            method="POST",
            path=self.path,
            json={"key_name": name},
            response_model=ApiKeyCreateResponse,
        )

    def revoke_key(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_key(self, key_id: str):
        """Get an API key."""

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


class ProjectResource(Resource):
    path = "projects"
    response_model = ProjectsResponse
    response_model_single = ProjectResponse

    def get_projects(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_project(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_project(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_project(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 new(self, project_id: str):
    #     """Create a new project."""

    #     project_create_response = self.api.request(
    #         method="POST",
    #         path=self.path,
    #         json={"project": {"name": project_id}},
    #         response_model=self.response_model,
    #     )
    #     return project_create_response.project

    # def get_project(self, project_id: str) -> ProjectResponse:
    #     """Get a project."""

    #     return self.api.request(
    #         "GET", f"projects/{project_id}", response_model=ProjectResponse
    #     )

    # def update_project(self, project: Project) -> ProjectResponse:
    #     """Update a project."""

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

    #     return self.api.request(
    #         "PATCH",
    #         f"projects/{project.id}",
    #         json={"project": payload.model_dump()},
    #         response_model=ProjectResponse,
    #     )

    # def delete_project(self, project_id: str) -> ProjectResponse:
    #     """Delete a project."""

    #     return self.api.request(
    #         "DELETE", f"projects/{project_id}", response_model=ProjectResponse
    #     )


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_databases(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_database(
        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_database(
        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_database(
        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