airbyte.sources.registry

  1# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  2"""Connectivity to the connector catalog registry."""
  3
  4from __future__ import annotations
  5
  6import json
  7import logging
  8import os
  9import warnings
 10from copy import copy
 11from enum import Enum
 12from pathlib import Path
 13
 14import requests
 15from pydantic import BaseModel
 16
 17from airbyte import exceptions as exc
 18from airbyte._util.meta import is_docker_installed
 19from airbyte.constants import AIRBYTE_OFFLINE_MODE
 20from airbyte.logs import warn_once
 21from airbyte.version import get_version
 22
 23
 24logger = logging.getLogger("airbyte")
 25
 26
 27__cache: dict[str, ConnectorMetadata] | None = None
 28
 29
 30_REGISTRY_ENV_VAR = "AIRBYTE_LOCAL_REGISTRY"
 31_REGISTRY_URL = "https://connectors.airbyte.com/files/registries/v0/oss_registry.json"
 32
 33_PYTHON_LANGUAGE = "python"
 34_MANIFEST_ONLY_LANGUAGE = "manifest-only"
 35
 36_PYTHON_LANGUAGE_TAG = f"language:{_PYTHON_LANGUAGE}"
 37_MANIFEST_ONLY_TAG = f"language:{_MANIFEST_ONLY_LANGUAGE}"
 38
 39
 40class InstallType(str, Enum):
 41    """The type of installation for a connector."""
 42
 43    YAML = "yaml"
 44    PYTHON = "python"
 45    DOCKER = "docker"
 46    JAVA = "java"
 47
 48
 49class Language(str, Enum):
 50    """The language of a connector."""
 51
 52    PYTHON = InstallType.PYTHON.value
 53    JAVA = InstallType.JAVA.value
 54    MANIFEST_ONLY = _MANIFEST_ONLY_LANGUAGE
 55
 56
 57class ConnectorMetadata(BaseModel):
 58    """Metadata for a connector."""
 59
 60    name: str
 61    """Connector name. For example, "source-google-sheets"."""
 62
 63    latest_available_version: str | None
 64    """The latest available version of the connector."""
 65
 66    pypi_package_name: str | None
 67    """The name of the PyPI package for the connector, if it exists."""
 68
 69    language: Language | None
 70    """The language of the connector."""
 71
 72    install_types: set[InstallType]
 73    """The supported install types for the connector."""
 74
 75    suggested_streams: list[str] | None = None
 76    """A list of suggested streams for the connector, if available."""
 77
 78    @property
 79    def default_install_type(self) -> InstallType:
 80        """Return the default install type for the connector."""
 81        if self.language == Language.MANIFEST_ONLY and InstallType.YAML in self.install_types:
 82            return InstallType.YAML
 83
 84        if InstallType.PYTHON in self.install_types:
 85            return InstallType.PYTHON
 86
 87        # Else: Java or Docker
 88        return InstallType.DOCKER
 89
 90
 91def _get_registry_url() -> str:
 92    if _REGISTRY_ENV_VAR in os.environ:
 93        return str(os.environ.get(_REGISTRY_ENV_VAR))
 94
 95    return _REGISTRY_URL
 96
 97
 98def _is_registry_disabled(url: str) -> bool:
 99    return url.upper() in {"0", "F", "FALSE"} or AIRBYTE_OFFLINE_MODE
100
101
102def _registry_entry_to_connector_metadata(entry: dict) -> ConnectorMetadata:
103    name = entry["dockerRepository"].replace("airbyte/", "")
104    latest_version: str | None = entry.get("dockerImageTag")
105    tags = entry.get("tags", [])
106    language: Language | None = None
107
108    if "language" in entry and entry["language"] is not None:
109        try:
110            language = Language(entry["language"])
111        except Exception:
112            warnings.warn(
113                message=f"Invalid language for connector {name}: {entry['language']}",
114                stacklevel=2,
115            )
116    if not language and _PYTHON_LANGUAGE_TAG in tags:
117        language = Language.PYTHON
118    if not language and _MANIFEST_ONLY_TAG in tags:
119        language = Language.MANIFEST_ONLY
120
121    remote_registries: dict = entry.get("remoteRegistries", {})
122    pypi_registry: dict = remote_registries.get("pypi", {})
123    pypi_package_name: str = pypi_registry.get("packageName", None)
124    pypi_enabled: bool = pypi_registry.get("enabled", False)
125    install_types: set[InstallType] = {
126        x
127        for x in [
128            InstallType.DOCKER,  # Always True
129            InstallType.PYTHON if language == Language.PYTHON and pypi_enabled else None,
130            InstallType.JAVA if language == Language.JAVA else None,
131            InstallType.YAML if language == Language.MANIFEST_ONLY else None,
132        ]
133        if x
134    }
135
136    return ConnectorMetadata(
137        name=name,
138        latest_available_version=latest_version,
139        pypi_package_name=pypi_package_name if pypi_enabled else None,
140        language=language,
141        install_types=install_types,
142        suggested_streams=entry.get("suggestedStreams", {}).get("streams", None),
143    )
144
145
146def _get_registry_cache(*, force_refresh: bool = False) -> dict[str, ConnectorMetadata]:
147    """Return the registry cache."""
148    global __cache
149    if __cache and not force_refresh:
150        return __cache
151
152    registry_url = _get_registry_url()
153
154    if _is_registry_disabled(registry_url):
155        return {}
156
157    if registry_url.startswith("http"):
158        response = requests.get(
159            registry_url,
160            headers={"User-Agent": f"PyAirbyte/{get_version()}"},
161        )
162        response.raise_for_status()
163        data = response.json()
164    else:
165        # Assume local file
166        with Path(registry_url).open(encoding="utf-8") as f:
167            data = json.load(f)
168
169    new_cache: dict[str, ConnectorMetadata] = {}
170
171    for connector in data["sources"]:
172        connector_metadata = _registry_entry_to_connector_metadata(connector)
173        new_cache[connector_metadata.name] = connector_metadata
174
175    for connector in data["destinations"]:
176        connector_metadata = _registry_entry_to_connector_metadata(connector)
177        new_cache[connector_metadata.name] = connector_metadata
178
179    if len(new_cache) == 0:
180        # This isn't necessarily fatal, since users can bring their own
181        # connector definitions.
182        warn_once(
183            message=f"Connector registry is empty: {registry_url}",
184            with_stack=False,
185        )
186
187    __cache = new_cache
188    return __cache
189
190
191def get_connector_metadata(name: str) -> ConnectorMetadata | None:
192    """Check the cache for the connector.
193
194    If the cache is empty, populate by calling update_cache.
195    """
196    registry_url = _get_registry_url()
197
198    if _is_registry_disabled(registry_url):
199        return None
200
201    cache = copy(_get_registry_cache())
202
203    if not cache:
204        raise exc.PyAirbyteInternalError(
205            message="Connector registry could not be loaded.",
206            context={
207                "registry_url": _get_registry_url(),
208            },
209        )
210    if name not in cache:
211        raise exc.AirbyteConnectorNotRegisteredError(
212            connector_name=name,
213            context={
214                "registry_url": _get_registry_url(),
215                "available_connectors": get_available_connectors(),
216            },
217        )
218    return cache[name]
219
220
221def get_available_connectors(install_type: InstallType | str | None = None) -> list[str]:
222    """Return a list of all available connectors.
223
224    Connectors will be returned in alphabetical order, with the standard prefix "source-".
225    """
226    if install_type is None:
227        # No install type specified. Filter for whatever is runnable.
228        if is_docker_installed():
229            logger.info("Docker is detected. Returning all connectors.")
230            # If Docker is available, return all connectors.
231            return sorted(conn.name for conn in _get_registry_cache().values())
232
233        logger.info("Docker was not detected. Returning only Python and Manifest-only connectors.")
234
235        # If Docker is not available, return only Python and Manifest-based connectors.
236        return sorted(
237            conn.name
238            for conn in _get_registry_cache().values()
239            if conn.language in {Language.PYTHON, Language.MANIFEST_ONLY}
240        )
241
242    if not isinstance(install_type, InstallType):
243        install_type = InstallType(install_type)
244
245    if install_type == InstallType.PYTHON:
246        return sorted(
247            conn.name
248            for conn in _get_registry_cache().values()
249            if conn.pypi_package_name is not None
250        )
251
252    if install_type == InstallType.JAVA:
253        warnings.warn(
254            message="Java connectors are not yet supported.",
255            stacklevel=2,
256        )
257        return sorted(
258            conn.name for conn in _get_registry_cache().values() if conn.language == Language.JAVA
259        )
260
261    if install_type == InstallType.DOCKER:
262        return sorted(conn.name for conn in _get_registry_cache().values())
263
264    if install_type == InstallType.YAML:
265        return sorted(
266            conn.name
267            for conn in _get_registry_cache().values()
268            if InstallType.YAML in conn.install_types
269        )
270
271    # pragma: no cover  # Should never be reached.
272    raise exc.PyAirbyteInputError(
273        message="Invalid install type.",
274        context={
275            "install_type": install_type,
276        },
277    )