content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  
  from typing import (
      Any, Dict, Iterator, Tuple, Union,
  )
  
  from databuilder.extractor.dashboard.apache_superset.apache_superset_extractor import (
      ApacheSupersetBaseExtractor, type_fields_mapping,
  )
  from databuilder.models.dashboard.dashboard_chart import DashboardChart
  from databuilder.models.dashboard.dashboard_query import DashboardQuery
  
  
  class ApacheSupersetChartExtractor(ApacheSupersetBaseExtractor):
      def chart_field_mappings(self) -> type_fields_mapping:
          result = [
              ('chart_id', 'id', lambda x: str(x), ''),
              ('chart_name', 'slice_name', None, ''),
              ('chart_type', 'viz_type', None, ''),
              ('chart_url', 'url', None, ''),  # currently not available in superset chart api
          ]
  
          return result
  
      def _get_extract_iter(self) -> Iterator[Union[DashboardQuery, DashboardChart, None]]:
          ids = self._get_resource_ids('dashboard')
  
          data = [self._get_dashboard_details(i) for i in ids]
  
          for entry in data:
              dashboard_id, dashboard_data = entry
  
              # Since Apache Superset doesn't support dashboard <> query <> chart relation we create a dummy 'bridge'
              # query node so we can connect charts to a dashboard
              dashboard_query_data = dict(dashboard_id=dashboard_id,
                                          query_name='default',
                                          query_id=self.dummy_query_id,
                                          url='',
                                          query_text='')
              dashboard_query_data.update(**self.common_params)
  
              yield DashboardQuery(**dashboard_query_data)
  
              charts = [s.get('__Slice__') for s in dashboard_data.get('slices', [])]
  
              for chart in charts:
                  if chart:
                      dashboard_chart_data = self.map_fields(chart, self.chart_field_mappings())
                      dashboard_chart_data.update(**{**dict(dashboard_id=dashboard_id, query_id=self.dummy_query_id),
                                                     **self.common_params})
  
                      yield DashboardChart(**dashboard_chart_data)
  
      def _get_dashboard_details(self, dashboard_id: str) -> Tuple[str, Dict[str, Any]]:
          url = self.build_full_url(f'api/v1/dashboard/export?q=[{dashboard_id}]')
  
          _data = self.execute_query(url)
  
          dashboard_data = _data.get('dashboards', [dict()])[0].get('__Dashboard__', dict())
  
          data = dashboard_id, dashboard_data
  
          return data
  
      @property
      def dummy_query_id(self) -> str:
          return '-1'
  