openmetadata-ingestion 1.9.0.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -12,12 +12,13 @@
 Base class for ingesting mlmodel services
 """
 from abc import ABC, abstractmethod
-from typing import Any, Iterable, List, Optional, Set
+from typing import Any, Iterable, List, Optional, Set, Tuple
 
 from pydantic import Field
 from typing_extensions import Annotated
 
 from metadata.generated.schema.api.data.createMlModel import CreateMlModelRequest
+from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.mlmodel import (
     MlFeature,
     MlHyperParameter,
@@ -39,6 +40,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
 from metadata.ingestion.models.delete_entity import DeleteEntity
+from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
@@ -92,6 +94,12 @@ class MlModelServiceTopology(ServiceTopology):
                 consumer=["mlmodel_service"],
                 use_cache=True,
             ),
+            NodeStage(
+                type_=AddLineageRequest,
+                processor="yield_mlmodel_lineage",
+                consumer=["mlmodel_service"],
+                nullable=True,
+            ),
         ],
     )
 
@@ -208,3 +216,69 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
 
     def prepare(self):
         """By default, nothing to prepare"""
+
+    def get_db_service_prefixes(self) -> List[str]:
+        """
+        Get the list of db service prefixes
+        """
+        return (
+            self.source_config.lineageInformation.dbServicePrefixes or []
+            if hasattr(self.source_config, "lineageInformation")
+            and self.source_config.lineageInformation
+            else []
+        )
+
+    def parse_db_service_prefix(
+        self, db_service_prefix: Optional[str]
+    ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
+        """
+        Parse the db service prefix
+        Returns:
+            Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]: service, database, schema, table
+        """
+        prefix_parts = (db_service_prefix or "").split(".")
+        return tuple(prefix_parts + ([None] * (4 - len(prefix_parts))))
+
+    def yield_mlmodel_lineage_details(
+        self, mlmodel_details: Any, db_service_prefix: Optional[str] = None
+    ) -> Iterable[Either[AddLineageRequest]]:
+        """
+        Get lineage between MLModel and source tables.
+        To be implemented by sources that support lineage.
+        """
+
+    def yield_mlmodel_lineage(
+        self, mlmodel_details: Any
+    ) -> Iterable[Either[OMetaLineageRequest]]:
+        """
+        Yields lineage if config is enabled.
+        We will look for the data in all the services we have informed.
+        """
+        db_service_prefixes = self.get_db_service_prefixes()
+        for db_service_prefix in db_service_prefixes or [None]:
+            for lineage in (
+                self.yield_mlmodel_lineage_details(mlmodel_details, db_service_prefix)
+                or []
+            ):
+                yield from self.yield_lineage_request(lineage)
+
+    def yield_lineage_request(
+        self, lineage: Optional[Either[AddLineageRequest]] = None
+    ) -> Iterable[Either[OMetaLineageRequest]]:
+        """
+        Method to yield lineage request
+        """
+        if lineage:
+            if lineage.right is not None:
+                yield Either(
+                    right=OMetaLineageRequest(
+                        lineage_request=lineage.right,
+                        override_lineage=(
+                            self.source_config.overrideLineage
+                            if hasattr(self.source_config, "overrideLineage")
+                            else False
+                        ),
+                    )
+                )
+            else:
+                yield lineage

metadata/ingestion/source/pipeline/airbyte/client.py

@@ -12,8 +12,17 @@
 Client to interact with airbyte apis
 """
 import json
-from typing import List
+import time
+from typing import List, Optional, Tuple
 
+import requests
+
+from metadata.generated.schema.entity.services.connections.pipeline.airbyte.basicAuth import (
+    BasicAuthentication,
+)
+from metadata.generated.schema.entity.services.connections.pipeline.airbyte.oauthClientAuth import (
+    Oauth20ClientCredentialsAuthentication,
+)
 from metadata.generated.schema.entity.services.connections.pipeline.airbyteConnection import (
     AirbyteConnection,
 )
@@ -21,26 +30,33 @@ from metadata.ingestion.ometa.client import REST, APIError, ClientConfig
 from metadata.utils.constants import AUTHORIZATION_HEADER, NO_ACCESS_TOKEN
 from metadata.utils.credentials import generate_http_basic_token
 from metadata.utils.helpers import clean_uri
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class AirbyteClient:
     """
-    Client handling API communication with Airbyte
+    Client for self-hosted Airbyte instances.
+    Supports Basic Authentication or no authentication.
     """
 
     def __init__(self, config: AirbyteConnection):
         self.config = config
+
         client_config: ClientConfig = ClientConfig(
             base_url=clean_uri(self.config.hostPort),
             api_version=self.config.apiVersion,
             auth_header=AUTHORIZATION_HEADER,
             auth_token=lambda: (NO_ACCESS_TOKEN, 0),
         )
-        if self.config.username:
+
+        if self.config.auth and isinstance(self.config.auth, BasicAuthentication):
             client_config.auth_token_mode = "Basic"
             client_config.auth_token = lambda: (
                 generate_http_basic_token(
-                    self.config.username, self.config.password.get_secret_value()
+                    self.config.auth.username,
+                    self.config.auth.password.get_secret_value(),
                 ),
                 0,
             )
@@ -96,3 +112,125 @@ class AirbyteClient:
         if response.get("exceptionStack"):
             raise APIError(response["message"])
         return response
+
+
+class AirbyteCloudClient(AirbyteClient):
+    """
+    Client for Airbyte Cloud instances.
+    Uses OAuth 2.0 Client Credentials authentication.
+    Overrides methods with Cloud-specific API endpoints.
+    """
+
+    def __init__(self, config: AirbyteConnection):
+        self.config = config
+        self._oauth_token: Optional[str] = None
+        self._oauth_token_expiry: float = 0
+
+        if not isinstance(config.auth, Oauth20ClientCredentialsAuthentication):
+            raise ValueError(
+                "AirbyteCloudClient requires OAuth 2.0 Client Credentials authentication"
+            )
+
+        client_config: ClientConfig = ClientConfig(
+            base_url=clean_uri(self.config.hostPort),
+            api_version=self.config.apiVersion,
+            auth_header=AUTHORIZATION_HEADER,
+            auth_token=self._get_oauth_token,
+        )
+        client_config.auth_token_mode = "Bearer"
+
+        self.client = REST(client_config)
+
+    def _fetch_oauth_token(self) -> Tuple[str, int]:
+        """
+        Fetch OAuth 2.0 access token using client credentials
+        """
+        token_url = f"{clean_uri(self.config.hostPort)}/v1/applications/token"
+        payload = {
+            "client_id": self.config.auth.clientId,
+            "client_secret": self.config.auth.clientSecret.get_secret_value(),
+        }
+
+        try:
+            response = requests.post(
+                token_url,
+                json=payload,
+                headers={"Content-Type": "application/json"},
+                timeout=30,
+            )
+            response.raise_for_status()
+            token_data = response.json()
+            access_token = token_data.get("access_token")
+
+            if not access_token:
+                raise APIError("No access_token in OAuth response")
+
+            logger.info("Successfully obtained OAuth access token for Airbyte Cloud")
+            return access_token, token_data.get("expires_in", 180)
+
+        except requests.exceptions.RequestException as exc:
+            logger.error(f"Failed to fetch OAuth token: {exc}")
+            raise APIError(f"OAuth token fetch failed: {exc}") from exc
+
+    def _get_oauth_token(self):
+        """
+        Get OAuth token with automatic refresh.
+        Returns tuple of (token, expiry_timestamp)
+        """
+        current_time = time.time()
+
+        if not self._oauth_token or current_time >= (self._oauth_token_expiry - 30):
+            self._oauth_token, expires_in = self._fetch_oauth_token()
+            self._oauth_token_expiry = current_time + expires_in
+
+        return (self._oauth_token, 0)
+
+    def list_workspaces(self) -> List[dict]:
+        """
+        Method returns the list of workspaces for Airbyte Cloud.
+        Uses the Cloud API endpoint which returns data in a different format.
+        """
+        response = self.client.get("/workspaces")
+        if response.get("exceptionStack"):
+            raise APIError(response["message"])
+        return response.get("data", [])
+
+    def list_connections(self, workflow_id: str) -> List[dict]:
+        """
+        Method returns the list of connections for a workspace in Airbyte Cloud.
+        Uses GET endpoint with workspaceIds query parameter.
+        """
+        response = self.client.get(f"/connections?workspaceIds={workflow_id}")
+        if response.get("exceptionStack"):
+            raise APIError(response["message"])
+        return response.get("data", [])
+
+    def list_jobs(self, connection_id: str) -> List[dict]:
+        """
+        Method returns the list of jobs for a connection in Airbyte Cloud.
+        Uses GET endpoint with connectionId query parameter.
+        """
+        response = self.client.get(f"/jobs?connectionId={connection_id}")
+        if response.get("exceptionStack"):
+            raise APIError(response["message"])
+        return response.get("data", [])
+
+    def get_source(self, source_id: str) -> dict:
+        """
+        Method returns source details for Airbyte Cloud.
+        Uses GET endpoint with path parameter instead of POST.
+        """
+        response = self.client.get(f"/sources/{source_id}")
+        if response.get("exceptionStack"):
+            raise APIError(response["message"])
+        return response
+
+    def get_destination(self, destination_id: str) -> dict:
+        """
+        Method returns destination details for Airbyte Cloud.
+        Uses GET endpoint with path parameter instead of POST.
+        """
+        response = self.client.get(f"/destinations/{destination_id}")
+        if response.get("exceptionStack"):
+            raise APIError(response["message"])
+        return response

metadata/ingestion/source/pipeline/airbyte/connection.py

@@ -12,11 +12,14 @@
 """
 Source connection handler
 """
-from typing import Optional
+from typing import Optional, Union
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
+from metadata.generated.schema.entity.services.connections.pipeline.airbyte.oauthClientAuth import (
+    Oauth20ClientCredentialsAuthentication,
+)
 from metadata.generated.schema.entity.services.connections.pipeline.airbyteConnection import (
     AirbyteConnection,
 )
@@ -25,20 +28,30 @@ from metadata.generated.schema.entity.services.connections.testConnectionResult
 )
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.pipeline.airbyte.client import AirbyteClient
+from metadata.ingestion.source.pipeline.airbyte.client import (
+    AirbyteClient,
+    AirbyteCloudClient,
+)
 from metadata.utils.constants import THREE_MIN
 
 
-def get_connection(connection: AirbyteConnection) -> AirbyteClient:
+def get_connection(
+    connection: AirbyteConnection,
+) -> Union[AirbyteClient, AirbyteCloudClient]:
     """
-    Create connection
+    Create connection - returns appropriate client based on auth type.
+    OAuth authentication indicates Airbyte Cloud, otherwise self-hosted instance.
     """
+    if connection.auth and isinstance(
+        connection.auth, Oauth20ClientCredentialsAuthentication
+    ):
+        return AirbyteCloudClient(connection)
     return AirbyteClient(connection)
 
 
 def test_connection(
     metadata: OpenMetadata,
-    client: AirbyteClient,
+    client: Union[AirbyteClient, AirbyteCloudClient],
     service_connection: AirbyteConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,

metadata/ingestion/source/pipeline/airbyte/metadata.py

@@ -12,6 +12,7 @@
 Airbyte source to extract metadata
 """
 
+from datetime import datetime
 from typing import Iterable, Optional
 
 from pydantic import BaseModel
@@ -45,6 +46,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.pipeline.airbyte.client import AirbyteCloudClient
 from metadata.ingestion.source.pipeline.openlineage.models import TableDetails
 from metadata.ingestion.source.pipeline.openlineage.utils import FQNNotFoundException
 from metadata.ingestion.source.pipeline.pipeline_service import PipelineServiceSource
@@ -83,6 +85,16 @@ class AirbyteSource(PipelineServiceSource):
     Pipeline metadata from Airflow's metadata db
     """
 
+    def __init__(self, config, metadata):
+        super().__init__(config, metadata)
+
+        if isinstance(self.client, AirbyteCloudClient):
+            self.airbyte_cloud = True
+            self.source_url_prefix = "https://cloud.airbyte.com"
+        else:
+            self.airbyte_cloud = False
+            self.source_url_prefix = clean_uri(self.service_connection.hostPort)
+
     @classmethod
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
@@ -116,7 +128,7 @@ class AirbyteSource(PipelineServiceSource):
         :return: Create Pipeline request with tasks
         """
         connection_url = (
-            f"{clean_uri(self.service_connection.hostPort)}/workspaces"
+            f"{self.source_url_prefix}/workspaces"
             f"/{pipeline_details.workspace.get('workspaceId')}"
             f"/connections/{pipeline_details.connection.get('connectionId')}"
         )
@@ -138,10 +150,13 @@ class AirbyteSource(PipelineServiceSource):
         """
         Method to get task & pipeline status
         """
+        if self.airbyte_cloud:
+            yield from self._yield_pipeline_status_cloud(pipeline_details)
+            return
 
         # Airbyte does not offer specific attempt link, just at pipeline level
         log_link = (
-            f"{self.service_connection.hostPort}workspaces/{pipeline_details.workspace.get('workspaceId')}"
+            f"{self.source_url_prefix}/workspaces/{pipeline_details.workspace.get('workspaceId')}"
             f"/connections/{pipeline_details.connection.get('connectionId')}/status"
         )
 
@@ -192,6 +207,79 @@ class AirbyteSource(PipelineServiceSource):
                     )
                 )
 
+    def _yield_pipeline_status_cloud(
+        self, pipeline_details: AirbytePipelineDetails
+    ) -> Iterable[Either[OMetaPipelineStatus]]:
+        """
+        Method to get task & pipeline status for Airbyte Cloud.
+        Handles flat job structure with ISO 8601 timestamps.
+        """
+        log_link = (
+            f"{self.source_url_prefix}/workspaces/{pipeline_details.workspace.get('workspaceId')}"
+            f"/connections/{pipeline_details.connection.get('connectionId')}/timeline"
+        )
+
+        for job in self.client.list_jobs(
+            pipeline_details.connection.get("connectionId")
+        ):
+            if not job:
+                continue
+
+            created_at = None
+            ended_at = None
+
+            if job.get("startTime"):
+                try:
+                    start_dt = datetime.fromisoformat(
+                        job["startTime"].replace("Z", "+00:00")
+                    )
+                    created_at = convert_timestamp_to_milliseconds(start_dt.timestamp())
+                except (ValueError, AttributeError) as exc:
+                    logger.warning(f"Failed to parse startTime: {exc}")
+
+            if job.get("lastUpdatedAt"):
+                try:
+                    end_dt = datetime.fromisoformat(
+                        job["lastUpdatedAt"].replace("Z", "+00:00")
+                    )
+                    ended_at = convert_timestamp_to_milliseconds(end_dt.timestamp())
+                except (ValueError, AttributeError) as exc:
+                    logger.warning(f"Failed to parse lastUpdatedAt: {exc}")
+
+            task_status = [
+                TaskStatus(
+                    name=str(pipeline_details.connection.get("connectionId")),
+                    executionStatus=STATUS_MAP.get(
+                        job["status"].lower(), StatusType.Pending
+                    ).value,
+                    startTime=created_at,
+                    endTime=ended_at,
+                    logLink=log_link,
+                )
+            ]
+
+            pipeline_status = PipelineStatus(
+                executionStatus=STATUS_MAP.get(
+                    job["status"].lower(), StatusType.Pending
+                ).value,
+                taskStatus=task_status,
+                timestamp=Timestamp(created_at) if created_at else None,
+            )
+
+            pipeline_fqn = fqn.build(
+                metadata=self.metadata,
+                entity_type=Pipeline,
+                service_name=self.context.get().pipeline_service,
+                pipeline_name=self.context.get().pipeline,
+            )
+
+            yield Either(
+                right=OMetaPipelineStatus(
+                    pipeline_fqn=pipeline_fqn,
+                    pipeline_status=pipeline_status,
+                )
+            )
+
     def _get_table_fqn(self, table_details: TableDetails) -> Optional[str]:
         """
         Get the FQN of the table
@@ -221,12 +309,24 @@ class AirbyteSource(PipelineServiceSource):
         :return: Lineage from inlets and outlets
         """
         pipeline_name = pipeline_details.connection.get("name")
+
+        logger.debug(
+            f"Processing lineage for pipeline: {pipeline_name}, "
+            f"connection_id: {pipeline_details.connection.get('connectionId')}, "
+            f"workspace_id: {pipeline_details.workspace.get('workspaceId')}"
+        )
+        logger.debug(f"Pipeline connection details: {pipeline_details.connection}")
+
         source_connection = self.client.get_source(
             pipeline_details.connection.get("sourceId")
         )
         destination_connection = self.client.get_destination(
             pipeline_details.connection.get("destinationId")
         )
+
+        logger.debug(f"Source connection response: {source_connection}")
+        logger.debug(f"Destination connection response: {destination_connection}")
+
         source_name = source_connection.get("sourceName")
         destination_name = destination_connection.get("destinationName")
 

metadata/ingestion/source/pipeline/airflow/connection.py

@@ -27,7 +27,7 @@ from metadata.generated.schema.entity.services.connections.database.mysqlConnect
     MysqlConnection as MysqlConnectionConfig,
 )
 from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
-    PostgresConnection,
+    PostgresConnection as PostgresConnectionConfig,
 )
 from metadata.generated.schema.entity.services.connections.database.sqliteConnection import (
     SQLiteConnection,
@@ -70,16 +70,16 @@ def _(_: BackendConnection) -> Engine:
 def _(airflow_connection: MysqlConnectionConfig) -> Engine:
     from metadata.ingestion.source.database.mysql.connection import MySQLConnection
 
-    return MySQLConnection(airflow_connection).get_client()
+    return MySQLConnection(airflow_connection)._get_client()
 
 
 @_get_connection.register
-def _(airflow_connection: PostgresConnection) -> Engine:
+def _(airflow_connection: PostgresConnectionConfig) -> Engine:
     from metadata.ingestion.source.database.postgres.connection import (
-        get_connection as get_postgres_connection,
+        PostgresConnection,
     )
 
-    return get_postgres_connection(airflow_connection)
+    return PostgresConnection(airflow_connection)._get_client()
 
 
 @_get_connection.register
@@ -131,7 +131,9 @@ def test_connection(
 
     def test_pipeline_details_access(session):
         try:
-            result = session.query(SerializedDagModel).first()
+            # Query only the dag_id column to avoid version compatibility issues
+            # The data_compressed column doesn't exist in Airflow 2.2.5
+            result = session.query(SerializedDagModel.dag_id).first()
             return result
         except Exception as e:
             raise AirflowPipelineDetailsAccessError(

metadata/ingestion/source/pipeline/airflow/metadata.py

@@ -371,7 +371,11 @@ class AirflowSource(PipelineServiceSource):
         offset = 0  # Start
 
         while True:
-            paginated_query = session_query.limit(limit).offset(offset)
+            paginated_query = (
+                session_query.order_by(SerializedDagModel.dag_id.asc())
+                .limit(limit)
+                .offset(offset)
+            )
             results = paginated_query.all()
             if not results:
                 break

metadata/ingestion/source/pipeline/databrickspipeline/connection.py

@@ -55,6 +55,7 @@ def get_connection(connection: DatabricksPipelineConnection) -> DatabricksClient
         get_connection_url_fn=get_connection_url,
         get_connection_args_fn=get_connection_args_common,
     )
+
     return DatabricksClient(connection, engine)