From 9f09598ccd99d288689abeab9736f62b56c606d0 Mon Sep 17 00:00:00 2001
From: LXW <1264174498@qq.com>
Date: Wed, 15 Nov 2023 09:58:50 +0800
Subject: [PATCH] (improvement) (semantic) add auth check and exception catch
 for batch download (#385)

Co-authored-by: jolunoluo
---
 .../query/service/DownloadServiceImpl.java    | 29 +++++++++++++++----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/semantic/query/src/main/java/com/tencent/supersonic/semantic/query/service/DownloadServiceImpl.java b/semantic/query/src/main/java/com/tencent/supersonic/semantic/query/service/DownloadServiceImpl.java
index b2db92e67..fe357099a 100644
--- a/semantic/query/src/main/java/com/tencent/supersonic/semantic/query/service/DownloadServiceImpl.java
+++ b/semantic/query/src/main/java/com/tencent/supersonic/semantic/query/service/DownloadServiceImpl.java
@@ -132,17 +132,35 @@ public class DownloadServiceImpl implements DownloadService {
             MetricSchemaResp metricSchemaResp = metrics.get(0);
             List<DimSchemaResp> dimensions = getMetricRelaDimensions(metricSchemaResp, dimensionRespMap);
             for (MetricSchemaResp metric : metrics) {
-                DataDownload downloadData = getSingleMetricDownloadData(metric, dimensions,
-                        batchDownloadReq.getDateInfo(), user);
-                WriteSheet writeSheet = EasyExcel.writerSheet("Sheet" + sheetCount)
-                        .head(downloadData.getHeaders()).build();
-                excelWriter.write(downloadData.getData(), writeSheet);
+                try {
+                    DataDownload downloadData = getSingleMetricDownloadData(metric, dimensions,
+                            batchDownloadReq.getDateInfo(), user);
+                    WriteSheet writeSheet = EasyExcel.writerSheet("Sheet" + sheetCount)
+                            .head(downloadData.getHeaders()).build();
+                    excelWriter.write(downloadData.getData(), writeSheet);
+                } catch (RuntimeException e) {
+                    EasyExcel.write(file).sheet("Sheet1").head(buildErrMessageHead())
+                            .doWrite(buildErrMessageData(e.getMessage()));
+                    return;
+                }
             }
             sheetCount++;
         }
         excelWriter.finish();
     }
 
+    private List<List<String>> buildErrMessageHead() {
+        List<List<String>> headers = Lists.newArrayList();
+        headers.add(Lists.newArrayList("异常提示"));
+        return headers;
+    }
+
+    private List<List<String>> buildErrMessageData(String errMsg) {
+        List<List<String>> data = Lists.newArrayList();
+        data.add(Lists.newArrayList(errMsg));
+        return data;
+    }
+
     public DataDownload getSingleMetricDownloadData(MetricSchemaResp metricSchemaResp, List<DimSchemaResp> dimensions,
                                                     DateConf dateConf, User user) throws Exception {
         QueryResultWithSchemaResp queryResult = getQueryResult(dimensions, metricSchemaResp, dateConf, user);
@@ -202,6 +220,7 @@ public class DownloadServiceImpl implements DownloadService {
         queryStructReq.setAggregators(Lists.newArrayList(aggregator));
         queryStructReq.setDateInfo(dateConf);
         queryStructReq.setModelId(metricResp.getModelId());
+        queryStructReq.setLimit(10000L);
         return queryService.queryByStructWithAuth(queryStructReq, user);
     }