openmetadata-ingestion 1.6.7.1__py3-none-any.whl → 1.7.0.0rc1__py3-none-any.whl

metadata/ingestion/source/messaging/messaging_service.py

@@ -48,6 +48,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_topic
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -125,6 +126,7 @@ class MessagingServiceSource(TopologyRunnerMixin, Source, ABC):
     context = TopologyContextManager(topology)
     topic_source_state: Set = set()
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,

metadata/ingestion/source/metadata/alationsink/metadata.py

@@ -56,6 +56,7 @@ from metadata.ingestion.source.metadata.alationsink.models import (
 )
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -71,6 +72,7 @@ class AlationsinkSource(Source):
     config: WorkflowSource
     alation_sink_client: AlationSinkClient
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,
@@ -117,9 +119,11 @@ class AlationsinkSource(Source):
                     ),
                 ),
                 db_username="Test",
-                title=om_database.displayName
-                if om_database.displayName
-                else model_str(om_database.name),
+                title=(
+                    om_database.displayName
+                    if om_database.displayName
+                    else model_str(om_database.name)
+                ),
                 description=model_str(om_database.description),
             )
         except Exception as exc:
@@ -140,9 +144,11 @@ class AlationsinkSource(Source):
                 key=fqn._build(  # pylint: disable=protected-access
                     str(alation_datasource_id), model_str(om_schema.name)
                 ),
-                title=om_schema.displayName
-                if om_schema.displayName
-                else model_str(om_schema.name),
+                title=(
+                    om_schema.displayName
+                    if om_schema.displayName
+                    else model_str(om_schema.name)
+                ),
                 description=model_str(om_schema.description),
             )
         except Exception as exc:
@@ -163,9 +169,11 @@ class AlationsinkSource(Source):
                 key=fqn._build(  # pylint: disable=protected-access
                     str(alation_datasource_id), schema_name, model_str(om_table.name)
                 ),
-                title=om_table.displayName
-                if om_table.displayName
-                else model_str(om_table.name),
+                title=(
+                    om_table.displayName
+                    if om_table.displayName
+                    else model_str(om_table.name)
+                ),
                 description=model_str(om_table.description),
                 table_type=TABLE_TYPE_MAPPER.get(om_table.tableType, "TABLE"),
                 sql=om_table.schemaDefinition,
@@ -273,16 +281,22 @@ class AlationsinkSource(Source):
                     table_name,
                     model_str(om_column.name),
                 ),
-                column_type=om_column.dataTypeDisplay.lower()
-                if om_column.dataTypeDisplay
-                else om_column.dataType.value.lower(),
-                title=om_column.displayName
-                if om_column.displayName
-                else model_str(om_column.name),
+                column_type=(
+                    om_column.dataTypeDisplay.lower()
+                    if om_column.dataTypeDisplay
+                    else om_column.dataType.value.lower()
+                ),
+                title=(
+                    om_column.displayName
+                    if om_column.displayName
+                    else model_str(om_column.name)
+                ),
                 description=model_str(om_column.description),
-                position=str(om_column.ordinalPosition)
-                if om_column.ordinalPosition
-                else None,
+                position=(
+                    str(om_column.ordinalPosition)
+                    if om_column.ordinalPosition
+                    else None
+                ),
                 index=self._get_column_index(
                     alation_datasource_id, om_column, table_constraints
                 ),

metadata/ingestion/source/metadata/amundsen/metadata.py

@@ -70,7 +70,7 @@ from metadata.ingestion.source.metadata.amundsen.queries import (
     NEO4J_AMUNDSEN_USER_QUERY,
 )
 from metadata.utils import fqn
-from metadata.utils.helpers import get_standard_chart_type
+from metadata.utils.helpers import get_standard_chart_type, retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.metadata_service_helper import SERVICE_TYPE_MAPPER
 from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_labels
@@ -116,6 +116,7 @@ class AmundsenSource(Source):
 
     dashboard_service: DashboardService
 
+    @retry_with_docker_host()
     def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
         super().__init__()
         self.config = config
@@ -248,9 +249,11 @@ class AmundsenSource(Source):
                 table_name = "default"
 
             database_request = CreateDatabaseRequest(
-                name=table_name
-                if hasattr(service_entity.connection.config, "supportsDatabase")
-                else "default",
+                name=(
+                    table_name
+                    if hasattr(service_entity.connection.config, "supportsDatabase")
+                    else "default"
+                ),
                 service=service_entity.fullyQualifiedName,
             )
             yield Either(right=database_request)

metadata/ingestion/source/metadata/atlas/metadata.py

@@ -50,7 +50,7 @@ from metadata.ingestion.source.connections import get_connection, get_test_conne
 from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
 from metadata.ingestion.source.metadata.atlas.client import AtlasClient
 from metadata.utils import fqn
-from metadata.utils.helpers import get_database_name_for_lineage
+from metadata.utils.helpers import get_database_name_for_lineage, retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.metadata_service_helper import SERVICE_TYPE_MAPPER
 from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_labels
@@ -70,6 +70,7 @@ class AtlasSource(Source):
     tables: Dict[str, Any]
     topics: Dict[str, Any]
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -51,6 +51,7 @@ from metadata.ingestion.models.topology import (
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
 from metadata.utils import fqn
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -113,6 +114,7 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
     context = TopologyContextManager(topology)
     mlmodel_source_state: Set = set()
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,

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

@@ -29,6 +29,7 @@ from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequ
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.pipeline import (
     Pipeline,
+    PipelineState,
     PipelineStatus,
     StatusType,
     Task,
@@ -122,6 +123,7 @@ class AirflowSource(PipelineServiceSource):
         metadata: OpenMetadata,
     ):
         super().__init__(config, metadata)
+        self.today = datetime.now().strftime("%Y-%m-%d")
         self._session = None
 
     @classmethod
@@ -367,6 +369,16 @@ class AirflowSource(PipelineServiceSource):
                 break
             for serialized_dag in results:
                 try:
+                    dag_model = (
+                        self.session.query(DagModel)
+                        .filter(DagModel.dag_id == serialized_dag[0])
+                        .one_or_none()
+                    )
+                    pipeline_state = (
+                        PipelineState.Active.value
+                        if dag_model and not dag_model.is_paused
+                        else PipelineState.Inactive.value
+                    )
                     data = serialized_dag[1]["dag"]
                     dag = AirflowDagDetails(
                         dag_id=serialized_dag[0],
@@ -375,6 +387,7 @@ class AirflowSource(PipelineServiceSource):
                         max_active_runs=data.get("max_active_runs", None),
                         description=data.get("_description", None),
                         start_date=data.get("start_date", None),
+                        state=pipeline_state,
                         tasks=list(
                             map(self._extract_serialized_task, data.get("tasks", []))
                         ),
@@ -425,6 +438,14 @@ class AirflowSource(PipelineServiceSource):
         """
         return pipeline_details.dag_id
 
+    def get_pipeline_state(
+        self, pipeline_details: AirflowDagDetails
+    ) -> Optional[PipelineState]:
+        """
+        Return the state of the DAG
+        """
+        return PipelineState[pipeline_details.state]
+
     def get_tasks_from_dag(self, dag: AirflowDagDetails, host_port: str) -> List[Task]:
         """
         Obtain the tasks from a SerializedDAG
@@ -480,12 +501,15 @@ class AirflowSource(PipelineServiceSource):
         try:
             # Airflow uses /dags/dag_id/grid to show pipeline / dag
             source_url = f"{clean_uri(self.service_connection.hostPort)}/dags/{pipeline_details.dag_id}/grid"
+            pipeline_state = self.get_pipeline_state(pipeline_details)
+
             pipeline_request = CreatePipelineRequest(
                 name=EntityName(pipeline_details.dag_id),
                 description=Markdown(pipeline_details.description)
                 if pipeline_details.description
                 else None,
                 sourceUrl=SourceUrl(source_url),
+                state=pipeline_state,
                 concurrency=pipeline_details.max_active_runs,
                 pipelineLocation=pipeline_details.fileloc,
                 startDate=pipeline_details.start_date.isoformat()
@@ -614,4 +638,5 @@ class AirflowSource(PipelineServiceSource):
                     )
 
     def close(self):
+        self.metadata.compute_percentile(Pipeline, self.today)
         self.session.close()

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

@@ -68,4 +68,5 @@ class AirflowDagDetails(AirflowBaseModel):
     start_date: Optional[datetime] = None
     tasks: List[AirflowTask]
     owner: Optional[str] = None
+    state: Optional[str] = None
     schedule_interval: Optional[str] = None

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

@@ -12,11 +12,19 @@
 Client to interact with Nifi apis
 """
 import traceback
-from typing import Any, Iterable, List, Optional
-
-import requests
-from requests import HTTPError
-
+from typing import Dict, Iterable, List
+
+from metadata.generated.schema.entity.services.connections.pipeline.nifi.basicAuth import (
+    NifiBasicAuth,
+)
+from metadata.generated.schema.entity.services.connections.pipeline.nifi.clientCertificateAuth import (
+    NifiClientCertificateAuth,
+)
+from metadata.generated.schema.entity.services.connections.pipeline.nifiConnection import (
+    NifiConnection,
+)
+from metadata.ingestion.ometa.client import REST, ClientConfig, HTTPError
+from metadata.utils.constants import AUTHORIZATION_HEADER, NO_ACCESS_TOKEN
 from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import ingestion_logger
 
@@ -26,6 +34,8 @@ IDENTIFIER = "identifier"
 PROCESS_GROUPS_STARTER = "/process-groups/"
 RESOURCES = "resources"
 REQUESTS_TIMEOUT = 60 * 5
+CONTENT_HEADER = {"Content-type": "application/x-www-form-urlencoded"}
+NIFI_API_BASE_ENDPOINT = "/nifi-api"
 
 
 class NifiClient:
@@ -33,39 +43,43 @@ class NifiClient:
     Wrapper on top of Nifi REST API
     """
 
-    # pylint: disable=too-many-arguments
-    def __init__(
-        self,
-        host_port: str,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
-        ca_file_path: Optional[str] = None,
-        client_cert_path: Optional[str] = None,
-        client_key_path: Optional[str] = None,
-        verify: bool = False,
-    ):
-        self._token = None
-        self._resources = None
-
-        self.content_headers = {"Content-Type": "application/x-www-form-urlencoded"}
-        self.api_endpoint = clean_uri(host_port) + "/nifi-api"
-        self.username = username
-        self.password = password
-
-        if all(setting for setting in [self.username, self.password]):
-            self.data = {"username": self.username, "password": self.password}
-            self.verify = verify
-            self.headers = {
-                "Authorization": f"Bearer {self.token}",
-                **self.content_headers,
+    client: REST
+
+    def __init__(self, connection: NifiConnection):
+        self.connection = connection
+        self._token, self._resources, self.data = None, None, None
+        self.api_endpoint = clean_uri(self.connection.hostPort) + NIFI_API_BASE_ENDPOINT
+
+        client_config = ClientConfig(
+            api_version="",
+            timeout=REQUESTS_TIMEOUT,
+            base_url=self.api_endpoint,
+            extra_headers=CONTENT_HEADER,
+        )
+
+        if isinstance(self.connection.nifiConfig, NifiBasicAuth):
+            self.verify = self.connection.nifiConfig.verifySSL
+            self.data = {
+                "username": self.connection.nifiConfig.username,
+                "password": self.connection.nifiConfig.password.get_secret_value()
+                if self.connection.nifiConfig.password
+                else None,
             }
-            self.client_cert = None
-        else:
-            self.data = None
-            self.verify = ca_file_path if ca_file_path else False
-            self.client_cert = (client_cert_path, client_key_path)
-            self.headers = self.content_headers
-            access = self.get("access")
+            client_config.verify = self.connection.nifiConfig.verifySSL
+            client_config.auth_header = AUTHORIZATION_HEADER
+            client_config.access_token = self.token
+
+            self.client = REST(client_config)
+        elif isinstance(self.connection.nifiConfig, NifiClientCertificateAuth):
+            ca_path = self.connection.nifiConfig.certificateAuthorityPath
+            cc_path = self.connection.nifiConfig.clientCertificatePath
+            ck_path = self.connection.nifiConfig.clientkeyPath
+
+            client_config.verify = ca_path if ca_path else False
+            client_config.cert = (cc_path, ck_path)
+
+            self.client = REST(client_config)
+            access = self.client.get("access")
             logger.debug(access)
 
     @property
@@ -76,13 +90,19 @@ class NifiClient:
         """
         if not self._token:
             try:
-                res = requests.post(
-                    f"{self.api_endpoint}/access/token",
-                    verify=self.verify,
-                    headers=self.content_headers,
-                    data=self.data,
-                    timeout=REQUESTS_TIMEOUT,
+                client = REST(
+                    ClientConfig(
+                        base_url=self.api_endpoint,
+                        verify=self.verify,
+                        extra_headers=CONTENT_HEADER,
+                        timeout=REQUESTS_TIMEOUT,
+                        api_version="",
+                        auth_token_mode=None,
+                        auth_token=lambda: (NO_ACCESS_TOKEN, 0),
+                    )
                 )
+
+                res = client.post("access/token", data=self.data)
                 self._token = res.text
 
                 if res.status_code not in (200, 201):
@@ -100,6 +120,7 @@ class NifiClient:
 
             except Exception as err:
                 logger.error(f"Fetching token failed due to - {err}")
+                logger.debug(traceback.format_exc())
                 raise err
 
         return self._token
@@ -110,7 +131,7 @@ class NifiClient:
         This can be expensive. Only query it once.
         """
         if not self._resources:
-            self._resources = self.get(RESOURCES)  # API endpoint
+            self._resources = self.client.get(RESOURCES)  # API endpoint
 
         # Get the first `resources` key from the dict
         try:
@@ -118,33 +139,6 @@ class NifiClient:
         except AttributeError:
             return []
 
-    def get(self, path: str) -> Optional[Any]:
-        """
-        GET call wrapper
-        """
-        try:
-            res = requests.get(
-                f"{self.api_endpoint}/{path}",
-                verify=self.verify,
-                headers=self.headers,
-                timeout=REQUESTS_TIMEOUT,
-                cert=self.client_cert,
-            )
-
-            return res.json()
-
-        except HTTPError as err:
-            logger.warning(f"Connection error calling the Nifi API - {err}")
-            raise err
-
-        except ValueError as err:
-            logger.warning(f"Cannot pick up the JSON from API response - {err}")
-            raise err
-
-        except Exception as err:
-            logger.warning(f"Unknown error calling Nifi API - {err}")
-            raise err
-
     def _get_process_group_ids(self) -> List[str]:
         return [
             elem.get(IDENTIFIER).replace(PROCESS_GROUPS_STARTER, "")
@@ -152,10 +146,10 @@ class NifiClient:
             if elem.get(IDENTIFIER).startswith(PROCESS_GROUPS_STARTER)
         ]
 
-    def get_process_group(self, id_: str) -> dict:
-        return self.get(f"flow/process-groups/{id_}")
+    def get_process_group(self, id_: str) -> Dict:
+        return self.client.get(f"flow/process-groups/{id_}")
 
-    def list_process_groups(self) -> Iterable[dict]:
+    def list_process_groups(self) -> Iterable[Dict]:
         """
         This will call the API endpoints
         one at a time.

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

@@ -17,9 +17,6 @@ from typing import Optional
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
-from metadata.generated.schema.entity.services.connections.pipeline.nifi.basicAuth import (
-    NifiBasicAuth,
-)
 from metadata.generated.schema.entity.services.connections.pipeline.nifiConnection import (
     NifiConnection,
 )
@@ -36,22 +33,8 @@ def get_connection(connection: NifiConnection) -> NifiClient:
     """
     Create connection
     """
-    if isinstance(connection.nifiConfig, NifiBasicAuth):
-        return NifiClient(
-            host_port=connection.hostPort,
-            username=connection.nifiConfig.username,
-            password=connection.nifiConfig.password.get_secret_value()
-            if connection.nifiConfig.password
-            else None,
-            verify=connection.nifiConfig.verifySSL,
-        )
 
-    return NifiClient(
-        host_port=connection.hostPort,
-        ca_file_path=connection.nifiConfig.certificateAuthorityPath,
-        client_cert_path=connection.nifiConfig.clientCertificatePath,
-        client_key_path=connection.nifiConfig.clientkeyPath,
-    )
+    return NifiClient(connection=connection)
 
 
 def test_connection(

metadata/ingestion/source/pipeline/openlineage/metadata.py

@@ -93,6 +93,7 @@ class OpenlineageSource(PipelineServiceSource):
         """Nothing to prepare"""
 
     def close(self) -> None:
+        self.metadata.compute_percentile(Pipeline, self.today)
         self.metadata.close()
 
     @classmethod