neon-api-python/neon_client/client.py

from .http_client import Neon_API_V2
from .resources import ResourceCollection

from . import models


class BaseNeonItem:
    def install_neon_client(self, neon):
        self.neon = neon

    def __repr__(self):
        return str(self)


class NeonUser(BaseNeonItem, models.CurrentUserAuthAccount):
    @classmethod
    def from_get_response(cls, r):
        """Create a NeonUser from an API response."""

        # TODO: is this the right way to do this?
        me = r.auth_accounts[0]

        return cls.model_validate(me.model_dump())

    def __str__(self):
        return f"<NeonUser email={self.email}>"


class NeonProject(BaseNeonItem, models.ProjectListItem):
    """A Neon project."""

    @classmethod
    def from_list_response(cls, r):
        """Create a list of NeonProjects from an API response."""

        def gen():
            for project in r.projects:
                p = cls.model_validate(project.model_dump())

                yield p

        return [g for g in gen()]

    @classmethod
    def from_get_response(cls, r):
        """Create a NeonProject from an API response."""

        return cls.model_validate(r.model_dump())

    def __str__(self):
        return f"<NeonProject id={self.id} name={self.name}>"

    def delete(self):
        """Delete this project."""

        return bool(self.neon.resources.projects.delete(self.id))


class ItemView:
    """A view into a single item."""

    def __init__(self, item, key_id=None, neon=None):
        self.neon = neon
        self._item = item
        self._key_id = key_id

    @property
    def item(self):
        if self._key_id:
            return getattr(self._item, self._key_id)

        return self._item

    def __getattr__(self, name):
        return getattr(self.item, name)

    def __setattr__(self, name, value):
        if name == "_item":
            return super().__setattr__(name, value)

        return setattr(self.item, name, value)

    def __str__(self):
        return str(self.item)

    def __repr__(self):
        return repr(self.item)

    def __eq__(self, other):
        return self.item == other

    def __ne__(self, other):
        return self.item != other


class CollectionView:
    """A view into a collection of items."""

    def __init__(
        self,
        collection,
        key_ids=None,
        collection_id=None,
        neon=None,
    ):
        self.neon = neon
        self.pagination = None

        if not key_ids:
            key_ids = []

        self._key_ids = key_ids
        if collection_id:
            self._collection = getattr(collection, collection_id)
            self.pagination = collection.pagination
        else:
            self._collection = collection

    def __iter__(self):
        return iter(self._collection)

    def __getitem__(self, key):
        for k in self._key_ids:
            for item in self._collection:
                if str(getattr(item, k)) == str(key):
                    return item

        return self._collection[key]

    def __len__(self):
        return len(self._collection)

    def __repr__(self):
        return repr(self._collection)

    def __str__(self):
        return str(self._collection)

    def __contains__(self, item):
        return item in self._collection

    def __eq__(self, other):
        return self._collection == other

    def __ne__(self, other):
        return self._collection != other


class NeonClient:
    def __init__(self, api_key: str, **kwargs):
        self.api = Neon_API_V2(api_key, **kwargs)
        self.resources = ResourceCollection(self.api)

    def me(self):
        return self.resources.users.get_current_user_info()

    def api_keys(self):
        return CollectionView(self.resources.api_keys.get_list(), key_ids=["id"])

    def projects(self, **kwargs):
        return CollectionView(
            self.resources.projects.get_list(**kwargs),
            key_ids=["id", "name"],
            collection_id="projects",
        )

    def project(self, project_id: str, **kwargs):
        return self.resources.projects.get(project_id, **kwargs)

    def project_create(self, **kwargs):
        return self.resources.projects.create(**kwargs)

    def project_delete(self, project_id: str, **kwargs):
        return self.resources.projects.delete(project_id, **kwargs)

    def databases(self, project_id: str, branch_id: str, **kwargs):
        return CollectionView(
            self.resources.databases.get_list(project_id, branch_id, **kwargs),
            key_ids=["id"],
            collection_id="databases",
        )

    def database(self, project_id: str, database_id: str, **kwargs):
        return self.resources.databases.get(project_id, database_id, **kwargs)

    def branches(self, project_id: str, **kwargs):
        return CollectionView(
            self.resources.branches.get_list(project_id, **kwargs),
            key_ids=["id", "name"],
        )

    def branch(self, project_id: str, branch_id: str):
        return self.resources.branches.get(project_id, branch_id)