acryl-datahub 1.2.0.1rc1__py3-none-any.whl → 1.2.0.2__py3-none-any.whl

datahub/ingestion/source/grafana/field_utils.py

@@ -0,0 +1,307 @@
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.source.grafana.models import Panel
+from datahub.metadata.schema_classes import (
+    NumberTypeClass,
+    SchemaFieldClass,
+    SchemaFieldDataTypeClass,
+    StringTypeClass,
+    TimeTypeClass,
+)
+from datahub.sql_parsing.sqlglot_lineage import (
+    create_lineage_sql_parsed_result,
+    infer_output_schema,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def extract_sql_column_fields(target: Dict[str, Any]) -> List[SchemaFieldClass]:
+    """Extract fields from SQL-style columns."""
+    fields = []
+    for col in target.get("sql", {}).get("columns", []):
+        for param in col.get("parameters", []):
+            if param.get("type") == "column" and param.get("name"):
+                field_type: Union[NumberTypeClass, StringTypeClass, TimeTypeClass] = (
+                    TimeTypeClass()
+                    if col["type"] == "time"
+                    else NumberTypeClass()
+                    if col["type"] == "number"
+                    else StringTypeClass()
+                )
+                fields.append(
+                    SchemaFieldClass(
+                        fieldPath=param["name"],
+                        type=SchemaFieldDataTypeClass(type=field_type),
+                        nativeDataType=col["type"],
+                    )
+                )
+    return fields
+
+
+def extract_prometheus_fields(target: Dict[str, Any]) -> List[SchemaFieldClass]:
+    """Extract fields from Prometheus expressions."""
+    expr = target.get("expr")
+    if expr:
+        legend = target.get("legendFormat", expr)
+        return [
+            SchemaFieldClass(
+                fieldPath=legend,
+                type=SchemaFieldDataTypeClass(type=NumberTypeClass()),
+                nativeDataType="prometheus_metric",
+            )
+        ]
+    return []
+
+
+def extract_raw_sql_fields(
+    target: Dict[str, Any],
+    panel: Optional[Panel] = None,
+    connection_to_platform_map: Optional[Dict[str, Any]] = None,
+    graph: Optional[DataHubGraph] = None,
+    report: Optional[Any] = None,
+) -> List[SchemaFieldClass]:
+    """Extract fields from raw SQL queries using DataHub's SQL parsing."""
+    raw_sql = target.get("rawSql", "")
+    if not raw_sql:
+        return []
+
+    # Determine upstream platform and environment from datasource mapping
+    platform = "unknown"
+    env = "PROD"
+    default_db = None
+    default_schema = None
+    platform_instance = None
+    schema_aware = False
+
+    if panel and panel.datasource_ref and connection_to_platform_map:
+        ds_type = panel.datasource_ref.type or "unknown"
+        ds_uid = panel.datasource_ref.uid or "unknown"
+
+        # Try to find mapping by datasource UID first, then by type
+        platform_config = connection_to_platform_map.get(
+            ds_uid
+        ) or connection_to_platform_map.get(ds_type)
+
+        if platform_config:
+            platform = platform_config.platform
+            env = getattr(platform_config, "env", env)
+            default_db = getattr(platform_config, "database", None)
+            default_schema = getattr(platform_config, "database_schema", None)
+            platform_instance = getattr(platform_config, "platform_instance", None)
+
+            # Enable schema-aware parsing if we have platform mapping and graph access
+            if graph and platform != "unknown":
+                schema_aware = True
+
+    # Track SQL parsing attempt
+    if report:
+        report.report_sql_parsing_attempt()
+
+    try:
+        # Use DataHub's standard SQL parsing approach
+        sql_parsing_result = create_lineage_sql_parsed_result(
+            query=raw_sql,
+            default_db=default_db,
+            default_schema=default_schema,
+            platform=platform,
+            platform_instance=platform_instance,
+            env=env,
+            schema_aware=schema_aware,
+            graph=graph,
+        )
+
+        # Extract the output schema from the parsing result
+        output_schema = infer_output_schema(sql_parsing_result)
+
+        if output_schema:
+            if report:
+                report.report_sql_parsing_success()
+            return output_schema
+        else:
+            # If sqlglot parsing succeeds but no schema is inferred,
+            # fall back to basic parsing
+            logger.debug(f"No schema inferred from SQL: {raw_sql}")
+            fallback_result = _extract_raw_sql_fields_fallback(target)
+            if fallback_result and report:
+                report.report_sql_parsing_success()
+            elif report:
+                report.report_sql_parsing_failure()
+            return fallback_result
+
+    except Exception as e:
+        logger.debug(f"Failed to parse SQL with DataHub parser: {raw_sql}, error: {e}")
+        if report:
+            report.report_sql_parsing_failure()
+        # Fallback to basic parsing for backwards compatibility
+        return _extract_raw_sql_fields_fallback(target)
+
+
+def _extract_raw_sql_fields_fallback(target: Dict[str, Any]) -> List[SchemaFieldClass]:
+    """Fallback basic SQL parsing for when sqlglot fails."""
+    raw_sql = target.get("rawSql", "").lower()
+    if not raw_sql:
+        return []
+
+    try:
+        sql = raw_sql.lower()
+        select_start = sql.index("select") + 6  # len("select")
+        from_start = sql.index("from")
+        select_part = sql[select_start:from_start].strip()
+
+        # Split by comma, handling nested parentheses
+        columns = []
+        current_column = ""
+        paren_count = 0
+
+        for char in select_part:
+            if char == "," and paren_count == 0:
+                if current_column.strip():
+                    columns.append(current_column.strip())
+                current_column = ""
+            else:
+                if char == "(":
+                    paren_count += 1
+                elif char == ")":
+                    paren_count -= 1
+                current_column += char
+
+        if current_column.strip():
+            columns.append(current_column.strip())
+
+        # For each column, extract the alias if it exists
+        fields = []
+        for col in columns:
+            # Check for alias with 'AS' keyword
+            if " as " in col:
+                field_name = col.split(" as ")[-1].strip()
+            else:
+                # If no alias, use the last part after last space
+                # This handles both simple columns and function calls without alias
+                field_name = col.split()[-1].strip()
+
+            # Clean up any remaining quotes or parentheses
+            field_name = field_name.strip("\"'()")
+
+            fields.append(
+                SchemaFieldClass(
+                    fieldPath=field_name,
+                    type=SchemaFieldDataTypeClass(type=StringTypeClass()),
+                    nativeDataType="sql_column",
+                )
+            )
+
+        return fields
+
+    except (IndexError, ValueError, StopIteration) as e:
+        logger.warning(f"Failed to parse SQL: {target.get('rawSql')}", e)
+        return []
+
+
+def extract_fields_from_panel(
+    panel: Panel,
+    connection_to_platform_map: Optional[Dict[str, Any]] = None,
+    graph: Optional[DataHubGraph] = None,
+    report: Optional[Any] = None,
+) -> List[SchemaFieldClass]:
+    """Extract all fields from a panel."""
+    fields = []
+    fields.extend(
+        extract_fields_from_targets(
+            panel.query_targets, panel, connection_to_platform_map, graph, report
+        )
+    )
+    fields.extend(get_fields_from_field_config(panel.field_config))
+    fields.extend(get_fields_from_transformations(panel.transformations))
+
+    # Track schema field extraction
+    if report:
+        if fields:
+            report.report_schema_fields_extracted()
+        else:
+            report.report_no_schema_fields()
+
+    return fields
+
+
+def extract_fields_from_targets(
+    targets: List[Dict[str, Any]],
+    panel: Optional[Panel] = None,
+    connection_to_platform_map: Optional[Dict[str, Any]] = None,
+    graph: Optional[DataHubGraph] = None,
+    report: Optional[Any] = None,
+) -> List[SchemaFieldClass]:
+    """Extract fields from panel targets."""
+    fields = []
+    for target in targets:
+        fields.extend(extract_sql_column_fields(target))
+        fields.extend(extract_prometheus_fields(target))
+        fields.extend(
+            extract_raw_sql_fields(
+                target, panel, connection_to_platform_map, graph, report
+            )
+        )
+        fields.extend(extract_time_format_fields(target))
+    return fields
+
+
+def extract_time_format_fields(target: Dict[str, Any]) -> List[SchemaFieldClass]:
+    """Extract fields from time series and table formats."""
+    if target.get("format") in {"time_series", "table"}:
+        return [
+            SchemaFieldClass(
+                fieldPath="time",
+                type=SchemaFieldDataTypeClass(type=TimeTypeClass()),
+                nativeDataType="timestamp",
+            )
+        ]
+    return []
+
+
+def get_fields_from_field_config(
+    field_config: Dict[str, Any],
+) -> List[SchemaFieldClass]:
+    """Extract fields from field configuration."""
+    fields = []
+    defaults = field_config.get("defaults", {})
+    unit = defaults.get("unit")
+    if unit:
+        fields.append(
+            SchemaFieldClass(
+                fieldPath=f"value_{unit}",
+                type=SchemaFieldDataTypeClass(type=NumberTypeClass()),
+                nativeDataType="value",
+            )
+        )
+    for override in field_config.get("overrides", []):
+        if override.get("matcher", {}).get("id") == "byName":
+            field_name = override.get("matcher", {}).get("options")
+            if field_name:
+                fields.append(
+                    SchemaFieldClass(
+                        fieldPath=field_name,
+                        type=SchemaFieldDataTypeClass(type=NumberTypeClass()),
+                        nativeDataType="metric",
+                    )
+                )
+    return fields
+
+
+def get_fields_from_transformations(
+    transformations: List[Dict[str, Any]],
+) -> List[SchemaFieldClass]:
+    """Extract fields from transformations."""
+    fields = []
+    for transform in transformations:
+        if transform.get("type") == "organize":
+            for field_name in transform.get("options", {}).get("indexByName", {}):
+                fields.append(
+                    SchemaFieldClass(
+                        fieldPath=field_name,
+                        type=SchemaFieldDataTypeClass(type=StringTypeClass()),
+                        nativeDataType="transformed",
+                    )
+                )
+    return fields

datahub/ingestion/source/grafana/grafana_api.py

@@ -0,0 +1,142 @@
+"""API client for Grafana metadata extraction"""
+
+import logging
+from typing import Dict, List, Optional, Union
+
+import requests
+import urllib3.exceptions
+from pydantic import SecretStr
+
+from datahub.ingestion.source.grafana.models import Dashboard, Folder
+from datahub.ingestion.source.grafana.report import GrafanaSourceReport
+
+logger = logging.getLogger(__name__)
+
+
+class GrafanaAPIClient:
+    """Client for making requests to Grafana API"""
+
+    def __init__(
+        self,
+        base_url: str,
+        token: SecretStr,
+        verify_ssl: bool,
+        page_size: int,
+        report: GrafanaSourceReport,
+    ) -> None:
+        self.base_url = base_url
+        self.verify_ssl = verify_ssl
+        self.page_size = page_size
+        self.report = report
+        self.session = self._create_session(token)
+
+    def _create_session(self, token: SecretStr) -> requests.Session:
+        session = requests.Session()
+        session.headers.update(
+            {
+                "Authorization": f"Bearer {token.get_secret_value()}",
+                "Accept": "application/json",
+                "Content-Type": "application/json",
+            }
+        )
+        session.verify = self.verify_ssl
+
+        # If SSL verification is disabled, suppress the warnings
+        if not self.verify_ssl:
+            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+            self.report.warning(
+                title="SSL Configuration Warning",
+                message="SSL Verification is recommended.",
+            )
+
+        return session
+
+    def get_folders(self) -> List[Folder]:
+        """Fetch all folders from Grafana with pagination."""
+        folders: List[Folder] = []
+        page = 1
+        per_page = self.page_size
+
+        while True:
+            try:
+                response = self.session.get(
+                    f"{self.base_url}/api/folders",
+                    params={"page": page, "limit": per_page},
+                )
+                response.raise_for_status()
+
+                batch = response.json()
+                if not batch:
+                    break
+
+                folders.extend(Folder.parse_obj(folder) for folder in batch)
+                page += 1
+            except requests.exceptions.RequestException as e:
+                self.report.report_failure(
+                    title="Folder Fetch Error",
+                    message="Failed to fetch folders on page",
+                    context=str(page),
+                    exc=e,
+                )
+                self.report.report_permission_warning()  # Likely a permission issue
+                break
+
+        return folders
+
+    def get_dashboard(self, uid: str) -> Optional[Dashboard]:
+        """Fetch a specific dashboard by UID"""
+        try:
+            response = self.session.get(f"{self.base_url}/api/dashboards/uid/{uid}")
+            response.raise_for_status()
+            return Dashboard.parse_obj(response.json())
+        except requests.exceptions.RequestException as e:
+            self.report.warning(
+                title="Dashboard Fetch Error",
+                message="Failed to fetch dashboard",
+                context=uid,
+                exc=e,
+            )
+            if e.response and e.response.status_code in (401, 403):
+                self.report.report_permission_warning()
+            return None
+
+    def get_dashboards(self) -> List[Dashboard]:
+        """Fetch all dashboards from search endpoint with pagination."""
+        dashboards: List[Dashboard] = []
+        page = 1
+        per_page = self.page_size
+
+        while True:
+            try:
+                params: Dict[str, Union[str, int]] = {
+                    "type": "dash-db",
+                    "page": page,
+                    "limit": per_page,
+                }
+                response = self.session.get(
+                    f"{self.base_url}/api/search",
+                    params=params,
+                )
+                response.raise_for_status()
+
+                batch = response.json()
+                if not batch:
+                    break
+
+                for result in batch:
+                    dashboard = self.get_dashboard(result["uid"])
+                    if dashboard:
+                        dashboards.append(dashboard)
+                page += 1
+            except requests.exceptions.RequestException as e:
+                self.report.report_failure(
+                    title="Dashboard Search Error",
+                    message="Failed to fetch dashboards on page",
+                    context=str(page),
+                    exc=e,
+                )
+                if e.response and e.response.status_code in (401, 403):
+                    self.report.report_permission_warning()
+                break
+
+        return dashboards

datahub/ingestion/source/grafana/grafana_config.py

@@ -0,0 +1,104 @@
+from typing import Dict, Optional
+
+from pydantic import Field, SecretStr, validator
+
+from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.source_common import (
+    DatasetLineageProviderConfigBase,
+    EnvConfigMixin,
+    PlatformInstanceConfigMixin,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionConfigBase,
+)
+from datahub.utilities import config_clean
+
+
+class PlatformConnectionConfig(
+    EnvConfigMixin,
+    PlatformInstanceConfigMixin,
+):
+    """Platform connection configuration for mapping Grafana datasources to their actual platforms."""
+
+    platform: str = Field(
+        description="The platform name (e.g., 'postgres', 'mysql', 'snowflake')"
+    )
+    database: Optional[str] = Field(default=None, description="Default database name")
+    database_schema: Optional[str] = Field(
+        default=None, description="Default schema name"
+    )
+
+
+class GrafanaSourceConfig(
+    DatasetLineageProviderConfigBase,
+    StatefulIngestionConfigBase,
+    PlatformInstanceConfigMixin,
+    EnvConfigMixin,
+):
+    """Configuration for Grafana source"""
+
+    platform: str = Field(default="grafana", hidden_from_docs=True)
+    url: str = Field(
+        description="Grafana URL in the format http://your-grafana-instance with no trailing slash"
+    )
+    service_account_token: SecretStr = Field(
+        description="Service account token for Grafana"
+    )
+    verify_ssl: bool = Field(
+        default=True,
+        description="Whether to verify SSL certificates when connecting to Grafana",
+    )
+
+    # API pagination configuration
+    page_size: int = Field(
+        default=100,
+        description="Number of items to fetch per API call when paginating through folders and dashboards",
+    )
+
+    # Extraction mode configuration
+    basic_mode: bool = Field(
+        default=False,
+        description="Enable basic extraction mode for users with limited permissions. "
+        "In basic mode, only dashboard metadata is extracted without detailed panel information, "
+        "lineage, or folder hierarchy. This requires only basic dashboard read permissions.",
+    )
+
+    # Content filtering
+    dashboard_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex pattern to filter dashboards for ingestion",
+    )
+    folder_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex pattern to filter folders for ingestion",
+    )
+
+    # Feature toggles
+    ingest_tags: bool = Field(
+        default=True, description="Whether to ingest dashboard and chart tags"
+    )
+    ingest_owners: bool = Field(
+        default=True, description="Whether to ingest dashboard ownership information"
+    )
+
+    include_lineage: bool = Field(
+        default=True,
+        description="Whether to extract lineage between charts and data sources. "
+        "When enabled, the source will parse SQL queries and datasource configurations "
+        "to build lineage relationships.",
+    )
+    include_column_lineage: bool = Field(
+        default=True,
+        description="Whether to extract column-level lineage from SQL queries. "
+        "Only applicable when include_lineage is enabled.",
+    )
+
+    # Platform connection mappings
+    connection_to_platform_map: Dict[str, PlatformConnectionConfig] = Field(
+        default_factory=dict,
+        description="Map of Grafana datasource types/UIDs to platform connection configs for lineage extraction",
+    )
+
+    @validator("url", allow_reuse=True)
+    def remove_trailing_slash(cls, v):
+        return config_clean.remove_trailing_slashes(v)