acryl-datahub 1.0.0.4rc7__py3-none-any.whl → 1.0.0.4rc8__py3-none-any.whl

datahub/ingestion/source/mode.py

@@ -1,5 +1,6 @@
 import dataclasses
 import logging
+import os
 import re
 import time
 from dataclasses import dataclass
@@ -9,6 +10,7 @@ from json import JSONDecodeError
 from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union
 
 import dateutil.parser as dp
+import psutil
 import pydantic
 import requests
 import sqlglot
@@ -114,8 +116,12 @@ from datahub.sql_parsing.sqlglot_lineage import (
 )
 from datahub.utilities import config_clean
 from datahub.utilities.lossy_collections import LossyList
+from datahub.utilities.perf_timer import PerfTimer
 
 logger: logging.Logger = logging.getLogger(__name__)
+# Default API limit for items returned per API call
+# Used for the default per_page value for paginated API requests
+DEFAULT_API_ITEMS_PER_PAGE = 30
 
 
 class SpaceKey(ContainerKey):
@@ -194,10 +200,25 @@ class ModeConfig(
         default=True, description="Tag measures and dimensions in the schema"
     )
 
+    items_per_page: int = Field(
+        default=DEFAULT_API_ITEMS_PER_PAGE,
+        description="Number of items per page for paginated API requests.",
+        hidden_from_docs=True,
+    )
+
     @validator("connect_uri")
     def remove_trailing_slash(cls, v):
         return config_clean.remove_trailing_slashes(v)
 
+    @validator("items_per_page")
+    def validate_items_per_page(cls, v):
+        if 1 <= v <= DEFAULT_API_ITEMS_PER_PAGE:
+            return v
+        else:
+            raise ValueError(
+                f"items_per_page must be between 1 and {DEFAULT_API_ITEMS_PER_PAGE}"
+            )
+
 
 class HTTPError429(HTTPError):
     pass
@@ -224,6 +245,20 @@ class ModeSourceReport(StaleEntityRemovalSourceReport):
     num_requests_exceeding_rate_limit: int = 0
     num_requests_retried_on_timeout: int = 0
     num_spaces_retrieved: int = 0
+    space_get_api_called: int = 0
+    report_get_api_called: int = 0
+    dataset_get_api_called: int = 0
+    query_get_api_called: int = 0
+    chart_get_api_called: int = 0
+    get_cache_hits: int = 0
+    get_cache_misses: int = 0
+    get_cache_size: int = 0
+    process_memory_used_mb: float = 0
+    space_get_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
+    report_get_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
+    dataset_get_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
+    query_get_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
+    chart_get_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
 
     def report_dropped_space(self, ent_name: str) -> None:
         self.filtered_spaces.append(ent_name)
@@ -583,34 +618,38 @@ class ModeSource(StatefulIngestionSourceBase):
         space_info = {}
         try:
             logger.debug(f"Retrieving spaces for {self.workspace_uri}")
-            for spaces_page in self._get_paged_request_json(
-                f"{self.workspace_uri}/spaces?filter=all", "spaces", 30
-            ):
-                logger.debug(
-                    f"Read {len(spaces_page)} spaces records from workspace {self.workspace_uri}"
-                )
-                self.report.num_spaces_retrieved += len(spaces_page)
-                for s in spaces_page:
-                    logger.debug(f"Space: {s.get('name')}")
-                    space_name = s.get("name", "")
-                    # Using both restricted and default_access_level because
-                    # there is a current bug with restricted returning False everytime
-                    # which has been reported to Mode team
-                    if self.config.exclude_restricted and (
-                        s.get("restricted")
-                        or s.get("default_access_level") == "restricted"
-                    ):
-                        logging.debug(
-                            f"Skipping space {space_name} due to exclude restricted"
-                        )
-                        continue
-                    if not self.config.space_pattern.allowed(space_name):
-                        self.report.report_dropped_space(space_name)
-                        logging.debug(
-                            f"Skipping space {space_name} due to space pattern"
-                        )
-                        continue
-                    space_info[s.get("token", "")] = s.get("name", "")
+            with self.report.space_get_timer:
+                for spaces_page in self._get_paged_request_json(
+                    f"{self.workspace_uri}/spaces?filter=all",
+                    "spaces",
+                    self.config.items_per_page,
+                ):
+                    self.report.space_get_api_called += 1
+                    logger.debug(
+                        f"Read {len(spaces_page)} spaces records from workspace {self.workspace_uri}"
+                    )
+                    self.report.num_spaces_retrieved += len(spaces_page)
+                    for s in spaces_page:
+                        logger.debug(f"Space: {s.get('name')}")
+                        space_name = s.get("name", "")
+                        # Using both restricted and default_access_level because
+                        # there is a current bug with restricted returning False everytime
+                        # which has been reported to Mode team
+                        if self.config.exclude_restricted and (
+                            s.get("restricted")
+                            or s.get("default_access_level") == "restricted"
+                        ):
+                            logging.debug(
+                                f"Skipping space {space_name} due to exclude restricted"
+                            )
+                            continue
+                        if not self.config.space_pattern.allowed(space_name):
+                            self.report.report_dropped_space(space_name)
+                            logging.debug(
+                                f"Skipping space {space_name} due to space pattern"
+                            )
+                            continue
+                        space_info[s.get("token", "")] = s.get("name", "")
         except ModeRequestError as e:
             self.report.report_failure(
                 title="Failed to Retrieve Spaces",
@@ -1414,48 +1453,61 @@ class ModeSource(StatefulIngestionSourceBase):
         mce = MetadataChangeEvent(proposedSnapshot=chart_snapshot)
         yield MetadataWorkUnit(id=chart_snapshot.urn, mce=mce)
 
-    @lru_cache(maxsize=None)
-    def _get_reports(self, space_token: str) -> List[dict]:
-        reports = []
+    def _get_reports(self, space_token: str) -> Iterator[List[dict]]:
         try:
-            reports_json = self._get_request_json(
-                f"{self.workspace_uri}/spaces/{space_token}/reports"
-            )
-            reports = reports_json.get("_embedded", {}).get("reports", {})
+            with self.report.report_get_timer:
+                for reports_page in self._get_paged_request_json(
+                    f"{self.workspace_uri}/spaces/{space_token}/reports?filter=all",
+                    "reports",
+                    self.config.items_per_page,
+                ):
+                    self.report.report_get_api_called += 1
+                    logger.debug(
+                        f"Read {len(reports_page)} reports records from workspace {self.workspace_uri} space {space_token}"
+                    )
+                    yield reports_page
         except ModeRequestError as e:
             self.report.report_failure(
                 title="Failed to Retrieve Reports for Space",
                 message="Unable to retrieve reports for space token.",
                 context=f"Space Token: {space_token}, Error: {str(e)}",
             )
-        return reports
 
-    @lru_cache(maxsize=None)
-    def _get_datasets(self, space_token: str) -> List[dict]:
+    def _get_datasets(self, space_token: str) -> Iterator[List[dict]]:
         """
         Retrieves datasets for a given space token.
         """
-        datasets = []
         try:
-            url = f"{self.workspace_uri}/spaces/{space_token}/datasets"
-            datasets_json = self._get_request_json(url)
-            datasets = datasets_json.get("_embedded", {}).get("reports", [])
+            with self.report.dataset_get_timer:
+                for dataset_page in self._get_paged_request_json(
+                    f"{self.workspace_uri}/spaces/{space_token}/datasets?filter=all",
+                    "reports",
+                    self.config.items_per_page,
+                ):
+                    self.report.dataset_get_api_called += 1
+                    logger.debug(
+                        f"Read {len(dataset_page)} datasets records from workspace {self.workspace_uri} space {space_token}"
+                    )
+                    yield dataset_page
         except ModeRequestError as e:
             self.report.report_failure(
                 title="Failed to Retrieve Datasets for Space",
                 message=f"Unable to retrieve datasets for space token {space_token}.",
                 context=f"Error: {str(e)}",
             )
-        return datasets
 
-    @lru_cache(maxsize=None)
-    def _get_queries(self, report_token: str) -> list:
-        queries = []
+    def _get_queries(self, report_token: str) -> List[dict]:
         try:
-            queries_json = self._get_request_json(
-                f"{self.workspace_uri}/reports/{report_token}/queries"
-            )
-            queries = queries_json.get("_embedded", {}).get("queries", {})
+            with self.report.query_get_timer:
+                # This endpoint does not handle pagination properly
+                queries = self._get_request_json(
+                    f"{self.workspace_uri}/reports/{report_token}/queries"
+                )
+                self.report.query_get_api_called += 1
+                logger.debug(
+                    f"Read {len(queries)} queries records from workspace {self.workspace_uri} report {report_token}"
+                )
+                return queries.get("_embedded", {}).get("queries", [])
         except ModeRequestError as e:
             if isinstance(e, HTTPError) and e.response.status_code == 404:
                 self.report.report_warning(
@@ -1469,35 +1521,39 @@ class ModeSource(StatefulIngestionSourceBase):
                     message="Unable to retrieve queries for report token.",
                     context=f"Report Token: {report_token}, Error: {str(e)}",
                 )
-        return queries
+            return []
 
     @lru_cache(maxsize=None)
-    def _get_last_query_run(
-        self, report_token: str, report_run_id: str, query_run_id: str
-    ) -> Dict:
+    def _get_last_query_run(self, report_token: str, report_run_id: str) -> list:
+        # This function is unused and may be subject to removal in a future revision of this source
+        query_runs = []
         try:
-            queries_json = self._get_request_json(
-                f"{self.workspace_uri}/reports/{report_token}/runs/{report_run_id}/query_runs{query_run_id}"
-            )
-            queries = queries_json.get("_embedded", {}).get("queries", {})
+            for query_run_page in self._get_paged_request_json(
+                f"{self.workspace_uri}/reports/{report_token}/runs/{report_run_id}/query_runs?filter=all",
+                "query_runs",
+                self.config.items_per_page,
+            ):
+                query_runs.extend(query_run_page)
         except ModeRequestError as e:
             self.report.report_failure(
                 title="Failed to Retrieve Queries for Report",
                 message="Unable to retrieve queries for report token.",
                 context=f"Report Token:{report_token}, Error: {str(e)}",
             )
-            return {}
-        return queries
+        return query_runs
 
-    @lru_cache(maxsize=None)
-    def _get_charts(self, report_token: str, query_token: str) -> list:
-        charts = []
+    def _get_charts(self, report_token: str, query_token: str) -> List[dict]:
         try:
-            charts_json = self._get_request_json(
-                f"{self.workspace_uri}/reports/{report_token}"
-                f"/queries/{query_token}/charts"
-            )
-            charts = charts_json.get("_embedded", {}).get("charts", {})
+            with self.report.chart_get_timer:
+                # This endpoint does not handle pagination properly
+                charts = self._get_request_json(
+                    f"{self.workspace_uri}/reports/{report_token}/queries/{query_token}/charts"
+                )
+                self.report.chart_get_api_called += 1
+                logger.debug(
+                    f"Read {len(charts)} charts records from workspace {self.workspace_uri} report {report_token} query {query_token}"
+                )
+                return charts.get("_embedded", {}).get("charts", [])
         except ModeRequestError as e:
             self.report.report_failure(
                 title="Failed to Retrieve Charts",
@@ -1506,7 +1562,7 @@ class ModeSource(StatefulIngestionSourceBase):
                 f"Query token: {query_token}, "
                 f"Error: {str(e)}",
             )
-        return charts
+            return []
 
     def _get_paged_request_json(
         self, url: str, key: str, per_page: int
@@ -1521,6 +1577,7 @@ class ModeSource(StatefulIngestionSourceBase):
             yield data
             page += 1
 
+    @lru_cache(maxsize=20480)
     def _get_request_json(self, url: str) -> Dict:
         r = tenacity.Retrying(
             wait=wait_exponential(
@@ -1568,6 +1625,17 @@ class ModeSource(StatefulIngestionSourceBase):
 
         return get_request()
 
+    @staticmethod
+    def _get_process_memory():
+        process = psutil.Process(os.getpid())
+        mem_info = process.memory_info()
+        return {
+            "rss": mem_info.rss / (1024 * 1024),
+            "vms": mem_info.vms / (1024 * 1024),
+            "shared": getattr(mem_info, "shared", 0) / (1024 * 1024),
+            "data": getattr(mem_info, "data", 0) / (1024 * 1024),
+        }
+
     @staticmethod
     def create_embed_aspect_mcp(
         entity_urn: str, embed_url: str
@@ -1603,115 +1671,116 @@ class ModeSource(StatefulIngestionSourceBase):
             yield from self.construct_space_container(space_token, space_name)
             space_container_key = self.gen_space_key(space_token)
 
-            reports = self._get_reports(space_token)
-            for report in reports:
-                logger.debug(
-                    f"Report: name: {report.get('name')} token: {report.get('token')}"
-                )
-                dashboard_tuple_from_report = self.construct_dashboard(
-                    space_token=space_token, report_info=report
-                )
-
-                if dashboard_tuple_from_report is None:
-                    continue
-                (
-                    dashboard_snapshot_from_report,
-                    browse_mcpw,
-                ) = dashboard_tuple_from_report
+            for report_page in self._get_reports(space_token):
+                for report in report_page:
+                    logger.debug(
+                        f"Report: name: {report.get('name')} token: {report.get('token')}"
+                    )
+                    dashboard_tuple_from_report = self.construct_dashboard(
+                        space_token=space_token, report_info=report
+                    )
 
-                mce = MetadataChangeEvent(
-                    proposedSnapshot=dashboard_snapshot_from_report
-                )
+                    if dashboard_tuple_from_report is None:
+                        continue
+                    (
+                        dashboard_snapshot_from_report,
+                        browse_mcpw,
+                    ) = dashboard_tuple_from_report
 
-                mcpw = MetadataChangeProposalWrapper(
-                    entityUrn=dashboard_snapshot_from_report.urn,
-                    aspect=SubTypesClass(typeNames=[BIAssetSubTypes.MODE_REPORT]),
-                )
-                yield mcpw.as_workunit()
-                yield from add_dataset_to_container(
-                    container_key=space_container_key,
-                    dataset_urn=dashboard_snapshot_from_report.urn,
-                )
-                yield browse_mcpw.as_workunit()
+                    mce = MetadataChangeEvent(
+                        proposedSnapshot=dashboard_snapshot_from_report
+                    )
 
-                usage_statistics = DashboardUsageStatisticsClass(
-                    timestampMillis=round(datetime.now().timestamp() * 1000),
-                    viewsCount=report.get("view_count", 0),
-                )
+                    mcpw = MetadataChangeProposalWrapper(
+                        entityUrn=dashboard_snapshot_from_report.urn,
+                        aspect=SubTypesClass(typeNames=[BIAssetSubTypes.MODE_REPORT]),
+                    )
+                    yield mcpw.as_workunit()
+                    yield from add_dataset_to_container(
+                        container_key=space_container_key,
+                        dataset_urn=dashboard_snapshot_from_report.urn,
+                    )
+                    yield browse_mcpw.as_workunit()
 
-                yield MetadataChangeProposalWrapper(
-                    entityUrn=dashboard_snapshot_from_report.urn,
-                    aspect=usage_statistics,
-                ).as_workunit()
+                    usage_statistics = DashboardUsageStatisticsClass(
+                        timestampMillis=round(datetime.now().timestamp() * 1000),
+                        viewsCount=report.get("view_count", 0),
+                    )
 
-                if self.config.ingest_embed_url is True:
-                    yield self.create_embed_aspect_mcp(
-                        entity_urn=dashboard_snapshot_from_report.urn,
-                        embed_url=f"{self.config.connect_uri}/{self.config.workspace}/reports/{report.get('token')}/embed",
+                    yield MetadataChangeProposalWrapper(
+                        entityUrn=dashboard_snapshot_from_report.urn,
+                        aspect=usage_statistics,
                     ).as_workunit()
 
-                yield MetadataWorkUnit(id=dashboard_snapshot_from_report.urn, mce=mce)
+                    if self.config.ingest_embed_url is True:
+                        yield self.create_embed_aspect_mcp(
+                            entity_urn=dashboard_snapshot_from_report.urn,
+                            embed_url=f"{self.config.connect_uri}/{self.config.workspace}/reports/{report.get('token')}/embed",
+                        ).as_workunit()
+
+                    yield MetadataWorkUnit(
+                        id=dashboard_snapshot_from_report.urn, mce=mce
+                    )
 
     def emit_chart_mces(self) -> Iterable[MetadataWorkUnit]:
         # Space/collection -> report -> query -> Chart
         for space_token in self.space_tokens:
-            reports = self._get_reports(space_token)
-            for report in reports:
-                report_token = report.get("token", "")
-
-                queries = self._get_queries(report_token)
-                for query in queries:
-                    query_mcps = self.construct_query_or_dataset(
-                        report_token,
-                        query,
-                        space_token=space_token,
-                        report_info=report,
-                        is_mode_dataset=False,
-                    )
-                    chart_fields: Dict[str, SchemaFieldClass] = {}
-                    for wu in query_mcps:
-                        if isinstance(
-                            wu.metadata, MetadataChangeProposalWrapper
-                        ) and isinstance(wu.metadata.aspect, SchemaMetadataClass):
-                            schema_metadata = wu.metadata.aspect
-                            for field in schema_metadata.fields:
-                                chart_fields.setdefault(field.fieldPath, field)
-
-                        yield wu
-
-                    charts = self._get_charts(report_token, query.get("token", ""))
-                    # build charts
-                    for i, chart in enumerate(charts):
-                        yield from self.construct_chart_from_api_data(
-                            i,
-                            chart,
-                            chart_fields,
+            for report_page in self._get_reports(space_token):
+                for report in report_page:
+                    report_token = report.get("token", "")
+
+                    queries = self._get_queries(report_token)
+                    for query in queries:
+                        query_mcps = self.construct_query_or_dataset(
+                            report_token,
                             query,
                             space_token=space_token,
                             report_info=report,
-                            query_name=query["name"],
+                            is_mode_dataset=False,
                         )
+                        chart_fields: Dict[str, SchemaFieldClass] = {}
+                        for wu in query_mcps:
+                            if isinstance(
+                                wu.metadata, MetadataChangeProposalWrapper
+                            ) and isinstance(wu.metadata.aspect, SchemaMetadataClass):
+                                schema_metadata = wu.metadata.aspect
+                                for field in schema_metadata.fields:
+                                    chart_fields.setdefault(field.fieldPath, field)
+
+                            yield wu
+
+                        charts = self._get_charts(report_token, query.get("token", ""))
+                        # build charts
+                        for i, chart in enumerate(charts):
+                            yield from self.construct_chart_from_api_data(
+                                i,
+                                chart,
+                                chart_fields,
+                                query,
+                                space_token=space_token,
+                                report_info=report,
+                                query_name=query["name"],
+                            )
 
     def emit_dataset_mces(self):
         """
         Emits MetadataChangeEvents (MCEs) for datasets within each space.
         """
         for space_token, _ in self.space_tokens.items():
-            datasets = self._get_datasets(space_token)
-
-            for report in datasets:
-                report_token = report.get("token", "")
-                queries = self._get_queries(report_token)
-                for query in queries:
-                    query_mcps = self.construct_query_or_dataset(
-                        report_token,
-                        query,
-                        space_token=space_token,
-                        report_info=report,
-                        is_mode_dataset=True,
-                    )
-                    for wu in query_mcps:
-                        yield wu
+            for dataset_page in self._get_datasets(space_token):
+                for report in dataset_page:
+                    report_token = report.get("token", "")
+                    queries = self._get_queries(report_token)
+                    for query in queries:
+                        query_mcps = self.construct_query_or_dataset(
+                            report_token,
+                            query,
+                            space_token=space_token,
+                            report_info=report,
+                            is_mode_dataset=True,
+                        )
+                        for wu in query_mcps:
+                            yield wu
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "ModeSource":
@@ -1730,6 +1799,12 @@ class ModeSource(StatefulIngestionSourceBase):
         yield from self.emit_dashboard_mces()
         yield from self.emit_dataset_mces()
         yield from self.emit_chart_mces()
+        cache_info = self._get_request_json.cache_info()
+        self.report.get_cache_hits = cache_info.hits
+        self.report.get_cache_misses = cache_info.misses
+        self.report.get_cache_size = cache_info.currsize
+        memory_used = self._get_process_memory()
+        self.report.process_memory_used_mb = round(memory_used["rss"], 2)
 
     def get_report(self) -> SourceReport:
         return self.report

datahub/ingestion/source/sql/mssql/source.py

@@ -13,6 +13,7 @@ from sqlalchemy.exc import ProgrammingError, ResourceClosedError
 
 import datahub.metadata.schema_classes as models
 from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.pattern_utils import UUID_REGEX
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
@@ -60,6 +61,15 @@ register_custom_type(sqlalchemy.dialects.mssql.SMALLMONEY, models.NumberTypeClas
 register_custom_type(sqlalchemy.dialects.mssql.SQL_VARIANT, models.UnionTypeClass)
 register_custom_type(sqlalchemy.dialects.mssql.UNIQUEIDENTIFIER, models.StringTypeClass)
 
+# Patterns copied from Snowflake source
+DEFAULT_TEMP_TABLES_PATTERNS = [
+    r".*\.FIVETRAN_.*_STAGING\..*",  # fivetran
+    r".*__DBT_TMP$",  # dbt
+    rf".*\.SEGMENT_{UUID_REGEX}",  # segment
+    rf".*\.STAGING_.*_{UUID_REGEX}",  # stitch
+    r".*\.(GE_TMP_|GE_TEMP_|GX_TEMP_)[0-9A-F]{8}",  # great expectations
+]
+
 
 class SQLServerConfig(BasicSQLAlchemyConfig):
     # defaults
@@ -114,6 +124,12 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
         default=False,
         description="Enable the container aspects ingestion for both pipelines and tasks. Note that this feature requires the corresponding model support in the backend, which was introduced in version 0.15.0.1.",
     )
+    temporary_tables_pattern: List[str] = Field(
+        default=DEFAULT_TEMP_TABLES_PATTERNS,
+        description="[Advanced] Regex patterns for temporary tables to filter in lineage ingestion. Specify regex to "
+        "match the entire table name in database.schema.table format. Defaults are to set in such a way "
+        "to ignore the temporary staging tables created by known ETL tools.",
+    )
 
     @pydantic.validator("uri_args")
     def passwords_match(cls, v, values, **kwargs):
@@ -179,6 +195,14 @@ class SQLServerSource(SQLAlchemySource):
         self.table_descriptions: Dict[str, str] = {}
         self.column_descriptions: Dict[str, str] = {}
         self.stored_procedures: FileBackedList[StoredProcedure] = FileBackedList()
+
+        self.report = SQLSourceReport()
+        if self.config.include_lineage and not self.config.convert_urns_to_lowercase:
+            self.report.warning(
+                title="Potential issue with lineage",
+                message="Lineage may not resolve accurately because 'convert_urns_to_lowercase' is False. To ensure lineage correct, set 'convert_urns_to_lowercase' to True.",
+            )
+
         if self.config.include_descriptions:
             for inspector in self.get_inspectors():
                 db_name: str = self.get_db_name(inspector)
@@ -774,6 +798,13 @@ class SQLServerSource(SQLAlchemySource):
                 )
 
     def is_temp_table(self, name: str) -> bool:
+        if any(
+            re.match(pattern, name, flags=re.IGNORECASE)
+            for pattern in self.config.temporary_tables_pattern
+        ):
+            logger.debug(f"temp table matched by pattern {name}")
+            return True
+
         try:
             parts = name.split(".")
             table_name = parts[-1]