acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/grafana/lineage.py

@@ -0,0 +1,202 @@
+import logging
+from typing import Dict, List, Optional, Tuple
+
+from datahub.emitter.mce_builder import (
+    make_dataset_urn_with_platform_instance,
+    make_schema_field_urn,
+)
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.source.grafana.grafana_config import PlatformConnectionConfig
+from datahub.ingestion.source.grafana.models import (
+    DatasourceRef,
+    GrafanaQueryTarget,
+    Panel,
+)
+from datahub.ingestion.source.grafana.report import GrafanaSourceReport
+from datahub.metadata.schema_classes import (
+    DatasetLineageTypeClass,
+    FineGrainedLineageClass,
+    FineGrainedLineageDownstreamTypeClass,
+    FineGrainedLineageUpstreamTypeClass,
+    UpstreamClass,
+    UpstreamLineageClass,
+)
+from datahub.sql_parsing.sqlglot_lineage import (
+    SqlParsingResult,
+    create_lineage_sql_parsed_result,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class LineageExtractor:
+    """Handles extraction of lineage information from Grafana panels"""
+
+    def __init__(
+        self,
+        platform: str,
+        platform_instance: Optional[str],
+        env: str,
+        connection_to_platform_map: Dict[str, PlatformConnectionConfig],
+        report: GrafanaSourceReport,
+        graph: Optional[DataHubGraph] = None,
+        include_column_lineage: bool = True,
+    ):
+        self.platform = platform
+        self.platform_instance = platform_instance
+        self.env = env
+        self.connection_map = connection_to_platform_map
+        self.graph = graph
+        self.report = report
+        self.include_column_lineage = include_column_lineage
+
+    def extract_panel_lineage(
+        self, panel: Panel
+    ) -> Optional[MetadataChangeProposalWrapper]:
+        """Extract lineage information from a panel."""
+        if not panel.datasource_ref:
+            return None
+
+        ds_type, ds_uid = self._extract_datasource_info(panel.datasource_ref)
+        raw_sql = self._extract_raw_sql(panel.query_targets)
+        ds_urn = self._build_dataset_urn(ds_type, ds_uid, panel.id)
+
+        # Handle platform-specific lineage
+        if ds_uid in self.connection_map:
+            if raw_sql:
+                parsed_sql = self._parse_sql(raw_sql, self.connection_map[ds_uid])
+                if parsed_sql:
+                    lineage = self._create_column_lineage(ds_urn, parsed_sql)
+                    if lineage:
+                        return lineage
+
+            # Fall back to basic lineage if SQL parsing fails or no column lineage created
+            return self._create_basic_lineage(
+                ds_uid, self.connection_map[ds_uid], ds_urn
+            )
+
+        return None
+
+    def _extract_datasource_info(
+        self, datasource_ref: "DatasourceRef"
+    ) -> Tuple[str, str]:
+        """Extract datasource type and UID."""
+        return datasource_ref.type or "unknown", datasource_ref.uid or "unknown"
+
+    def _extract_raw_sql(
+        self, query_targets: List["GrafanaQueryTarget"]
+    ) -> Optional[str]:
+        """Extract raw SQL from panel query targets."""
+        for target in query_targets:
+            if target.get("rawSql"):
+                return target["rawSql"]
+        return None
+
+    def _build_dataset_urn(self, ds_type: str, ds_uid: str, panel_id: str) -> str:
+        """Build dataset URN."""
+        dataset_name = f"{ds_type}.{ds_uid}.{panel_id}"
+        return make_dataset_urn_with_platform_instance(
+            platform=self.platform,
+            name=dataset_name,
+            platform_instance=self.platform_instance,
+            env=self.env,
+        )
+
+    def _create_basic_lineage(
+        self, ds_uid: str, platform_config: PlatformConnectionConfig, ds_urn: str
+    ) -> MetadataChangeProposalWrapper:
+        """Create basic upstream lineage."""
+        name = (
+            f"{platform_config.database}.{ds_uid}"
+            if platform_config.database
+            else ds_uid
+        )
+
+        upstream_urn = make_dataset_urn_with_platform_instance(
+            platform=platform_config.platform,
+            name=name,
+            platform_instance=platform_config.platform_instance,
+            env=platform_config.env,
+        )
+
+        logger.info(f"Generated upstream URN: {upstream_urn}")
+
+        return MetadataChangeProposalWrapper(
+            entityUrn=ds_urn,
+            aspect=UpstreamLineageClass(
+                upstreams=[
+                    UpstreamClass(
+                        dataset=upstream_urn,
+                        type=DatasetLineageTypeClass.TRANSFORMED,
+                    )
+                ]
+            ),
+        )
+
+    def _parse_sql(
+        self, sql: str, platform_config: PlatformConnectionConfig
+    ) -> Optional[SqlParsingResult]:
+        """Parse SQL query for lineage information."""
+        if not self.graph:
+            logger.warning("No DataHub graph specified for SQL parsing.")
+            return None
+
+        try:
+            return create_lineage_sql_parsed_result(
+                query=sql,
+                platform=platform_config.platform,
+                platform_instance=platform_config.platform_instance,
+                env=platform_config.env,
+                default_db=platform_config.database,
+                default_schema=platform_config.database_schema,
+                graph=self.graph,
+            )
+        except ValueError as e:
+            logger.error(f"SQL parsing error for query: {sql}", exc_info=e)
+        except Exception as e:
+            logger.exception(f"Unexpected error during SQL parsing: {sql}", exc_info=e)
+
+        return None
+
+    def _create_column_lineage(
+        self,
+        dataset_urn: str,
+        parsed_sql: SqlParsingResult,
+    ) -> Optional[MetadataChangeProposalWrapper]:
+        """Create column-level lineage"""
+        if not parsed_sql.column_lineage or not self.include_column_lineage:
+            return None
+
+        upstream_lineages = []
+        for col_lineage in parsed_sql.column_lineage:
+            upstream_lineages.append(
+                FineGrainedLineageClass(
+                    downstreamType=FineGrainedLineageDownstreamTypeClass.FIELD,
+                    downstreams=[
+                        make_schema_field_urn(
+                            dataset_urn, col_lineage.downstream.column
+                        )
+                    ],
+                    upstreamType=FineGrainedLineageUpstreamTypeClass.FIELD_SET,
+                    upstreams=[
+                        make_schema_field_urn(upstream_dataset, col.column)
+                        for col in col_lineage.upstreams
+                        for upstream_dataset in parsed_sql.in_tables
+                    ],
+                )
+            )
+
+        return MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn,
+            aspect=UpstreamLineageClass(
+                upstreams=[
+                    UpstreamClass(
+                        dataset=table,
+                        type=DatasetLineageTypeClass.TRANSFORMED,
+                    )
+                    for table in parsed_sql.in_tables
+                ],
+                fineGrainedLineages=upstream_lineages,
+            ),
+        )

datahub/ingestion/source/grafana/models.py

@@ -0,0 +1,137 @@
+"""Grafana data models for DataHub ingestion.
+
+References:
+- Grafana HTTP API: https://grafana.com/docs/grafana/latest/developers/http_api/
+- Dashboard API: https://grafana.com/docs/grafana/latest/developers/http_api/dashboard/
+- Folder API: https://grafana.com/docs/grafana/latest/developers/http_api/folder/
+- Search API: https://grafana.com/docs/grafana/latest/developers/http_api/other/#search-api
+- Dashboard JSON structure: https://grafana.com/docs/grafana/latest/dashboards/build-dashboards/view-dashboard-json-model/
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from datahub.emitter.mcp_builder import ContainerKey
+
+logger = logging.getLogger(__name__)
+# Grafana-specific type definitions for better type safety
+GrafanaQueryTarget = Dict[
+    str, Any
+]  # Query targets: refId, expr/query, datasource, hide, etc.
+GrafanaFieldConfig = Dict[
+    str, Any
+]  # Field config: defaults, overrides, display settings
+GrafanaTransformation = Dict[str, Any]  # Transformations: id, options
+
+
+class _GrafanaBaseModel(BaseModel):
+    model_config = ConfigDict(coerce_numbers_to_str=True)
+
+
+class DatasourceRef(_GrafanaBaseModel):
+    """Reference to a Grafana datasource."""
+
+    type: Optional[str] = None  # Datasource type (prometheus, mysql, postgres, etc.)
+    uid: Optional[str] = None  # Datasource unique identifier
+    name: Optional[str] = None  # Datasource display name
+
+
+class Panel(_GrafanaBaseModel):
+    """Represents a Grafana dashboard panel."""
+
+    id: str
+    title: str
+    description: str = ""
+    type: Optional[str] = None
+    # Query targets - each contains refId (A,B,C...), query/expr, datasource ref, etc.
+    query_targets: List[GrafanaQueryTarget] = Field(
+        default_factory=list, alias="targets"
+    )
+    # Datasource reference - contains type, uid, name
+    datasource_ref: Optional[DatasourceRef] = Field(default=None, alias="datasource")
+    # Field configuration - display settings, defaults, overrides
+    field_config: GrafanaFieldConfig = Field(default_factory=dict, alias="fieldConfig")
+    # Data transformations - each contains id and transformation-specific options
+    transformations: List[GrafanaTransformation] = Field(default_factory=list)
+
+
+class Dashboard(_GrafanaBaseModel):
+    """Represents a Grafana dashboard."""
+
+    uid: str
+    title: str
+    description: str = ""
+    version: Optional[str] = None
+    panels: List[Panel]
+    tags: List[str]
+    timezone: Optional[str] = None
+    refresh: Optional[str] = None
+    schema_version: Optional[str] = Field(default=None, alias="schemaVersion")
+    folder_id: Optional[str] = Field(default=None, alias="meta.folderId")
+    created_by: Optional[str] = None
+
+    @staticmethod
+    def extract_panels(panels_data: List[Dict[str, Any]]) -> List[Panel]:
+        """Extract panels, including nested ones."""
+        panels: List[Panel] = []
+        for panel_data in panels_data:
+            if panel_data.get("type") == "row" and "panels" in panel_data:
+                panels.extend(
+                    Panel.parse_obj(p)
+                    for p in panel_data["panels"]
+                    if p.get("type") != "row"
+                )
+            elif panel_data.get("type") != "row":
+                panels.append(Panel.parse_obj(panel_data))
+        return panels
+
+    @classmethod
+    def parse_obj(cls, data: Dict[str, Any]) -> "Dashboard":
+        """Custom parsing to handle nested panel extraction."""
+        dashboard_data = data.get("dashboard", {})
+        _panel_data = dashboard_data.get("panels", [])
+        panels = []
+        try:
+            panels = cls.extract_panels(_panel_data)
+        except Exception as e:
+            logger.warning(
+                f"Error extracting panels from dashboard for dashboard panels {_panel_data} : {e}"
+            )
+
+        # Extract meta.folderId from nested structure
+        meta = dashboard_data.get("meta", {})
+        folder_id = meta.get("folderId")
+
+        # Create dashboard data without meta to avoid conflicts
+        dashboard_dict = {**dashboard_data, "panels": panels, "folder_id": folder_id}
+        if "meta" in dashboard_dict:
+            del dashboard_dict["meta"]
+
+        # Handle refresh field type mismatch - convert boolean to string
+        if "refresh" in dashboard_dict and isinstance(dashboard_dict["refresh"], bool):
+            dashboard_dict["refresh"] = str(dashboard_dict["refresh"])
+
+        return super().parse_obj(dashboard_dict)
+
+
+class Folder(_GrafanaBaseModel):
+    """Represents a Grafana folder."""
+
+    id: str
+    title: str
+    description: Optional[str] = ""
+
+
+class FolderKey(ContainerKey):
+    """Key for identifying a Grafana folder."""
+
+    folder_id: str
+
+
+class DashboardContainerKey(ContainerKey):
+    """Key for identifying a Grafana dashboard."""
+
+    dashboard_id: str
+    folder_id: Optional[str] = None  # Reference to parent folder

datahub/ingestion/source/grafana/report.py

@@ -0,0 +1,90 @@
+from dataclasses import dataclass
+
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalSourceReport,
+)
+
+
+@dataclass
+class GrafanaSourceReport(StaleEntityRemovalSourceReport):
+    # Entity counters
+    dashboards_scanned: int = 0
+    charts_scanned: int = 0
+    folders_scanned: int = 0
+    datasets_scanned: int = 0
+
+    # Lineage counters
+    panels_with_lineage: int = 0
+    panels_without_lineage: int = 0
+    lineage_extraction_failures: int = 0
+    sql_parsing_attempts: int = 0
+    sql_parsing_successes: int = 0
+    sql_parsing_failures: int = 0
+
+    # Schema extraction counters
+    panels_with_schema_fields: int = 0
+    panels_without_schema_fields: int = 0
+
+    # Warning counters
+    permission_warnings: int = 0
+    datasource_warnings: int = 0
+    panel_parsing_warnings: int = 0
+
+    def report_dashboard_scanned(self) -> None:
+        self.dashboards_scanned += 1
+
+    def report_chart_scanned(self) -> None:
+        self.charts_scanned += 1
+
+    def report_folder_scanned(self) -> None:
+        self.folders_scanned += 1
+
+    def report_dataset_scanned(self) -> None:
+        self.datasets_scanned += 1
+
+    # Lineage reporting methods
+    def report_lineage_extracted(self) -> None:
+        """Report successful lineage extraction for a panel"""
+        self.panels_with_lineage += 1
+
+    def report_no_lineage(self) -> None:
+        """Report that no lineage was found for a panel"""
+        self.panels_without_lineage += 1
+
+    def report_lineage_extraction_failure(self) -> None:
+        """Report failure to extract lineage for a panel"""
+        self.lineage_extraction_failures += 1
+
+    def report_sql_parsing_attempt(self) -> None:
+        """Report attempt to parse SQL"""
+        self.sql_parsing_attempts += 1
+
+    def report_sql_parsing_success(self) -> None:
+        """Report successful SQL parsing"""
+        self.sql_parsing_successes += 1
+
+    def report_sql_parsing_failure(self) -> None:
+        """Report failed SQL parsing"""
+        self.sql_parsing_failures += 1
+
+    # Schema field reporting methods
+    def report_schema_fields_extracted(self) -> None:
+        """Report that schema fields were extracted for a panel"""
+        self.panels_with_schema_fields += 1
+
+    def report_no_schema_fields(self) -> None:
+        """Report that no schema fields were found for a panel"""
+        self.panels_without_schema_fields += 1
+
+    # Warning reporting methods
+    def report_permission_warning(self) -> None:
+        """Report a permission-related warning"""
+        self.permission_warnings += 1
+
+    def report_datasource_warning(self) -> None:
+        """Report a datasource-related warning"""
+        self.datasource_warnings += 1
+
+    def report_panel_parsing_warning(self) -> None:
+        """Report a panel parsing warning"""
+        self.panel_parsing_warnings += 1

datahub/ingestion/source/grafana/types.py

@@ -0,0 +1,16 @@
+from datahub.metadata.schema_classes import (
+    ChartTypeClass,
+)
+
+CHART_TYPE_MAPPINGS = {
+    "graph": ChartTypeClass.LINE,
+    "timeseries": ChartTypeClass.LINE,
+    "table": ChartTypeClass.TABLE,
+    "stat": ChartTypeClass.TEXT,
+    "gauge": ChartTypeClass.TEXT,
+    "bargauge": ChartTypeClass.TEXT,
+    "bar": ChartTypeClass.BAR,
+    "pie": ChartTypeClass.PIE,
+    "heatmap": ChartTypeClass.TABLE,
+    "histogram": ChartTypeClass.BAR,
+}

datahub/ingestion/source/hex/api.py

@@ -5,7 +5,9 @@ from typing import Any, Dict, Generator, List, Optional, Union
 
 import requests
 from pydantic import BaseModel, Field, ValidationError, validator
+from requests.adapters import HTTPAdapter
 from typing_extensions import assert_never
+from urllib3.util.retry import Retry
 
 from datahub.ingestion.api.source import SourceReport
 from datahub.ingestion.source.hex.constants import (
@@ -220,6 +222,7 @@ class HexApi:
         self.base_url = base_url
         self.report = report
         self.page_size = page_size
+        self.session = self._create_retry_session()
 
     def _list_projects_url(self):
         return f"{self.base_url}/projects"
@@ -227,6 +230,28 @@ class HexApi:
     def _auth_header(self):
         return {"Authorization": f"Bearer {self.token}"}
 
+    def _create_retry_session(self) -> requests.Session:
+        """Create a requests session with retry logic for rate limiting.
+
+        Hex API rate limit: 60 requests per minute
+        https://learn.hex.tech/docs/api/api-overview#kernel-and-rate-limits
+        """
+        session = requests.Session()
+
+        # Configure retry strategy for 429 (Too Many Requests) with exponential backoff
+        retry_strategy = Retry(
+            total=5,  # Maximum number of retries
+            status_forcelist=[429],  # Only retry on 429 status code
+            backoff_factor=2,  # Exponential backoff: 2, 4, 8, 16, 32 seconds
+            raise_on_status=True,  # Raise exception after max retries
+        )
+
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        session.mount("http://", adapter)
+        session.mount("https://", adapter)
+
+        return session
+
     def fetch_projects(
         self,
         include_components: bool = True,
@@ -259,7 +284,7 @@ class HexApi:
         logger.debug(f"Fetching projects page with params: {params}")
         self.report.fetch_projects_page_calls += 1
         try:
-            response = requests.get(
+            response = self.session.get(
                 url=self._list_projects_url(),
                 headers=self._auth_header(),
                 params=params,
@@ -350,6 +375,7 @@ class HexApi:
                 description=hex_item.description,
                 created_at=hex_item.created_at,
                 last_edited_at=hex_item.last_edited_at,
+                last_published_at=hex_item.last_published_at,
                 status=status,
                 categories=categories,
                 collections=collections,
@@ -364,6 +390,7 @@ class HexApi:
                 description=hex_item.description,
                 created_at=hex_item.created_at,
                 last_edited_at=hex_item.last_edited_at,
+                last_published_at=hex_item.last_published_at,
                 status=status,
                 categories=categories,
                 collections=collections,

datahub/ingestion/source/hex/hex.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, Iterable, List, Optional
@@ -22,6 +23,7 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.source import MetadataWorkUnitProcessor
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.hex.api import HexApi, HexApiReport
 from datahub.ingestion.source.hex.constants import (
     DATAHUB_API_PAGE_SIZE_DEFAULT,
@@ -44,7 +46,6 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
     StatefulIngestionSourceBase,
 )
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.sdk.main_client import DataHubClient
 
 
@@ -68,7 +69,7 @@ class HexSourceConfig(
     )
     include_components: bool = Field(
         default=True,
-        desciption="Include Hex Components in the ingestion",
+        description="Include Hex Components in the ingestion",
     )
     page_size: int = Field(
         default=HEX_API_PAGE_SIZE_DEFAULT,
@@ -121,7 +122,11 @@ class HexSourceConfig(
 
     @root_validator(pre=True)
     def validate_lineage_times(cls, data: Dict[str, Any]) -> Dict[str, Any]:
-        # lineage_end_time default = now
+        # In-place update of the input dict would cause state contamination. This was discovered through test failures
+        # in test_hex.py where the same dict is reused.
+        # So a deepcopy is performed first.
+        data = deepcopy(data)
+
         if "lineage_end_time" not in data or data["lineage_end_time"] is None:
             data["lineage_end_time"] = datetime.now(tz=timezone.utc)
         # if string is given, parse it
@@ -166,7 +171,6 @@ class HexSourceConfig(
 class HexReport(
     StaleEntityRemovalSourceReport,
     HexApiReport,
-    IngestionStageReport,
     HexQueryFetcherReport,
 ):
     pass
@@ -174,11 +178,18 @@ class HexReport(
 
 @platform_name("Hex")
 @config_class(HexSourceConfig)
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.INCUBATING)
 @capability(SourceCapability.DESCRIPTIONS, "Supported by default")
 @capability(SourceCapability.OWNERSHIP, "Supported by default")
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.USAGE_STATS,
+    "Supported by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.HEX_PROJECT,
+    ],
+)
 class HexSource(StatefulIngestionSourceBase):
     def __init__(self, config: HexSourceConfig, ctx: PipelineContext):
         super().__init__(config, ctx)

datahub/ingestion/source/hex/mapper.py

@@ -122,7 +122,7 @@ class Mapper:
             lastModified=self._change_audit_stamps(
                 created_at=project.created_at, last_edited_at=project.last_edited_at
             ),
-            externalUrl=f"{self._base_url}/{self._workspace_name}/hex/{project.id}",
+            externalUrl=self._get_project_or_component_external_url(project),
             customProperties=dict(id=project.id),
             datasetEdges=self._dataset_edges(project.upstream_datasets),
             # TODO: support schema field upstream, maybe InputFields?
@@ -173,7 +173,7 @@ class Mapper:
             lastModified=self._change_audit_stamps(
                 created_at=component.created_at, last_edited_at=component.last_edited_at
             ),
-            externalUrl=f"{self._base_url}/{self._workspace_name}/hex/{component.id}",
+            externalUrl=self._get_project_or_component_external_url(component),
             customProperties=dict(id=component.id),
         )
 
@@ -242,6 +242,20 @@ class Mapper:
         assert isinstance(dashboard_urn, DashboardUrn)
         return dashboard_urn
 
+    def _get_project_or_component_external_url(
+        self,
+        project_or_component: Union[Project, Component],
+    ) -> Optional[str]:
+        if project_or_component.last_published_at is None:
+            return (
+                f"{self._base_url}/{self._workspace_name}/hex/{project_or_component.id}"
+            )
+        else:
+            # published Projects/Components have a different URL that everybody, not just editors, can access
+            return (
+                f"{self._base_url}/{self._workspace_name}/app/{project_or_component.id}"
+            )
+
     def _change_audit_stamps(
         self, created_at: Optional[datetime], last_edited_at: Optional[datetime]
     ) -> ChangeAuditStampsClass:

datahub/ingestion/source/hex/model.py

@@ -46,6 +46,7 @@ class Project:
     title: str
     description: Optional[str]
     last_edited_at: Optional[datetime] = None
+    last_published_at: Optional[datetime] = None
     created_at: Optional[datetime] = None
     status: Optional[Status] = None
     categories: Optional[List[Category]] = None  # TODO: emit category description!
@@ -67,6 +68,7 @@ class Component:
     title: str
     description: Optional[str]
     last_edited_at: Optional[datetime] = None
+    last_published_at: Optional[datetime] = None
     created_at: Optional[datetime] = None
     status: Optional[Status] = None
     categories: Optional[List[Category]] = None

datahub/ingestion/source/hex/query_fetcher.py

@@ -97,7 +97,7 @@ class HexQueryFetcher:
             if not query_urns or not entities_by_urn:
                 self.report.warning(
                     title="No Queries found with Hex as origin",
-                    message="No lineage because of no Queries found with Hex as origin in the given time range; you may consider extending the time range to fetch more queries.",
+                    message="No lineage because of no Queries found with Hex as origin in the given time range. You may need to set use_queries_v2: true on your warehouse ingestion or you may consider extending the time range to fetch more queries.",
                     context=str(
                         dict(
                             workspace_name=self.workspace_name,