ss-tools/backend/src/scripts/test_dataset_dashboard_relations.py

#!/usr/bin/env python3
# #region test_dataset_dashboard_relations_script [C:2] [TYPE Module] [SEMANTICS pydantic, dashboard, dataset, test, superset]
# @defgroup Module Module group.
# @BRIEF Tests and inspects dataset-to-dashboard relationship responses from Superset API.
# @RATIONALE Added '# DIAGNOSTIC SCRIPT — not a test' header comment. Script already had if __name__ guard.

"""
Script to test dataset-to-dashboard relationships from Superset API.

Usage:
    cd backend && .venv/bin/python3 src/scripts/test_dataset_dashboard_relations.py
"""

import json
from pathlib import Path
import sys

# Add src to path (parent of scripts directory)
sys.path.append(str(Path(__file__).parent.parent.parent))

from src.core.config_manager import ConfigManager
from src.core.cot_logger import seed_trace_id
from src.core.logger import logger
from src.core.superset_client import SupersetClient


def test_dashboard_dataset_relations():
    """Test fetching dataset-to-dashboard relationships."""
    seed_trace_id()

    # Load environment from existing config
    config_manager = ConfigManager()
    environments = config_manager.get_environments()

    if not environments:
        logger.explore("No environments configured", payload={}, error="ConfigManager returned empty environment list")
        return

    # Use first available environment
    env = environments[0]
    logger.reason("Using environment for test", payload={"env_name": env.name, "env_url": env.url})

    client = SupersetClient(env)

    try:
        # Authenticate
        logger.reason("Authenticating to Superset")
        client.authenticate()
        logger.reason("Authentication successful", payload={"env_name": env.name})

        # Test dashboard ID 13
        dashboard_id = 13
        logger.reason("Fetching dashboard", payload={"dashboard_id": dashboard_id})
        dashboard = client.network.request(method="GET", endpoint=f"/dashboard/{dashboard_id}")

        print("\nDashboard structure:")
        print(f"  ID: {dashboard.get('id')}")
        print(f"  Title: {dashboard.get('dashboard_title')}")
        print(f"  Published: {dashboard.get('published')}")

        # Check for slices/charts
        if 'slices' in dashboard:
            logger.reason("Found slices in dashboard", payload={"dashboard_id": dashboard_id, "slice_count": len(dashboard['slices'])})
            for i, slice_data in enumerate(dashboard['slices'][:5]):  # Show first 5
                print(f"    Slice {i+1}:")
                print(f"      ID: {slice_data.get('slice_id')}")
                print(f"      Name: {slice_data.get('slice_name')}")
                # Check for datasource_id
                if 'datasource_id' in slice_data:
                    print(f"      Datasource ID: {slice_data['datasource_id']}")
                if 'datasource_name' in slice_data:
                    print(f"      Datasource Name: {slice_data['datasource_name']}")
                if 'datasource_type' in slice_data:
                    print(f"      Datasource Type: {slice_data['datasource_type']}")
        else:
            logger.reflect("No slices field in dashboard response", payload={"dashboard_id": dashboard_id})
            logger.reason("Dashboard available fields", payload={"fields": list(dashboard.keys())})

        # Test dataset ID 26
        dataset_id = 26
        logger.reason("Fetching dataset", payload={"dataset_id": dataset_id})
        dataset = client.get_dataset(dataset_id)

        print("\nDataset structure:")
        print(f"  ID: {dataset.get('id')}")
        print(f"  Table Name: {dataset.get('table_name')}")
        print(f"  Schema: {dataset.get('schema')}")
        print(f"  Database: {dataset.get('database', {}).get('database_name', 'Unknown')}")

        # Check for dashboards that use this dataset
        logger.reason("Finding dashboards for dataset", payload={"dataset_id": dataset_id})

        # Method: Use Superset's related_objects API
        try:
            logger.reason("Querying related objects endpoint", payload={"dataset_id": dataset_id})
            related_objects = client.network.request(
                method="GET",
                endpoint=f"/dataset/{dataset_id}/related_objects"
            )

            logger.reason("Related objects response type", payload={"response_type": str(type(related_objects))})
            logger.reason("Related objects keys", payload={"keys": list(related_objects.keys()) if isinstance(related_objects, dict) else "N/A"})

            # Check for dashboards in related objects
            if 'dashboards' in related_objects:
                dashboards = related_objects['dashboards']
                logger.reason("Found dashboards for dataset", payload={"dataset_id": dataset_id, "count": len(dashboards)})

                for dash in dashboards:
                    if isinstance(dash, dict):
                        logger.reason("Dashboard entry", payload={"dashboard_id": dash.get("id"), "title": dash.get("dashboard_title", dash.get("title", "Unknown"))})
                    else:
                        logger.reason("Dashboard entry", payload={"dashboard": str(dash)})
            elif 'result' in related_objects:
                # Some Superset versions use 'result' wrapper
                result = related_objects['result']
                if 'dashboards' in result:
                    dashboards = result['dashboards']
                    logger.reason("Found dashboards for dataset (result wrapper)", payload={"dataset_id": dataset_id, "count": len(dashboards)})

                    for dash in dashboards:
                        logger.reason("Dashboard entry", payload={"dashboard_id": dash.get("id"), "title": dash.get("dashboard_title", dash.get("title", "Unknown"))})
                else:
                    logger.reflect("No dashboards key in result", payload={"keys": list(result.keys())})
            else:
                logger.reflect("No dashboards key in response", payload={"available_keys": list(related_objects.keys())})
                logger.reason("Full related_objects response")
                print(json.dumps(related_objects, indent=2, default=str)[:1000])

        except Exception as e:
            logger.explore("Error fetching related objects", payload={"dataset_id": dataset_id}, error=str(e))
            import traceback
            traceback.print_exc()

        # Method 2: Try to use the position_json from dashboard
        logger.reason("Analyzing dashboard position JSON")
        if 'position_json' in dashboard:
            position_data = json.loads(dashboard['position_json'])
            logger.reason("Position data type", payload={"type": str(type(position_data))})

            # Look for datasource references
            datasource_ids = set()
            if isinstance(position_data, dict):
                for key, value in position_data.items():
                    if 'datasource' in key.lower() or key == 'DASHBOARD_VERSION_KEY':
                        logger.reason("Position entry", payload={"key": key, "value_type": str(type(value))})
            elif isinstance(position_data, list):
                logger.reason("Position data item count", payload={"count": len(position_data)})
                for item in position_data[:3]:  # Show first 3
                    logger.reason("Position item", payload={"item_type": str(type(item)), "keys": list(item.keys()) if isinstance(item, dict) else "N/A"})
                    if isinstance(item, dict):
                        if 'datasource_id' in item:
                            datasource_ids.add(item['datasource_id'])

            if datasource_ids:
                logger.reason("Found datasource IDs in position data", payload={"datasource_ids": sorted(datasource_ids)})

        # Save full response for analysis
        output_file = Path(__file__).parent / "dataset_dashboard_analysis.json"
        with open(output_file, 'w') as f:
            json.dump({
                'dashboard': dashboard,
                'dataset': dataset
            }, f, indent=2, default=str)
        logger.reason("Full response saved", payload={"output_file": str(output_file)})

    except Exception as e:
        logger.explore("Error in test dashboard dataset relations", payload={}, error=str(e))
        raise


if __name__ == "__main__":
    test_dashboard_dataset_relations()

# #endregion test_dataset_dashboard_relations_script
# #endregion test_dataset_dashboard_relations_script