acryl-datahub 1.2.0.2rc1__py3-none-any.whl → 1.2.0.2rc3__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,120 @@
+"""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/
+"""
+
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from datahub.emitter.mcp_builder import ContainerKey
+
+# 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 DatasourceRef(BaseModel):
+    """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(BaseModel):
+    """Represents a Grafana dashboard panel."""
+
+    id: str
+    title: str
+    description: str = ""
+    type: Optional[str]
+    # 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(BaseModel):
+    """Represents a Grafana dashboard."""
+
+    uid: str
+    title: str
+    description: str = ""
+    version: Optional[str]
+    panels: List[Panel]
+    tags: List[str]
+    timezone: Optional[str]
+    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", {})
+        panels = cls.extract_panels(dashboard_data.get("panels", []))
+
+        # 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"]
+
+        return super().parse_obj(dashboard_dict)
+
+
+class Folder(BaseModel):
+    """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,91 @@
+from dataclasses import dataclass
+
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalSourceReport,
+)
+from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
+
+
+@dataclass
+class GrafanaSourceReport(StaleEntityRemovalSourceReport, IngestionStageReport):
+    # 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/hex.py

@@ -22,6 +22,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,
@@ -179,6 +180,13 @@ class HexReport(
 @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/looker/looker_source.py

@@ -51,6 +51,7 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import (
     BIAssetSubTypes,
     BIContainerSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.looker import looker_usage
 from datahub.ingestion.source.looker.looker_common import (
@@ -127,6 +128,14 @@ logger = logging.getLogger(__name__)
     "Enabled by default, configured using `extract_usage_history`",
 )
 @capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.LOOKML_MODEL,
+        SourceCapabilityModifier.LOOKER_FOLDER,
+    ],
+)
 class LookerDashboardSource(TestableSource, StatefulIngestionSourceBase):
     """
     This plugin extracts the following:

datahub/ingestion/source/looker/lookml_source.py

@@ -27,6 +27,7 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import (
     BIContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.git.git_import import GitClone
 from datahub.ingestion.source.looker.looker_common import (
@@ -273,6 +274,13 @@ class LookerManifest:
     SourceCapability.LINEAGE_FINE,
     "Enabled by default, configured using `extract_column_level_lineage`",
 )
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.LOOKML_PROJECT,
+    ],
+)
 class LookMLSource(StatefulIngestionSourceBase):
     """
     This plugin extracts the following:

datahub/ingestion/source/mongodb.py

@@ -36,7 +36,10 @@ 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 DatasetContainerSubTypes
+from datahub.ingestion.source.common.subtypes import (
+    DatasetContainerSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.schema_inference.object import (
     SchemaDescription,
     construct_schema,
@@ -249,6 +252,13 @@ def construct_schema_pymongo(
 @support_status(SupportStatus.CERTIFIED)
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.DATABASE,
+    ],
+)
 @dataclass
 class MongoDBSource(StatefulIngestionSourceBase):
     """

datahub/ingestion/source/redshift/redshift.py

@@ -47,6 +47,7 @@ from datahub.ingestion.source.common.data_reader import DataReader
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.redshift.config import RedshiftConfig
 from datahub.ingestion.source.redshift.datashares import RedshiftDatasharesHelper
@@ -126,7 +127,13 @@ logger: logging.Logger = logging.getLogger(__name__)
 @platform_name("Redshift")
 @config_class(RedshiftConfig)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.DATABASE,
+    ],
+)
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")

datahub/ingestion/source/s3/source.py

@@ -41,6 +41,7 @@ from datahub.ingestion.source.aws.s3_util import (
     get_key_prefix,
     strip_s3_prefix,
 )
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.data_lake_common.data_lake_utils import (
     ContainerWUCreator,
     add_partition_columns_to_schema,
@@ -196,7 +197,14 @@ class TableData:
 @platform_name("S3 / Local Files", id="s3")
 @config_class(DataLakeSourceConfig)
 @support_status(SupportStatus.INCUBATING)
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.FOLDER,
+        SourceCapabilityModifier.S3_BUCKET,
+    ],
+)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(
     SourceCapability.SCHEMA_METADATA, "Can infer schema from supported file types"

datahub/ingestion/source/sql/athena.py

@@ -347,12 +347,18 @@ class Partitionitem:
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Supported for S3 tables",
-    subtype_modifier=[SourceCapabilityModifier.TABLE],
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.TABLE,
+    ],
 )
 @capability(
     SourceCapability.LINEAGE_FINE,
     "Supported for S3 tables",
-    subtype_modifier=[SourceCapabilityModifier.TABLE],
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.TABLE,
+    ],
 )
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 class AthenaSource(SQLAlchemySource):

datahub/ingestion/source/sql/clickhouse.py

@@ -32,6 +32,7 @@ from datahub.ingestion.api.decorators import (
     support_status,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.sql.sql_common import (
     SqlWorkUnit,
     logger,
@@ -383,6 +384,14 @@ clickhouse_datetime_format = "%Y-%m-%d %H:%M:%S"
     SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Enabled by default to get lineage for views via `include_view_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.TABLE,
+    ],
+)
 class ClickHouseSource(TwoTierSQLAlchemySource):
     """
     This plugin extracts the following:

datahub/ingestion/source/sql_queries.py

@@ -73,9 +73,8 @@ class SqlQueriesSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
         default=None,
     )
     override_dialect: Optional[str] = Field(
-        description="DEPRECATED: This field is ignored. SQL dialect detection is now handled automatically by the SQL parsing aggregator based on the platform.",
+        description="The SQL dialect to use when parsing queries. Overrides automatic dialect detection.",
         default=None,
-        hidden_from_docs=True,
     )
 
 
@@ -230,6 +229,7 @@ class SqlQueriesSource(Source):
                     session_id=query_entry.session_id,
                     default_db=self.config.default_db,
                     default_schema=self.config.default_schema,
+                    override_dialect=self.config.override_dialect,
                 )
                 self.aggregator.add_observed_query(observed_query)