Refactor code and add new functions for handling datetime conversion

Makefile

@@ -9,7 +9,7 @@ ci:
 record:
 	pytest --record-mode=rewrite tests/
 
-gen-model: fetch-v2-schema
+schema: fetch-v2-schema
 	datamodel-codegen \
 	--input v2.json \
 	--collapse-root-models \

neon_client/client.py

@@ -1,10 +1,11 @@
 import os
+from datetime import datetime
 import typing as t
 
 import requests
 
 from . import schema
-from .utils import compact_mapping
+from .utils import compact_mapping, to_iso8601
 from .exceptions import NeonClientException
 
 
@@ -479,3 +480,28 @@ class NeonAPI:
         return self.request(
             "GET", f"projects/{ project_id }/operations/{ operation_id }"
         )
+
+    # @returns_model(schema.ProjectsConsumptionResponse)
+    def consumption(
+        self,
+        *,
+        cursor: str | None = None,
+        limit: int | None = None,
+        from_date: datetime | str | None = None,
+        to_date: datetime | str | None = None,
+    ) -> t.Dict[str, t.Any]:
+        """Get a list of consumption metrics for all projects."""
+
+        # Convert datetime objects to ISO 8601 strings.
+        from_date = (
+            to_iso8601(from_date) if isinstance(from_date, datetime) else from_date
+        )
+        to_date = to_iso8601(to_date) if isinstance(to_date, datetime) else to_date
+
+        # Construct the request parameters.
+        r_params = compact_mapping(
+            {"cursor": cursor, "limit": limit, "from": from_date, "to": to_date}
+        )
+
+        # Make the request.
+        return self.request("GET", "consumption/projects", params=r_params)

neon_client/schema.py

@@ -1,10 +1,10 @@
 # generated by datamodel-codegen:
 #   filename:  v2.json
-#   timestamp: 2024-01-30T15:09:41+00:00
+#   timestamp: 2024-01-31T14:58:13+00:00
 
 from __future__ import annotations
 
-# from dataclasses import dataclass
+from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
 from typing import Optional
@@ -13,8 +13,8 @@ from pydantic.dataclasses import dataclass
 
 
 class Provisioner(Enum):
-    k8s_pod = "k8s-pod"
-    k8s_neonvm = "k8s-neonvm"
+    k8s_pod = 'k8s-pod'
+    k8s_neonvm = 'k8s-neonvm'
 
 
 @dataclass
@@ -37,7 +37,7 @@ class ApiKeyCreateResponse:
     id: int
     key: str
     name: str
-    created_at: datetime
+    created_at: str
 
 
 @dataclass
@@ -46,41 +46,41 @@ class ApiKeyRevokeResponse:
     name: str
     revoked: bool
     last_used_from_addr: str
-    last_used_at: Optional[datetime] = None
+    last_used_at: Optional[str] = None
 
 
 @dataclass
 class ApiKeysListResponseItem:
     id: int
     name: str
-    created_at: datetime
+    created_at: str
     last_used_from_addr: str
-    last_used_at: Optional[datetime] = None
+    last_used_at: Optional[str] = None
 
 
 class OperationAction(Enum):
-    create_compute = "create_compute"
-    create_timeline = "create_timeline"
-    start_compute = "start_compute"
-    suspend_compute = "suspend_compute"
-    apply_config = "apply_config"
-    check_availability = "check_availability"
-    delete_timeline = "delete_timeline"
-    create_branch = "create_branch"
-    tenant_ignore = "tenant_ignore"
-    tenant_attach = "tenant_attach"
-    tenant_detach = "tenant_detach"
-    tenant_reattach = "tenant_reattach"
-    replace_safekeeper = "replace_safekeeper"
-    disable_maintenance = "disable_maintenance"
-    apply_storage_config = "apply_storage_config"
+    create_compute = 'create_compute'
+    create_timeline = 'create_timeline'
+    start_compute = 'start_compute'
+    suspend_compute = 'suspend_compute'
+    apply_config = 'apply_config'
+    check_availability = 'check_availability'
+    delete_timeline = 'delete_timeline'
+    create_branch = 'create_branch'
+    tenant_ignore = 'tenant_ignore'
+    tenant_attach = 'tenant_attach'
+    tenant_detach = 'tenant_detach'
+    tenant_reattach = 'tenant_reattach'
+    replace_safekeeper = 'replace_safekeeper'
+    disable_maintenance = 'disable_maintenance'
+    apply_storage_config = 'apply_storage_config'
 
 
 class OperationStatus(Enum):
-    running = "running"
-    finished = "finished"
-    failed = "failed"
-    scheduling = "scheduling"
+    running = 'running'
+    finished = 'finished'
+    failed = 'failed'
+    scheduling = 'scheduling'
 
 
 @dataclass
@@ -94,8 +94,8 @@ class Branch:
 class ProjectPermission:
     id: str
     granted_to_email: str
-    granted_at: datetime
-    revoked_at: Optional[datetime] = None
+    granted_at: str
+    revoked_at: Optional[str] = None
 
 
 @dataclass
@@ -118,16 +118,16 @@ class ProjectConsumption:
     written_data_bytes: int
     compute_time_seconds: int
     active_time_seconds: int
-    updated_at: datetime
-    period_start: datetime
-    period_end: datetime
+    updated_at: str
+    period_start: str
+    period_end: str
     previous_period_id: str
-    data_storage_bytes_hour_updated_at: Optional[datetime] = None
-    synthetic_storage_size_updated_at: Optional[datetime] = None
-    data_transfer_bytes_updated_at: Optional[datetime] = None
-    written_data_bytes_updated_at: Optional[datetime] = None
-    compute_time_seconds_updated_at: Optional[datetime] = None
-    active_time_seconds_updated_at: Optional[datetime] = None
+    data_storage_bytes_hour_updated_at: Optional[str] = None
+    synthetic_storage_size_updated_at: Optional[str] = None
+    data_transfer_bytes_updated_at: Optional[str] = None
+    written_data_bytes_updated_at: Optional[str] = None
+    compute_time_seconds_updated_at: Optional[str] = None
+    active_time_seconds_updated_at: Optional[str] = None
 
 
 @dataclass
@@ -148,8 +148,8 @@ class ProjectLimits:
 
 
 class BranchState(Enum):
-    init = "init"
-    ready = "ready"
+    init = 'init'
+    ready = 'ready'
 
 
 @dataclass
@@ -186,18 +186,18 @@ class ConnectionDetails:
 
 
 class EndpointState(Enum):
-    init = "init"
-    active = "active"
-    idle = "idle"
+    init = 'init'
+    active = 'active'
+    idle = 'idle'
 
 
 class EndpointType(Enum):
-    read_only = "read_only"
-    read_write = "read_write"
+    read_only = 'read_only'
+    read_write = 'read_write'
 
 
 class EndpointPoolerMode(Enum):
-    transaction = "transaction"
+    transaction = 'transaction'
 
 
 @dataclass
@@ -240,8 +240,8 @@ class ExplainData:
 class Role:
     branch_id: str
     name: str
-    created_at: datetime
-    updated_at: datetime
+    created_at: str
+    updated_at: str
     password: Optional[str] = None
     protected: Optional[bool] = None
 
@@ -272,14 +272,14 @@ class RolePasswordResponse:
 
 
 class Brand(Enum):
-    amex = "amex"
-    diners = "diners"
-    discover = "discover"
-    jcb = "jcb"
-    mastercard = "mastercard"
-    unionpay = "unionpay"
-    unknown = "unknown"
-    visa = "visa"
+    amex = 'amex'
+    diners = 'diners'
+    discover = 'discover'
+    jcb = 'jcb'
+    mastercard = 'mastercard'
+    unionpay = 'unionpay'
+    unknown = 'unknown'
+    visa = 'visa'
 
 
 @dataclass
@@ -307,14 +307,14 @@ class BillingAccountUpdateRequest:
 
 
 class BillingSubscriptionType(Enum):
-    UNKNOWN = "UNKNOWN"
-    free = "free"
-    pro = "pro"
-    direct_sales = "direct_sales"
-    aws_marketplace = "aws_marketplace"
-    free_v2 = "free_v2"
-    launch = "launch"
-    scale = "scale"
+    UNKNOWN = 'UNKNOWN'
+    free = 'free'
+    pro = 'pro'
+    direct_sales = 'direct_sales'
+    aws_marketplace = 'aws_marketplace'
+    free_v2 = 'free_v2'
+    launch = 'launch'
+    scale = 'scale'
 
 
 @dataclass
@@ -323,8 +323,8 @@ class Database:
     branch_id: str
     name: str
     owner_name: str
-    created_at: datetime
-    updated_at: datetime
+    created_at: str
+    updated_at: str
 
 
 @dataclass
@@ -399,10 +399,10 @@ class ProjectOwnerData:
 
 
 class SupportTicketSeverity(Enum):
-    low = "low"
-    normal = "normal"
-    high = "high"
-    critical = "critical"
+    low = 'low'
+    normal = 'normal'
+    high = 'high'
+    critical = 'critical'
 
 
 @dataclass
@@ -417,13 +417,13 @@ class Operation:
     action: OperationAction
     status: OperationStatus
     failures_count: int
-    created_at: datetime
-    updated_at: datetime
+    created_at: str
+    updated_at: str
     total_duration_ms: int
     branch_id: Optional[str] = None
     endpoint_id: Optional[str] = None
     error: Optional[str] = None
-    retry_at: Optional[datetime] = None
+    retry_at: Optional[str] = None
 
 
 @dataclass
@@ -462,14 +462,14 @@ class Branch1:
     active_time_seconds: int
     written_data_bytes: int
     data_transfer_bytes: int
-    created_at: datetime
-    updated_at: datetime
+    created_at: str
+    updated_at: str
     parent_id: Optional[str] = None
     parent_lsn: Optional[str] = None
     parent_timestamp: Optional[str] = None
     pending_state: Optional[BranchState] = None
     logical_size: Optional[int] = None
-    last_reset_at: Optional[datetime] = None
+    last_reset_at: Optional[str] = None
 
 
 @dataclass
@@ -509,7 +509,7 @@ class StatementResult:
 class BillingAccount:
     payment_source: PaymentSource
     subscription_type: BillingSubscriptionType
-    quota_reset_at_last: datetime
+    quota_reset_at_last: str
     email: str
     address_city: str
     address_country: str
@@ -594,15 +594,15 @@ class ProjectListItem:
     active_time: int
     cpu_used_sec: int
     creation_source: str
-    created_at: datetime
-    updated_at: datetime
+    created_at: str
+    updated_at: str
     owner_id: str
     default_endpoint_settings: Optional[DefaultEndpointSettings] = None
     settings: Optional[ProjectSettingsData] = None
-    maintenance_starts_at: Optional[datetime] = None
+    maintenance_starts_at: Optional[str] = None
     synthetic_storage_size: Optional[int] = None
-    quota_reset_at: Optional[datetime] = None
-    compute_last_active_at: Optional[datetime] = None
+    quota_reset_at: Optional[str] = None
+    compute_last_active_at: Optional[str] = None
 
 
 @dataclass
@@ -625,18 +625,18 @@ class Project:
     store_passwords: bool
     creation_source: str
     history_retention_seconds: int
-    created_at: datetime
-    updated_at: datetime
-    consumption_period_start: datetime
-    consumption_period_end: datetime
+    created_at: str
+    updated_at: str
+    consumption_period_start: str
+    consumption_period_end: str
     owner_id: str
     default_endpoint_settings: Optional[DefaultEndpointSettings] = None
     settings: Optional[ProjectSettingsData] = None
-    maintenance_starts_at: Optional[datetime] = None
+    maintenance_starts_at: Optional[str] = None
     synthetic_storage_size: Optional[int] = None
-    quota_reset_at: Optional[datetime] = None
+    quota_reset_at: Optional[str] = None
     owner: Optional[ProjectOwnerData] = None
-    compute_last_active_at: Optional[datetime] = None
+    compute_last_active_at: Optional[str] = None
 
 
 @dataclass
@@ -699,8 +699,8 @@ class Endpoint:
     disabled: bool
     passwordless_access: bool
     creation_source: str
-    created_at: datetime
-    updated_at: datetime
+    created_at: str
+    updated_at: str
     proxy_host: str
     suspend_timeout_seconds: int
     provisioner: Provisioner

neon_client/utils.py

@@ -1,4 +1,23 @@
+import datetime
+
+
 def compact_mapping(obj):
-    """Compact a mapping by removing None values."""
+    """Compact a dict/mapping by removing all None values."""
 
     return {k: v for k, v in obj.items() if v is not None}
+
+
+def to_iso8601(dt):
+    """Convert a datetime object to an
+    `ISO 8601 <https://www.iso.org/iso-8601-date-and-time-format.html>`_ string.
+    """
+
+    return dt.strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def from_iso8601(s):
+    """Convert an `ISO 8601 <https://www.iso.org/iso-8601-date-and-time-format.html>`_
+    string to a datetime object.
+    """
+
+    return datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ")