diff --git a/pebblo/app/models/db_models.py b/pebblo/app/models/db_models.py index ed9dce0e..128dc967 100644 --- a/pebblo/app/models/db_models.py +++ b/pebblo/app/models/db_models.py @@ -149,6 +149,7 @@ class AiDataSource(BaseModel): loadId: str metadata: Metadata sourcePath: str + sourceSize: int sourceType: str loader: str diff --git a/pebblo/app/service/loader/loader_doc_service.py b/pebblo/app/service/loader/loader_doc_service.py index 5549b0a6..e66a4466 100644 --- a/pebblo/app/service/loader/loader_doc_service.py +++ b/pebblo/app/service/loader/loader_doc_service.py @@ -275,6 +275,7 @@ def _get_or_create_data_source(self): "modifiedAt": get_current_time(), }, "sourcePath": loader_details.get("source_path"), + "sourceSize": loader_details.get("source_aggregate_size", 0), "sourceType": loader_details.get("source_type"), "loader": loader_details.get("loader"), } diff --git a/pebblo/app/service/local_ui/loader_apps.py b/pebblo/app/service/local_ui/loader_apps.py index e58e9534..2c0dfa47 100644 --- a/pebblo/app/service/local_ui/loader_apps.py +++ b/pebblo/app/service/local_ui/loader_apps.py @@ -230,6 +230,7 @@ def _update_loader_datasource( "name": ds_data["loader"], "sourcePath": ds_data["sourcePath"], "sourceType": ds_data["sourceType"], + "sourceSize": ds_data.get("sourceSize", "-"), "findingsEntities": entity_count, "findingsTopics": topic_count, "totalSnippetCount": total_snippet_count, @@ -246,6 +247,19 @@ def _update_loader_datasource( self.loader_details.get("loader_data_source_list", []) ) + def _get_data_source_name(self, document_detail: dict) -> str: + """ + Get data source name for given data source id from db. + """ + data_source_id = document_detail.get("dataSourceId") + source_name = "-" + if data_source_id: + _, datasource = self.db.query(AiDataSourceTable, {"id": data_source_id}) + if datasource: + source_name = datasource[0].data.get("loader", "-") + + return source_name + def _get_documents_with_findings(self, app_data: AiDataLoaderTable) -> None: """ Fetch required data for DocumentWithFindings @@ -268,9 +282,12 @@ def _get_documents_with_findings(self, app_data: AiDataLoaderTable) -> None: for topic, topic_data in document_detail.get("topics", {}).items(): topic_count += topic_data.get("count", 0) + source_name = self._get_data_source_name(document_detail) + if document_detail["sourcePath"] in loader_document_with_findings: document_obj = { "appName": document_detail["appName"], + "sourceName": source_name, "owner": document_detail.get("owner", "-"), "sourceSize": document_detail.get("sourceSize", 0), "sourceFilePath": document_detail["sourcePath"],