fix(superset): API datamodel fields + SQL lineage, cache datasource/database fetches

ingestion/src/metadata/ingestion/source/dashboard/superset/api_source.py

@@ -31,13 +31,15 @@
     FullyQualifiedEntityName,
     Markdown,
     SourceUrl,
+    SqlQuery,
 )
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.source.dashboard.superset.mixin import SupersetSourceMixin
 from metadata.ingestion.source.dashboard.superset.models import (
     ChartResult,
     DashboardResult,
     FetchChart,
+    SupersetDatasource,
 )
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_datamodel
@@ -124,11 +126,54 @@ def yield_dashboard(self, dashboard_details: DashboardResult) -> Iterable[Either
                 )
             )
 
-    def _get_datasource_fqn_for_lineage(self, chart_json: ChartResult, db_service_prefix: Optional[str]):  # noqa: UP045
+    def _get_datasource_fqn_for_lineage(self, chart_json, db_service_prefix: str | None):
+        # A SQL-parsed source table carries a table_name; a raw chart carries only a datasource_id.
+        if getattr(chart_json, "table_name", None):
+            return self._get_source_table_fqn(chart_json, db_service_prefix)
         return (
             self._get_datasource_fqn(chart_json.datasource_id, db_service_prefix) if chart_json.datasource_id else None
         )
 
+    def _get_input_tables(self, chart: ChartResult):
+        # Virtual (SQL) datasets: parse the dataset SQL to reach real source tables.
+        datasource_id = getattr(chart, "datasource_id", None)
+        datasource = self.client.fetch_datasource(datasource_id) if datasource_id else None
+        dataset_sql = datasource.result.sql if datasource and datasource.result else None
+        if dataset_sql:
+            enriched = FetchChart(
+                sql=dataset_sql,
+                schema=datasource.result.table_schema,
+                datasource_id=datasource_id,
+            )
+            result = self._parse_lineage_from_dataset_sql(enriched)
+        else:
+            result = super()._get_input_tables(chart)
+        return result
+
+    def _resolve_lineage_database_name(self, datasource_json: SupersetDatasource, db_service_name: str) -> str | None:
+        database_json = self.client.fetch_database(datasource_json.result.database.id)
+        default_database_name = database_json.result.parameters.database if database_json.result.parameters else None
+        db_service_entity = self.metadata.get_by_name(entity=DatabaseService, fqn=db_service_name)
+        return get_database_name_for_lineage(db_service_entity, default_database_name)
+
+    def _get_source_table_fqn(self, chart_json: FetchChart, db_service_prefix: Optional[str]) -> Optional[str]:  # noqa: UP045
+        (
+            db_service_name,
+            prefix_database_name,
+            prefix_schema_name,
+            prefix_table_name,
+        ) = self.parse_db_service_prefix(db_service_prefix)
+        database_name = None
+        if db_service_name and chart_json.datasource_id:
+            datasource_json = self.client.fetch_datasource(chart_json.datasource_id)
+            database_name = self._resolve_lineage_database_name(datasource_json, db_service_name)
+        return build_es_fqn_search_string(
+            database_name=prefix_database_name or database_name,
+            schema_name=prefix_schema_name or chart_json.table_schema,
+            service_name=db_service_name or "*",
+            table_name=prefix_table_name or chart_json.table_name,
+        )
+
     def yield_dashboard_chart(self, dashboard_details: DashboardResult) -> Iterable[Either[CreateChartRequest]]:
         """Method to fetch charts linked to dashboard"""
         for chart_id in self._get_charts_of_dashboard(dashboard_details):
@@ -168,12 +213,7 @@ def _get_datasource_fqn(self, datasource_id: str, db_service_prefix: Optional[st
             if datasource_json:
                 database_name = None
                 if db_service_prefix:
-                    database_json = self.client.fetch_database(datasource_json.result.database.id)
-                    default_database_name = (
-                        database_json.result.parameters.database if database_json.result.parameters else None
-                    )
-                    db_service_entity = self.metadata.get_by_name(entity=DatabaseService, fqn=db_service_name)
-                    database_name = get_database_name_for_lineage(db_service_entity, default_database_name)
+                    database_name = self._resolve_lineage_database_name(datasource_json, db_service_name)
 
                     if prefix_database_name and database_name and prefix_database_name.lower() != database_name.lower():
                         logger.debug(f"Database {database_name} does not match prefix {prefix_database_name}")
@@ -228,11 +268,19 @@ def yield_datamodel(self, dashboard_details: DashboardResult) -> Iterable[Either
                             datasource_json.result.table_name,
                             "Data model filtered out.",
                         )
-                    data_model_request = CreateDashboardDataModelRequest(
+                    result = datasource_json.result
+                    data_model_request = CreateDashboardDataModelRequest(  # pyright: ignore[reportCallIssue]
                         name=EntityName(str(datasource_json.id)),
-                        displayName=datasource_json.result.table_name,
+                        displayName=result.table_name,
+                        description=Markdown(result.description) if result.description else None,
+                        sql=SqlQuery(result.sql) if result.sql else None,
+                        sourceUrl=(
+                            SourceUrl(f"{clean_uri(str(self.service_connection.hostPort))}{result.url}")
+                            if result.url
+                            else None
+                        ),
                         service=FullyQualifiedEntityName(self.context.get().dashboard_service),
-                        columns=self.get_column_info(datasource_json.result.columns),
+                        columns=self.get_column_info(result.columns),
                         dataModelType=DataModelType.SupersetDataModel.value,
                     )
                     yield Either(right=data_model_request)

ingestion/src/metadata/ingestion/source/dashboard/superset/client.py

@@ -14,6 +14,9 @@
 
 import json
 import traceback
+from operator import attrgetter
+
+from cachetools import LRUCache, cachedmethod
 
 from metadata.generated.schema.entity.services.connections.dashboard.supersetConnection import (
     SupersetConnection,
@@ -102,6 +105,8 @@ def __init__(self, config: SupersetConnection):
             verify=get_verify_ssl(config.connection.sslConfig),
         )
         self.client = TrackedREST(client_config, source_name="superset")
+        self._datasource_cache = LRUCache(maxsize=512)
+        self._database_cache = LRUCache(maxsize=512)
 
     def get_dashboard_count(self) -> int:
         resp_dashboards = self.client.get("/dashboard/?q=(page:0,page_size:1)")
@@ -214,6 +219,7 @@ def fetch_charts_with_id(self, chart_id: str):
         response = self.client.get(f"/chart/{chart_id}")
         return response  # noqa: RET504
 
+    @cachedmethod(attrgetter("_datasource_cache"))
     def fetch_datasource(self, datasource_id: str) -> SupersetDatasource:
         """
         Fetch data source
@@ -235,6 +241,7 @@ def fetch_datasource(self, datasource_id: str) -> SupersetDatasource:
 
         return SupersetDatasource()
 
+    @cachedmethod(attrgetter("_database_cache"))
     def fetch_database(self, database_id: str) -> ListDatabaseResult:
         """
         Fetch database

ingestion/src/metadata/ingestion/source/dashboard/superset/mixin.py

@@ -47,7 +47,7 @@
 from metadata.ingestion.source.dashboard.dashboard_service import DashboardServiceSource
 from metadata.ingestion.source.dashboard.superset.models import (
     DashboardResult,
-    DataSourceResult,
+    DSColumns,
     FetchChart,
     FetchColumn,
     FetchDashboard,
@@ -232,6 +232,7 @@ def _parse_lineage_from_dataset_sql(self, chart_json: FetchChart) -> List[Tuple[
                         table_name=table_name,
                         schema=table_schema,
                         sqlalchemy_uri=chart_json.sqlalchemy_uri,
+                        datasource_id=chart_json.datasource_id,
                     ),
                     column_mapping,
                 )
@@ -380,7 +381,7 @@ def parse_row_data_type(self, col_parse: dict) -> List[Column]:  # noqa: UP006
             return col_parse["children"]
         return []
 
-    def get_column_info(self, data_source: List[Union[DataSourceResult, FetchColumn]]) -> Optional[List[Column]]:  # noqa: UP006, UP007, UP045
+    def get_column_info(self, data_source: list[DSColumns | FetchColumn]) -> list[Column] | None:
         """
         Args:
             data_source: DataSource
@@ -400,7 +401,7 @@ def get_column_info(self, data_source: List[Union[DataSourceResult, FetchColumn]
                         dataType=col_parse["dataType"],
                         arrayDataType=self.parse_array_data_type(col_parse),
                         children=self.parse_row_data_type(col_parse),
-                        name=truncate_column_name(str(field.id)),
+                        name=truncate_column_name(field.column_name or str(field.id)),
                         displayName=field.column_name,
                         description=field.description,
                         dataLength=int(col_parse.get("dataLength", 0)),

ingestion/tests/integration/superset/test_superset.py

@@ -64,6 +64,7 @@
 from metadata.ingestion.source.dashboard.superset.db_source import SupersetDBSource
 from metadata.ingestion.source.dashboard.superset.metadata import SupersetSource
 from metadata.ingestion.source.dashboard.superset.models import (
+    ChartResult,
     DatabaseResult,
     DataSourceResult,
     FetchChart,
@@ -218,6 +219,20 @@
 )
 MOCK_DATABASE_RESPONSE = ListDatabaseResult(result=DatabaseResult(database_name="examples", id=1, parameters=None))
 
+MOCK_DATASOURCE_RESPONSE_WITH_SQL = SupersetDatasource(
+    id=99,
+    result=DataSourceResult.model_validate(
+        {
+            "table_name": "sample_table",
+            "sql": "SELECT id FROM sample_table",
+            "description": "rollup dataset",
+            "url": "/tablemodelview/edit/99",
+            "schema": "main",
+            "columns": [{"id": 11, "column_name": "Population", "type": "INT"}],
+        }
+    ),
+)
+
 
 def setup_sample_data(postgres_container):
     engine = sqlalchemy.create_engine(postgres_container.get_connection_url())
@@ -631,6 +646,62 @@ def test_broken_column_type_in_datamodel(self):
         self.superset_db.prepare()
         parsed_datasource = self.superset_db.get_column_info(MOCK_DATASOURCE)
         assert parsed_datasource[0].dataType.value == "INT"
+        # column name is the real column_name, not the numeric superset column id
+        assert parsed_datasource[0].name.root == "Population"
+
+    def test_datamodel_fields_api(self):
+        """
+        API datamodel carries sql, description and sourceUrl from the dataset payload
+        """
+        self.superset_api.all_charts = {69: MOCK_CHART}
+        with patch.object(
+            self.superset_api.client,
+            "fetch_datasource",
+            return_value=MOCK_DATASOURCE_RESPONSE_WITH_SQL,
+        ):
+            data_model = next(self.superset_api.yield_datamodel(MOCK_DASHBOARD)).right
+        assert data_model.sql.root == "SELECT id FROM sample_table"
+        assert data_model.description.root == "rollup dataset"
+        assert str(data_model.sourceUrl.root).endswith("/tablemodelview/edit/99")
+        assert data_model.columns[0].name.root == "Population"
+
+    def test_api_get_input_tables_parses_dataset_sql(self):
+        """
+        API _get_input_tables parses the virtual dataset SQL to reach the real source tables
+        """
+        with patch.object(
+            self.superset_api.client,
+            "fetch_datasource",
+            return_value=MOCK_DATASOURCE_RESPONSE_WITH_SQL,
+        ):
+            result = self.superset_api._get_input_tables(ChartResult(datasource_id=99))
+        source_tables = [fetch_chart.table_name for fetch_chart, _ in result]
+        self.assertIn("sample_table", source_tables)
+
+    def test_api_get_source_table_fqn_uses_parsed_table(self):
+        """
+        SQL-parsed source table fqn uses the parsed table name, not the datasource's own table
+        """
+        with (
+            patch.object(OpenMetadata, "get_by_name", return_value=MOCK_DB_POSTGRES_SERVICE),
+            patch.object(self.superset_api.client, "fetch_datasource", return_value=MOCK_DATASOURCE_RESPONSE),
+            patch.object(self.superset_api.client, "fetch_database", return_value=MOCK_DATABASE_RESPONSE),
+        ):
+            fqn = self.superset_api._get_source_table_fqn(  # pylint: disable=protected-access
+                FetchChart(table_name="orders", schema="main", datasource_id=1),
+                MOCK_DB_POSTGRES_SERVICE.name.root,
+            )
+        self.assertEqual(fqn, "test_postgres.*.main.orders")
+
+    def test_api_fetch_datasource_is_cached(self):
+        """
+        Repeated fetch_datasource for the same id resolves from cache, hitting the network once
+        """
+        with patch.object(self.superset_api.client.client, "get", return_value={"id": 1}) as mock_get:
+            self.superset_api.client.fetch_datasource(7)
+            self.superset_api.client.fetch_datasource(7)
+            self.superset_api.client.fetch_datasource(8)
+        self.assertEqual(mock_get.call_count, 2)
 
     def test_is_table_to_table_lineage(self):
         table = Table(name="table_name", schema=Schema(name="schema_name"))