From 681a9c732e983b82bf1a82d0f334d271193e7794 Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Wed, 31 Jan 2024 10:00:59 -0500 Subject: [PATCH] Refactor code and add new functions for handling datetime conversion --- Makefile | 2 +- neon_client/client.py | 28 ++++++- neon_client/schema.py | 182 +++++++++++++++++++++--------------------- neon_client/utils.py | 21 ++++- 4 files changed, 139 insertions(+), 94 deletions(-) diff --git a/Makefile b/Makefile index c3ce9f3..6715875 100644 --- a/Makefile +++ b/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 \ diff --git a/neon_client/client.py b/neon_client/client.py index 2800aff..4768ba0 100644 --- a/neon_client/client.py +++ b/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) diff --git a/neon_client/schema.py b/neon_client/schema.py index 126a744..09934e6 100644 --- a/neon_client/schema.py +++ b/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 diff --git a/neon_client/utils.py b/neon_client/utils.py index 729978a..cec41cc 100644 --- a/neon_client/utils.py +++ b/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 `_ string. + """ + + return dt.strftime("%Y-%m-%dT%H:%M:%SZ") + + +def from_iso8601(s): + """Convert an `ISO 8601 `_ + string to a datetime object. + """ + + return datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ")