airbyte.cloud.organizations

  1# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
  2"""PyAirbyte classes and methods for Airbyte Cloud organizations."""
  3
  4from __future__ import annotations
  5
  6import logging
  7from typing import Any
  8
  9import requests
 10
 11from airbyte._util import api_util
 12from airbyte.cloud._credentials import _AirbyteCredentials
 13from airbyte.cloud.models import CloudOrganizationBillingInfo
 14from airbyte.exceptions import AirbyteError
 15from airbyte.secrets.base import SecretString
 16
 17
 18logger = logging.getLogger(__name__)
 19
 20
 21class CloudOrganization:
 22    """Information about an organization in Airbyte Cloud.
 23
 24    This class provides lazy loading of organization attributes including billing status.
 25    It is typically created via `CloudWorkspace.get_organization()`.
 26    """
 27
 28    def __init__(
 29        self,
 30        organization_id: str,
 31        organization_name: str | None = None,
 32        email: str | None = None,
 33        *,
 34        client_id: str | SecretString | None = None,
 35        client_secret: str | SecretString | None = None,
 36        bearer_token: str | SecretString | None = None,
 37        public_api_root: str | None = None,
 38        config_api_root: str | None = None,
 39    ) -> None:
 40        """Initialize a `CloudOrganization`."""
 41        self.organization_id = organization_id
 42        """The organization ID."""
 43
 44        self._organization_name = organization_name
 45        """Display name of the organization."""
 46
 47        self._email = email
 48        """Email associated with the organization."""
 49
 50        self._credentials = _AirbyteCredentials(
 51            client_id=SecretString(client_id) if client_id else None,
 52            client_secret=SecretString(client_secret) if client_secret else None,
 53            bearer_token=SecretString(bearer_token) if bearer_token else None,
 54            public_api_root=public_api_root or api_util.CLOUD_API_ROOT,
 55            config_api_root=config_api_root,
 56            organization_id=organization_id,
 57        )
 58        self._organization_info: dict[str, Any] | None = None
 59        self._organization_info_fetch_failed: bool = False
 60
 61    def _fetch_organization_info(self, *, force_refresh: bool = False) -> dict[str, Any]:
 62        """Fetch and cache organization info including billing status."""
 63        if force_refresh:
 64            self._organization_info_fetch_failed = False
 65
 66        if self._organization_info_fetch_failed and self._organization_info is None:
 67            return {}
 68
 69        if not force_refresh and self._organization_info is not None:
 70            return self._organization_info
 71
 72        try:
 73            self._organization_info = api_util.get_organization_info(
 74                organization_id=self.organization_id,
 75                api_root=self._credentials.public_api_root,
 76                config_api_root=self._credentials.config_api_root,
 77                client_id=self._credentials.client_id,
 78                client_secret=self._credentials.client_secret,
 79                bearer_token=self._credentials.bearer_token,
 80            )
 81        except Exception as ex:
 82            logger.debug("Failed to fetch organization info.", exc_info=ex)
 83            if self._organization_info is None:
 84                self._organization_info_fetch_failed = True
 85            return self._organization_info or {}
 86        else:
 87            return self._organization_info
 88
 89    @property
 90    def organization_name(self) -> str | None:
 91        """Display name of the organization."""
 92        if self._organization_name is not None:
 93            return self._organization_name
 94        info = self._fetch_organization_info()
 95        return info.get("organizationName")
 96
 97    @property
 98    def email(self) -> str | None:
 99        """Email associated with the organization."""
100        if self._email is not None:
101            return self._email
102        info = self._fetch_organization_info()
103        return info.get("email")
104
105    def get_billing_status(self) -> CloudOrganizationBillingInfo:
106        """Fetch billing status for the organization or raise on failure."""
107        try:
108            info = api_util.get_organization_info(
109                organization_id=self.organization_id,
110                api_root=self._credentials.public_api_root,
111                config_api_root=self._credentials.config_api_root,
112                client_id=self._credentials.client_id,
113                client_secret=self._credentials.client_secret,
114                bearer_token=self._credentials.bearer_token,
115            )
116        except (requests.RequestException, ValueError) as ex:
117            raise AirbyteError(
118                message="Failed to retrieve organization billing information.",
119                context={"organization_id": self.organization_id},
120            ) from ex
121        billing = info.get("billing")
122        if not isinstance(billing, dict):
123            raise AirbyteError(
124                message="Organization info did not include billing details.",
125                context={"organization_id": self.organization_id},
126            )
127        payment_status = billing.get("paymentStatus")
128        subscription_status = billing.get("subscriptionStatus")
129        return CloudOrganizationBillingInfo(
130            payment_status=payment_status if isinstance(payment_status, str) else None,
131            subscription_status=(
132                subscription_status if isinstance(subscription_status, str) else None
133            ),
134            is_account_locked=api_util.is_account_locked(payment_status, subscription_status),
135        )
136
137    @property
138    def payment_status(self) -> str | None:
139        """Payment status of the organization."""
140        info = self._fetch_organization_info()
141        return (info.get("billing") or {}).get("paymentStatus")
142
143    @property
144    def subscription_status(self) -> str | None:
145        """Subscription status of the organization."""
146        info = self._fetch_organization_info()
147        return (info.get("billing") or {}).get("subscriptionStatus")
148
149    @property
150    def is_account_locked(self) -> bool:
151        """Whether the account is locked due to billing issues."""
152        return api_util.is_account_locked(self.payment_status, self.subscription_status)