acryl-datahub 1.0.0rc18__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,
+}