acryl-datahub 0.15.0.6rc3__py3-none-any.whl → 1.0.0__py3-none-any.whl

datahub/ingestion/source/iceberg/iceberg_common.py

@@ -1,11 +1,15 @@
 import logging
+import threading
 from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
 from humanfriendly import format_timespan
 from pydantic import Field, validator
 from pyiceberg.catalog import Catalog, load_catalog
+from pyiceberg.catalog.rest import RestCatalog
+from requests.adapters import HTTPAdapter
 from sortedcontainers import SortedList
+from urllib3.util import Retry
 
 from datahub.configuration.common import AllowDenyPattern, ConfigModel
 from datahub.configuration.source_common import DatasetSourceConfigMixin
@@ -25,6 +29,23 @@ from datahub.utilities.stats_collections import TopKDict, int_top_k_dict
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_REST_TIMEOUT = 120
+DEFAULT_REST_RETRY_POLICY = {"total": 3, "backoff_factor": 0.1}
+
+
+class TimeoutHTTPAdapter(HTTPAdapter):
+    def __init__(self, *args, **kwargs):
+        if "timeout" in kwargs:
+            self.timeout = kwargs["timeout"]
+            del kwargs["timeout"]
+        super().__init__(*args, **kwargs)
+
+    def send(self, request, **kwargs):
+        timeout = kwargs.get("timeout")
+        if timeout is None and hasattr(self, "timeout"):
+            kwargs["timeout"] = self.timeout
+        return super().send(request, **kwargs)
+
 
 class IcebergProfilingConfig(ConfigModel):
     enabled: bool = Field(
@@ -145,7 +166,26 @@ class IcebergSourceConfig(StatefulIngestionConfigBase, DatasetSourceConfigMixin)
         logger.debug(
             "Initializing the catalog %s with config: %s", catalog_name, catalog_config
         )
-        return load_catalog(name=catalog_name, **catalog_config)
+        catalog = load_catalog(name=catalog_name, **catalog_config)
+        if isinstance(catalog, RestCatalog):
+            logger.debug(
+                "Recognized REST catalog type being configured, attempting to configure HTTP Adapter for the session"
+            )
+            retry_policy: Dict[str, Any] = DEFAULT_REST_RETRY_POLICY.copy()
+            retry_policy.update(catalog_config.get("connection", {}).get("retry", {}))
+            retries = Retry(**retry_policy)
+            logger.debug(f"Retry policy to be set: {retry_policy}")
+            timeout = catalog_config.get("connection", {}).get(
+                "timeout", DEFAULT_REST_TIMEOUT
+            )
+            logger.debug(f"Timeout to be set: {timeout}")
+            catalog._session.mount(
+                "http://", TimeoutHTTPAdapter(timeout=timeout, max_retries=retries)
+            )
+            catalog._session.mount(
+                "https://", TimeoutHTTPAdapter(timeout=timeout, max_retries=retries)
+            )
+        return catalog
 
 
 class TopTableTimings:
@@ -156,18 +196,21 @@ class TopTableTimings:
     def __init__(self, size: int = 10):
         self._size = size
         self.top_entites = SortedList(key=lambda x: -x.get(self._VALUE_FIELD, 0))
+        self._lock = threading.Lock()
 
     def add(self, entity: Dict[str, Any]) -> None:
         if self._VALUE_FIELD not in entity:
             return
-        self.top_entites.add(entity)
-        if len(self.top_entites) > self._size:
-            self.top_entites.pop()
+        with self._lock:
+            self.top_entites.add(entity)
+            if len(self.top_entites) > self._size:
+                self.top_entites.pop()
 
     def __str__(self) -> str:
-        if len(self.top_entites) == 0:
-            return "no timings reported"
-        return str(list(self.top_entites))
+        with self._lock:
+            if len(self.top_entites) == 0:
+                return "no timings reported"
+            return str(list(self.top_entites))
 
 
 class TimingClass:
@@ -175,24 +218,31 @@ class TimingClass:
 
     def __init__(self):
         self.times = SortedList()
+        self._lock = threading.Lock()
 
     def add_timing(self, t: float) -> None:
-        self.times.add(t)
+        with self._lock:
+            self.times.add(t)
 
     def __str__(self) -> str:
-        if len(self.times) == 0:
-            return "no timings reported"
-        total = sum(self.times)
-        avg = total / len(self.times)
-        return str(
-            {
-                "average_time": format_timespan(avg, detailed=True, max_units=3),
-                "min_time": format_timespan(self.times[0], detailed=True, max_units=3),
-                "max_time": format_timespan(self.times[-1], detailed=True, max_units=3),
-                # total_time does not provide correct information in case we run in more than 1 thread
-                "total_time": format_timespan(total, detailed=True, max_units=3),
-            }
-        )
+        with self._lock:
+            if len(self.times) == 0:
+                return "no timings reported"
+            total = sum(self.times)
+            avg = total / len(self.times)
+            return str(
+                {
+                    "average_time": format_timespan(avg, detailed=True, max_units=3),
+                    "min_time": format_timespan(
+                        self.times[0], detailed=True, max_units=3
+                    ),
+                    "max_time": format_timespan(
+                        self.times[-1], detailed=True, max_units=3
+                    ),
+                    # total_time does not provide correct information in case we run in more than 1 thread
+                    "total_time": format_timespan(total, detailed=True, max_units=3),
+                }
+            )
 
 
 @dataclass

datahub/ingestion/source/identity/okta.py

@@ -5,7 +5,7 @@ import urllib
 from collections import defaultdict
 from dataclasses import dataclass, field
 from time import sleep
-from typing import Dict, Iterable, List, Optional, Union
+from typing import Dict, Iterable, List, Optional, Set, Union
 
 import nest_asyncio
 from okta.client import Client as OktaClient
@@ -14,7 +14,6 @@ from okta.models import Group, GroupProfile, User, UserProfile, UserStatus
 from pydantic import validator
 from pydantic.fields import Field
 
-from datahub.configuration.common import ConfigModel
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
@@ -56,7 +55,7 @@ logger = logging.getLogger(__name__)
 nest_asyncio.apply()
 
 
-class OktaConfig(StatefulIngestionConfigBase, ConfigModel):
+class OktaConfig(StatefulIngestionConfigBase):
     # Required: Domain of the Okta deployment. Example: dev-33231928.okta.com
     okta_domain: str = Field(
         description="The location of your Okta Domain, without a protocol. Can be found in Okta Developer console. e.g. dev-33231928.okta.com",
@@ -77,6 +76,10 @@ class OktaConfig(StatefulIngestionConfigBase, ConfigModel):
         default=True,
         description="Whether group membership should be ingested into DataHub. ingest_groups must be True if this is True.",
     )
+    ingest_groups_users: bool = Field(
+        default=True,
+        description="Only ingest users belonging to the selected groups. This option is only useful when `ingest_users` is set to False and `ingest_group_membership` to True.",
+    )
 
     # Optional: Customize the mapping to DataHub Username from an attribute appearing in the Okta User
     # profile. Reference: https://developer.okta.com/docs/reference/api/users/
@@ -344,6 +347,7 @@ class OktaSource(StatefulIngestionSourceBase):
                     aspect=StatusClass(removed=False),
                 ).as_workunit()
 
+        okta_users: Set[User] = set()
         # Step 2: Populate GroupMembership Aspects for CorpUsers
         datahub_corp_user_urn_to_group_membership: Dict[str, GroupMembershipClass] = (
             defaultdict(lambda: GroupMembershipClass(groups=[]))
@@ -372,6 +376,9 @@ class OktaSource(StatefulIngestionSourceBase):
                         self.report.report_failure("okta_user_mapping", error_str)
                         continue
 
+                    if self.config.ingest_groups_users:
+                        okta_users.add(okta_user)
+
                     # Update the GroupMembership aspect for this group member.
                     datahub_corp_user_urn_to_group_membership[
                         datahub_corp_user_urn
@@ -379,7 +386,10 @@ class OktaSource(StatefulIngestionSourceBase):
 
         # Step 3: Produce MetadataWorkUnits for CorpUsers.
         if self.config.ingest_users:
-            okta_users = self._get_okta_users(event_loop)
+            # we can just throw away collected okta users so far and fetch them all
+            okta_users = set(self._get_okta_users(event_loop))
+
+        if okta_users:
             filtered_okta_users = filter(self._filter_okta_user, okta_users)
             datahub_corp_user_snapshots = self._map_okta_users(filtered_okta_users)
             for user_count, datahub_corp_user_snapshot in enumerate(
@@ -558,9 +568,7 @@ class OktaSource(StatefulIngestionSourceBase):
         if (
             self.config.include_deprovisioned_users is False
             and okta_user.status == UserStatus.DEPROVISIONED
-        ):
-            return False
-        elif (
+        ) or (
             self.config.include_suspended_users is False
             and okta_user.status == UserStatus.SUSPENDED
         ):
@@ -658,6 +666,27 @@ class OktaSource(StatefulIngestionSourceBase):
             self.config.okta_profile_to_username_regex,
         )
 
+    def _map_okta_user_profile_custom_properties(
+        self, profile: UserProfile
+    ) -> Dict[str, str]:
+        # filter out the common fields that are already mapped to the CorpUserInfo aspect and the private ones
+        return {
+            k: str(v)
+            for k, v in profile.__dict__.items()
+            if v
+            and k
+            not in [
+                "displayName",
+                "firstName",
+                "lastName",
+                "email",
+                "title",
+                "countryCode",
+                "department",
+            ]
+            and not k.startswith("_")
+        }
+
     # Converts Okta User Profile into a CorpUserInfo.
     def _map_okta_user_profile(self, profile: UserProfile) -> CorpUserInfoClass:
         # TODO: Extract user's manager if provided.
@@ -675,6 +704,7 @@ class OktaSource(StatefulIngestionSourceBase):
             title=profile.title,
             countryCode=profile.countryCode,
             departmentName=profile.department,
+            customProperties=self._map_okta_user_profile_custom_properties(profile),
         )
 
     def _make_corp_group_urn(self, name: str) -> str:

datahub/ingestion/source/kafka/kafka.py

@@ -272,7 +272,7 @@ class KafkaSource(StatefulIngestionSourceBase, TestableSource):
             return schema_registry_class.create(config, report)
         except Exception as e:
             logger.debug(e, exc_info=e)
-            raise ImportError(config.schema_registry_class)
+            raise ImportError(config.schema_registry_class) from e
 
     def __init__(self, config: KafkaSourceConfig, ctx: PipelineContext):
         super().__init__(config, ctx)

datahub/ingestion/source/kafka_connect/common.py

@@ -110,7 +110,7 @@ class ConnectorManifest:
 
     name: str
     type: str
-    config: Dict
+    config: Dict[str, str]
     tasks: Dict
     url: Optional[str] = None
     flow_property_bag: Optional[Dict[str, str]] = None
@@ -141,12 +141,7 @@ def get_dataset_name(
     database_name: Optional[str],
     source_table: str,
 ) -> str:
-    if database_name:
-        dataset_name = database_name + "." + source_table
-    else:
-        dataset_name = source_table
-
-    return dataset_name
+    return database_name + "." + source_table if database_name else source_table
 
 
 def get_platform_instance(

datahub/ingestion/source/kafka_connect/kafka_connect.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Iterable, List, Optional, Type
+from typing import Dict, Iterable, List, Optional, Type
 
 import jpype
 import jpype.imports
@@ -121,7 +121,11 @@ class KafkaConnectSource(StatefulIngestionSourceBase):
                     connector_manifest.config, self.config.provided_configs
                 )
             connector_manifest.url = connector_url
-            connector_manifest.topic_names = self._get_connector_topics(connector_name)
+            connector_manifest.topic_names = self._get_connector_topics(
+                connector_name=connector_name,
+                config=connector_manifest.config,
+                connector_type=connector_manifest.type,
+            )
             connector_class_value = connector_manifest.config.get(CONNECTOR_CLASS) or ""
 
             class_type: Type[BaseConnector] = BaseConnector
@@ -203,7 +207,9 @@ class KafkaConnectSource(StatefulIngestionSourceBase):
 
         return response.json()
 
-    def _get_connector_topics(self, connector_name: str) -> List[str]:
+    def _get_connector_topics(
+        self, connector_name: str, config: Dict[str, str], connector_type: str
+    ) -> List[str]:
         try:
             response = self.session.get(
                 f"{self.config.connect_uri}/connectors/{connector_name}/topics",
@@ -215,7 +221,21 @@ class KafkaConnectSource(StatefulIngestionSourceBase):
             )
             return []
 
-        return response.json()[connector_name]["topics"]
+        processed_topics = response.json()[connector_name]["topics"]
+
+        if connector_type == SINK:
+            try:
+                return SinkTopicFilter().filter_stale_topics(processed_topics, config)
+            except Exception as e:
+                self.report.warning(
+                    title="Error parsing sink conector topics configuration",
+                    message="Some stale lineage tasks might show up for connector",
+                    context=connector_name,
+                    exc=e,
+                )
+                return processed_topics
+        else:
+            return processed_topics
 
     def construct_flow_workunit(self, connector: ConnectorManifest) -> MetadataWorkUnit:
         connector_name = connector.name
@@ -359,3 +379,76 @@ class KafkaConnectSource(StatefulIngestionSourceBase):
         return builder.make_dataset_urn_with_platform_instance(
             platform, name, platform_instance, self.config.env
         )
+
+
+class SinkTopicFilter:
+    """Helper class to filter Kafka Connect topics based on configuration."""
+
+    def filter_stale_topics(
+        self,
+        processed_topics: List[str],
+        sink_config: Dict[str, str],
+    ) -> List[str]:
+        """
+        Kafka-connect's /topics API returns the set of topic names the connector has been using
+        since its creation or since the last time its set of active topics was reset. This means-
+        if a topic was ever used by a connector, it will be returned, even if it is no longer used.
+        To remove these stale topics from the list, we double-check the list returned by the API
+        against the sink connector's config.
+        Sink connectors configure exactly one of `topics` or `topics.regex`
+        https://kafka.apache.org/documentation/#sinkconnectorconfigs_topics
+
+        Args:
+            processed_topics: List of topics currently being processed
+            sink_config: Configuration dictionary for the sink connector
+
+        Returns:
+            List of filtered topics that match the configuration
+
+        Raises:
+            ValueError: If sink connector configuration is missing both 'topics' and 'topics.regex' fields
+
+        """
+        # Absence of topics config is a defensive NOOP,
+        # although this should never happen in real world
+        if not self.has_topic_config(sink_config):
+            logger.warning(
+                f"Found sink without topics config {sink_config.get(CONNECTOR_CLASS)}"
+            )
+            return processed_topics
+
+        # Handle explicit topic list
+        if sink_config.get("topics"):
+            return self._filter_by_topic_list(processed_topics, sink_config["topics"])
+        else:
+            # Handle regex pattern
+            return self._filter_by_topic_regex(
+                processed_topics, sink_config["topics.regex"]
+            )
+
+    def has_topic_config(self, sink_config: Dict[str, str]) -> bool:
+        """Check if sink config has either topics or topics.regex."""
+        return bool(sink_config.get("topics") or sink_config.get("topics.regex"))
+
+    def _filter_by_topic_list(
+        self, processed_topics: List[str], topics_config: str
+    ) -> List[str]:
+        """Filter topics based on explicit topic list from config."""
+        config_topics = [
+            topic.strip() for topic in topics_config.split(",") if topic.strip()
+        ]
+        return [topic for topic in processed_topics if topic in config_topics]
+
+    def _filter_by_topic_regex(
+        self, processed_topics: List[str], regex_pattern: str
+    ) -> List[str]:
+        """Filter topics based on regex pattern from config."""
+        from java.util.regex import Pattern
+
+        regex_matcher = Pattern.compile(regex_pattern)
+
+        return [
+            topic
+            for topic in processed_topics
+            if regex_matcher.matcher(topic).matches()
+        ]

datahub/ingestion/source/kafka_connect/sink_connectors.py

@@ -175,7 +175,7 @@ class BigQuerySinkConnector(BaseConnector):
     class BQParser:
         project: str
         target_platform: str
-        sanitizeTopics: str
+        sanitizeTopics: bool
         transforms: list
         topicsToTables: Optional[str] = None
         datasets: Optional[str] = None
@@ -187,7 +187,7 @@ class BigQuerySinkConnector(BaseConnector):
         connector_manifest: ConnectorManifest,
     ) -> BQParser:
         project = connector_manifest.config["project"]
-        sanitizeTopics = connector_manifest.config.get("sanitizeTopics", "false")
+        sanitizeTopics = connector_manifest.config.get("sanitizeTopics") or "false"
         transform_names = (
             self.connector_manifest.config.get("transforms", "").split(",")
             if self.connector_manifest.config.get("transforms")

datahub/ingestion/source/kafka_connect/source_connectors.py

@@ -107,9 +107,9 @@ class ConfluentJDBCSourceConnector(BaseConnector):
         assert database_name
         db_connection_url = f"{url_instance.drivername}://{url_instance.host}:{url_instance.port}/{database_name}"
 
-        topic_prefix = self.connector_manifest.config.get("topic.prefix", None)
+        topic_prefix = self.connector_manifest.config.get("topic.prefix") or ""
 
-        query = self.connector_manifest.config.get("query", None)
+        query = self.connector_manifest.config.get("query") or ""
 
         transform_names = (
             self.connector_manifest.config.get("transforms", "").split(",")
@@ -447,13 +447,10 @@ class DebeziumSourceConnector(BaseConnector):
     ) -> DebeziumParser:
         connector_class = connector_manifest.config.get(CONNECTOR_CLASS, "")
 
-        if connector_class == "io.debezium.connector.mysql.MySqlConnector":
-            parser = self.DebeziumParser(
-                source_platform="mysql",
-                server_name=self.get_server_name(connector_manifest),
-                database_name=None,
-            )
-        elif connector_class == "MySqlConnector":
+        if (
+            connector_class == "io.debezium.connector.mysql.MySqlConnector"
+            or connector_class == "MySqlConnector"
+        ):
             parser = self.DebeziumParser(
                 source_platform="mysql",
                 server_name=self.get_server_name(connector_manifest),

datahub/ingestion/source/looker/looker_common.py

@@ -923,7 +923,7 @@ class LookerExplore:
             tags=cast(List, dict.get("tags")) if dict.get("tags") is not None else [],
         )
 
-    @classmethod  # noqa: C901
+    @classmethod
     def from_api(  # noqa: C901
         cls,
         model: str,
@@ -931,7 +931,7 @@ class LookerExplore:
         client: LookerAPI,
         reporter: SourceReport,
         source_config: LookerDashboardSourceConfig,
-    ) -> Optional["LookerExplore"]:  # noqa: C901
+    ) -> Optional["LookerExplore"]:
         try:
             explore = client.lookml_model_explore(model, explore_name)
             views: Set[str] = set()
@@ -1183,7 +1183,7 @@ class LookerExplore:
         base_url = remove_port_from_url(base_url)
         return f"{base_url}/embed/explore/{self.model_name}/{self.name}"
 
-    def _to_metadata_events(  # noqa: C901
+    def _to_metadata_events(
         self,
         config: LookerCommonConfig,
         reporter: SourceReport,

datahub/ingestion/source/looker/looker_file_loader.py

@@ -33,14 +33,14 @@ class LookerViewFileLoader:
         base_projects_folder: Dict[str, pathlib.Path],
         reporter: LookMLSourceReport,
         source_config: LookMLSourceConfig,
-        manifest_constants: Dict[str, LookerConstant] = {},
+        manifest_constants: Optional[Dict[str, LookerConstant]] = None,
     ) -> None:
         self.viewfile_cache: Dict[str, Optional[LookerViewFile]] = {}
         self._root_project_name = root_project_name
         self._base_projects_folder = base_projects_folder
         self.reporter = reporter
         self.source_config = source_config
-        self.manifest_constants = manifest_constants
+        self.manifest_constants = manifest_constants or {}
 
     def _load_viewfile(
         self, project_name: str, path: str, reporter: LookMLSourceReport

datahub/ingestion/source/looker/looker_lib_wrapper.py

@@ -205,8 +205,9 @@ class LookerAPI:
     def folder_ancestors(
         self,
         folder_id: str,
-        fields: Union[str, List[str]] = ["id", "name", "parent_id"],
+        fields: Optional[Union[str, List[str]]] = None,
     ) -> Sequence[Folder]:
+        fields = fields or ["id", "name", "parent_id"]
         self.client_stats.folder_calls += 1
         try:
             return self.client.folder_ancestors(

datahub/ingestion/source/looker/looker_source.py

@@ -383,7 +383,7 @@ class LookerDashboardSource(TestableSource, StatefulIngestionSourceBase):
 
         self.reachable_explores[(model, explore)].append(via)
 
-    def _get_looker_dashboard_element(  # noqa: C901
+    def _get_looker_dashboard_element(
         self, element: DashboardElement
     ) -> Optional[LookerDashboardElement]:
         # Dashboard elements can use raw usage_queries against explores

datahub/ingestion/source/looker/looker_template_language.py

@@ -464,9 +464,10 @@ def process_lookml_template_language(
     source_config: LookMLSourceConfig,
     view_lkml_file_dict: dict,
     reporter: LookMLSourceReport,
-    manifest_constants: Dict[str, "LookerConstant"] = {},
+    manifest_constants: Optional[Dict[str, "LookerConstant"]] = None,
     resolve_constants: bool = False,
 ) -> None:
+    manifest_constants = manifest_constants or {}
     if "views" not in view_lkml_file_dict:
         return
 
@@ -507,9 +508,10 @@ def load_and_preprocess_file(
     path: Union[str, pathlib.Path],
     source_config: LookMLSourceConfig,
     reporter: LookMLSourceReport,
-    manifest_constants: Dict[str, "LookerConstant"] = {},
+    manifest_constants: Optional[Dict[str, "LookerConstant"]] = None,
     resolve_constants: bool = False,
 ) -> dict:
+    manifest_constants = manifest_constants or {}
     parsed = load_lkml(path)
 
     process_lookml_template_language(

datahub/ingestion/source/looker/lookml_source.py

@@ -501,7 +501,7 @@ class LookMLSource(StatefulIngestionSourceBase):
             raise ValueError(
                 f"Could not locate a project name for model {model_name}. Consider configuring a static project name "
                 f"in your config file"
-            )
+            ) from None
 
     def get_manifest_if_present(self, folder: pathlib.Path) -> Optional[LookerManifest]:
         manifest_file = folder / "manifest.lkml"
@@ -1006,8 +1006,9 @@ class LookMLSource(StatefulIngestionSourceBase):
     def report_skipped_unreachable_views(
         self,
         viewfile_loader: LookerViewFileLoader,
-        processed_view_map: Dict[str, Set[str]] = {},
+        processed_view_map: Optional[Dict[str, Set[str]]] = None,
     ) -> None:
+        processed_view_map = processed_view_map or {}
         view_files: Dict[str, List[pathlib.Path]] = {}
         for project, folder_path in self.base_projects_folder.items():
             folder = pathlib.Path(folder_path)