vultr-api/vultr_api/client.py

"""
Vultr API v2 Client
"""

import requests
from typing import Optional, Dict, Any, List
from urllib.parse import urljoin

from .resources.account import AccountResource
from .resources.instances import InstancesResource
from .resources.dns import DNSResource
from .resources.firewall import FirewallResource
from .resources.ssh_keys import SSHKeysResource
from .resources.startup_scripts import StartupScriptsResource
from .resources.snapshots import SnapshotsResource
from .resources.block_storage import BlockStorageResource
from .resources.reserved_ips import ReservedIPsResource
from .resources.vpc import VPCResource
from .resources.load_balancers import LoadBalancersResource
from .resources.bare_metal import BareMetalResource
from .resources.plans import PlansResource
from .resources.regions import RegionsResource
from .resources.os import OSResource
from .resources.backups import BackupsResource


class VultrAPIError(Exception):
    """Vultr API Error"""
    def __init__(self, message: str, status_code: int = None, response: dict = None):
        self.message = message
        self.status_code = status_code
        self.response = response
        super().__init__(self.message)


class VultrClient:
    """
    Vultr API v2 Client

    Usage:
        client = VultrClient(api_key="your-api-key")

        # Get account info
        account = client.account.get()

        # List instances
        instances = client.instances.list()

        # Create instance
        instance = client.instances.create(
            region="ewr",
            plan="vc2-1c-1gb",
            os_id=215
        )
    """

    BASE_URL = "https://api.vultr.com/v2/"

    def __init__(self, api_key: str, timeout: int = 30):
        """
        Initialize Vultr API client

        Args:
            api_key: Vultr API key (get from https://my.vultr.com/settings/#settingsapi)
            timeout: Request timeout in seconds
        """
        self.api_key = api_key
        self.timeout = timeout
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        })

        # Initialize resources
        self.account = AccountResource(self)
        self.instances = InstancesResource(self)
        self.dns = DNSResource(self)
        self.firewall = FirewallResource(self)
        self.ssh_keys = SSHKeysResource(self)
        self.startup_scripts = StartupScriptsResource(self)
        self.snapshots = SnapshotsResource(self)
        self.block_storage = BlockStorageResource(self)
        self.reserved_ips = ReservedIPsResource(self)
        self.vpc = VPCResource(self)
        self.load_balancers = LoadBalancersResource(self)
        self.bare_metal = BareMetalResource(self)
        self.plans = PlansResource(self)
        self.regions = RegionsResource(self)
        self.os = OSResource(self)
        self.backups = BackupsResource(self)

    def _request(
        self,
        method: str,
        endpoint: str,
        params: Optional[Dict] = None,
        data: Optional[Dict] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Make API request

        Args:
            method: HTTP method (GET, POST, PATCH, DELETE)
            endpoint: API endpoint (e.g., "instances")
            params: Query parameters
            data: Request body data

        Returns:
            Response JSON or None for 204 responses
        """
        url = urljoin(self.BASE_URL, endpoint)

        try:
            response = self.session.request(
                method=method,
                url=url,
                params=params,
                json=data,
                timeout=self.timeout,
            )

            if response.status_code == 204:
                return None

            if response.status_code >= 400:
                error_data = None
                try:
                    error_data = response.json()
                except:
                    pass

                error_msg = f"API Error: {response.status_code}"
                if error_data and "error" in error_data:
                    error_msg = error_data["error"]

                raise VultrAPIError(
                    message=error_msg,
                    status_code=response.status_code,
                    response=error_data,
                )

            if response.text:
                return response.json()
            return None

        except requests.RequestException as e:
            raise VultrAPIError(f"Request failed: {str(e)}")

    def get(self, endpoint: str, params: Optional[Dict] = None) -> Optional[Dict]:
        """GET request"""
        return self._request("GET", endpoint, params=params)

    def post(self, endpoint: str, data: Optional[Dict] = None) -> Optional[Dict]:
        """POST request"""
        return self._request("POST", endpoint, data=data)

    def patch(self, endpoint: str, data: Optional[Dict] = None) -> Optional[Dict]:
        """PATCH request"""
        return self._request("PATCH", endpoint, data=data)

    def put(self, endpoint: str, data: Optional[Dict] = None) -> Optional[Dict]:
        """PUT request"""
        return self._request("PUT", endpoint, data=data)

    def delete(self, endpoint: str) -> None:
        """DELETE request"""
        self._request("DELETE", endpoint)

    def paginate(
        self,
        endpoint: str,
        key: str,
        params: Optional[Dict] = None,
        per_page: int = 100,
    ) -> List[Dict]:
        """
        Paginate through all results

        Args:
            endpoint: API endpoint
            key: Response key containing items (e.g., "instances")
            params: Additional query parameters
            per_page: Items per page (max 500)

        Returns:
            List of all items
        """
        params = params or {}
        params["per_page"] = per_page

        all_items = []
        cursor = None

        while True:
            if cursor:
                params["cursor"] = cursor

            response = self.get(endpoint, params=params)

            if not response:
                break

            items = response.get(key, [])
            all_items.extend(items)

            meta = response.get("meta", {})
            links = meta.get("links", {})
            cursor = links.get("next", "")

            if not cursor:
                break

        return all_items