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
  9from airbyte._util import api_util
 10from airbyte.cloud._credentials import _AirbyteCredentials
 11from airbyte.secrets.base import SecretString
 12
 13
 14logger = logging.getLogger(__name__)
 15
 16
 17class CloudOrganization:
 18    """Information about an organization in Airbyte Cloud.
 19
 20    This class provides lazy loading of organization attributes including billing status.
 21    It is typically created via `CloudWorkspace.get_organization()`.
 22    """
 23
 24    def __init__(
 25        self,
 26        organization_id: str,
 27        organization_name: str | None = None,
 28        email: str | None = None,
 29        *,
 30        client_id: str | SecretString | None = None,
 31        client_secret: str | SecretString | None = None,
 32        bearer_token: str | SecretString | None = None,
 33        public_api_root: str | None = None,
 34        config_api_root: str | None = None,
 35    ) -> None:
 36        """Initialize a `CloudOrganization`."""
 37        self.organization_id = organization_id
 38        """The organization ID."""
 39
 40        self._organization_name = organization_name
 41        """Display name of the organization."""
 42
 43        self._email = email
 44        """Email associated with the organization."""
 45
 46        self._credentials = _AirbyteCredentials(
 47            client_id=SecretString(client_id) if client_id else None,
 48            client_secret=SecretString(client_secret) if client_secret else None,
 49            bearer_token=SecretString(bearer_token) if bearer_token else None,
 50            public_api_root=public_api_root or api_util.CLOUD_API_ROOT,
 51            config_api_root=config_api_root,
 52            organization_id=organization_id,
 53        )
 54        self._organization_info: dict[str, Any] | None = None
 55        self._organization_info_fetch_failed: bool = False
 56
 57    def _fetch_organization_info(self, *, force_refresh: bool = False) -> dict[str, Any]:
 58        """Fetch and cache organization info including billing status."""
 59        if force_refresh:
 60            self._organization_info_fetch_failed = False
 61
 62        if self._organization_info_fetch_failed and self._organization_info is None:
 63            return {}
 64
 65        if not force_refresh and self._organization_info is not None:
 66            return self._organization_info
 67
 68        try:
 69            self._organization_info = api_util.get_organization_info(
 70                organization_id=self.organization_id,
 71                api_root=self._credentials.public_api_root,
 72                config_api_root=self._credentials.config_api_root,
 73                client_id=self._credentials.client_id,
 74                client_secret=self._credentials.client_secret,
 75                bearer_token=self._credentials.bearer_token,
 76            )
 77        except Exception as ex:
 78            logger.debug("Failed to fetch organization info.", exc_info=ex)
 79            if self._organization_info is None:
 80                self._organization_info_fetch_failed = True
 81            return self._organization_info or {}
 82        else:
 83            return self._organization_info
 84
 85    @property
 86    def organization_name(self) -> str | None:
 87        """Display name of the organization."""
 88        if self._organization_name is not None:
 89            return self._organization_name
 90        info = self._fetch_organization_info()
 91        return info.get("organizationName")
 92
 93    @property
 94    def email(self) -> str | None:
 95        """Email associated with the organization."""
 96        if self._email is not None:
 97            return self._email
 98        info = self._fetch_organization_info()
 99        return info.get("email")
100
101    @property
102    def payment_status(self) -> str | None:
103        """Payment status of the organization."""
104        info = self._fetch_organization_info()
105        return (info.get("billing") or {}).get("paymentStatus")
106
107    @property
108    def subscription_status(self) -> str | None:
109        """Subscription status of the organization."""
110        info = self._fetch_organization_info()
111        return (info.get("billing") or {}).get("subscriptionStatus")
112
113    @property
114    def is_account_locked(self) -> bool:
115        """Whether the account is locked due to billing issues."""
116        return api_util.is_account_locked(self.payment_status, self.subscription_status)