Source code

001package gudusoft.gsqlparser.dlineage.util;
002
003import gudusoft.gsqlparser.dlineage.dataflow.metadata.MetadataReader;
004import gudusoft.gsqlparser.dlineage.dataflow.metadata.sqlflow.SqlflowObjectAdapter;
005import gudusoft.gsqlparser.dlineage.dataflow.model.Option;
006import gudusoft.gsqlparser.dlineage.dataflow.model.SqlInfo;
007import gudusoft.gsqlparser.dlineage.metadata.Sqlflow;
008import gudusoft.gsqlparser.util.SQLUtil;
009import gudusoft.gsqlparser.util.json.JSON;
010
011import java.io.File;
012import java.io.FileWriter;
013import java.io.IOException;
014import java.util.List;
015import java.util.Map;
016
017import org.slf4j.Logger;
018import org.slf4j.LoggerFactory;
019
020public class LargeFileDetector {
021
022    private static final Logger logger = LoggerFactory.getLogger(LargeFileDetector.class);
023
024    public enum DetectionType {
025        SQL_INFO_COUNT,
026        SQL_TOTAL_SIZE,
027        SHARDED_SOURCE_COUNT,
028        SHARDED_SOURCE_SIZE,
029        MANIFEST_QUERY_COUNT,
030        SINGLE_FILE_SIZE,
031        NONE
032    }
033
034    public static class LargeFileDetectionResult {
035        public final boolean shouldDelegate;
036        public final String reason;
037        public final DetectionType type;
038
039        public LargeFileDetectionResult(boolean shouldDelegate, String reason, DetectionType type) {
040            this.shouldDelegate = shouldDelegate;
041            this.reason = reason;
042            this.type = type;
043        }
044    }
045
046    public static LargeFileDetectionResult detect(SqlInfo[] sqlInfos, File[] originalFiles, Option option) {
047        LargeFileDetectionResult r = checkManifestScale(sqlInfos, option);
048        if (r.shouldDelegate) return r;
049
050        r = checkSqlInfoCount(sqlInfos, option);
051        if (r.shouldDelegate) return r;
052
053        r = checkSqlTotalSize(sqlInfos, option);
054        if (r.shouldDelegate) return r;
055
056        r = checkSingleFileSize(originalFiles, option);
057        if (r.shouldDelegate) return r;
058
059        return checkTotalFileSize(originalFiles, option);
060    }
061
062    public static void persistSqlToTempFile(SqlInfo info, String sql) {
063        try {
064            File tempFile = File.createTempFile("gsp-dataflow-", ".sql");
065            try (FileWriter writer = new FileWriter(tempFile)) {
066                writer.write(sql);
067            }
068            info.setFilePath(tempFile.getAbsolutePath());
069            info.setFileName(tempFile.getName());
070            info.setSql(null);
071            logger.info("Large SQL persisted to temp file: " + tempFile.getAbsolutePath());
072        } catch (IOException e) {
073            logger.error("Failed to persist large SQL to temp file", e);
074        }
075    }
076
077    public static File serializeToTempFile(Map<String, Object> map) throws IOException {
078        File tempFile = File.createTempFile("gsp-dataflow-", ".sql");
079        try (FileWriter writer = new FileWriter(tempFile)) {
080            JSON.toJSONString(map, writer);
081        }
082        return tempFile;
083    }
084
085    public static File writeSqlflowToTempFile(Sqlflow sqlflow) {
086        try {
087            return serializeToTempFile(SqlflowObjectAdapter.toMap(sqlflow));
088        } catch (IOException e) {
089            throw new RuntimeException("Failed to serialize sqlflow to temp file", e);
090        }
091    }
092
093    private static LargeFileDetectionResult checkSqlInfoCount(SqlInfo[] sqlInfos, Option option) {
094        if (sqlInfos == null) return notDelegated();
095        int count = 0;
096        for (SqlInfo info : sqlInfos) {
097            if (info != null) count++;
098        }
099        if (count >= option.getLargeSqlInfoCountThreshold()) {
100            return delegated(DetectionType.SQL_INFO_COUNT,
101                    "SqlInfo count " + count + " >= " + option.getLargeSqlInfoCountThreshold());
102        }
103        return notDelegated();
104    }
105
106    private static LargeFileDetectionResult checkSqlTotalSize(SqlInfo[] sqlInfos, Option option) {
107        if (sqlInfos == null) return notDelegated();
108        long total = 0;
109        for (SqlInfo info : sqlInfos) {
110            if (info != null && info.getSql() != null) {
111                total += info.getSql().length();
112            }
113        }
114        if (total >= option.getLargeSqlTotalSizeThreshold()) {
115            return delegated(DetectionType.SQL_TOTAL_SIZE,
116                    "Total SQL size " + total + " >= " + option.getLargeSqlTotalSizeThreshold());
117        }
118        return notDelegated();
119    }
120
121    private static LargeFileDetectionResult checkSingleFileSize(File[] originalFiles, Option option) {
122        if (originalFiles == null) return notDelegated();
123        long thresholdBytes = option.getLargeFileSplitSizeMB() * 1024L * 1024L;
124        for (File f : originalFiles) {
125            if (f != null && f.isFile() && f.length() >= thresholdBytes) {
126                return delegated(DetectionType.SINGLE_FILE_SIZE,
127                        "File size " + f.length() + " >= " + thresholdBytes
128                                + " for " + f.getName());
129            }
130        }
131        return notDelegated();
132    }
133
134    private static LargeFileDetectionResult checkTotalFileSize(File[] originalFiles, Option option) {
135        if (originalFiles == null) return notDelegated();
136        long total = 0;
137        for (File f : originalFiles) {
138            if (f != null && f.isFile()) {
139                total += f.length();
140            }
141        }
142        if (total >= option.getLargeSqlTotalSizeThreshold()) {
143            return delegated(DetectionType.SQL_TOTAL_SIZE,
144                    "Total file size " + total + " >= " + option.getLargeSqlTotalSizeThreshold());
145        }
146        return notDelegated();
147    }
148
149    private static LargeFileDetectionResult checkManifestScale(SqlInfo[] sqlInfos, Option option) {
150        if (sqlInfos == null) return notDelegated();
151        for (SqlInfo info : sqlInfos) {
152            if (info == null) continue;
153            String sql = info.getSql();
154            if (sql == null && info.getFilePath() != null) {
155                File f = new File(info.getFilePath());
156                if (f.exists() && f.isFile()) {
157                    sql = SQLUtil.getFileContent(f);
158                }
159            }
160            if (sql == null) continue;
161            sql = sql.trim();
162            if (!sql.startsWith("{")) continue;
163
164            try {
165                if (MetadataReader.isSqlflowSharded(sql)) {
166                    return checkShardedManifest(sql, info.getFilePath(), option);
167                }
168                if (MetadataReader.isSqlflow(sql) || MetadataReader.isGrabit(sql)) {
169                    return checkManifestQueryCount(sql, option);
170                }
171            } catch (Exception e) {
172                // JSON 解析失败，跳过
173            }
174        }
175        return notDelegated();
176    }
177
178    static LargeFileDetectionResult checkShardedManifest(String manifestSql, String manifestPath, Option option) {
179        Map manifest = (Map) JSON.parseObject(manifestSql);
180        String baseDir = null;
181        if (manifestPath != null) {
182            File manifestFile = new File(manifestPath);
183            File parent = manifestFile.getParentFile();
184            if (parent != null) {
185                baseDir = parent.getAbsolutePath();
186            }
187        }
188        if (baseDir == null) return notDelegated();
189
190        List<Map> servers = (List<Map>) manifest.get("servers");
191        if (servers == null) return notDelegated();
192
193        int sourceFileCount = 0;
194        long sourceTotalSize = 0;
195
196        for (Map server : servers) {
197            List<Map> shards = (List<Map>) server.get("schemas");
198            if (shards == null || shards.isEmpty()) {
199                shards = (List<Map>) server.get("databases");
200            }
201            if (shards == null) continue;
202
203            for (Map shard : shards) {
204                Map source = (Map) shard.get("source");
205                if (source == null) continue;
206                String sourcePath = (String) source.get("path");
207                if (sourcePath == null) continue;
208
209                File sourceFile = new File(baseDir, sourcePath);
210                if (sourceFile.exists()) {
211                    sourceFileCount++;
212                    sourceTotalSize += sourceFile.length();
213                }
214            }
215        }
216
217        if (sourceFileCount >= option.getLargeShardCountThreshold()) {
218            return delegated(DetectionType.SHARDED_SOURCE_COUNT,
219                    "Sharded source count " + sourceFileCount + " >= " + option.getLargeShardCountThreshold());
220        }
221        long thresholdBytes = option.getLargeFileSplitSizeMB() * 1024L * 1024L;
222        if (sourceTotalSize >= thresholdBytes) {
223            return delegated(DetectionType.SHARDED_SOURCE_SIZE,
224                    "Sharded source total size " + sourceTotalSize + " >= " + thresholdBytes);
225        }
226        return notDelegated();
227    }
228
229    static LargeFileDetectionResult checkManifestQueryCount(String manifestSql, Option option) {
230        Map manifest = (Map) JSON.parseObject(manifestSql);
231        int queryCount = 0;
232
233        List queries = (List) manifest.get("queries");
234        if (queries != null) queryCount += queries.size();
235
236        List<Map> servers = (List<Map>) manifest.get("servers");
237        if (servers != null) {
238            for (Map server : servers) {
239                List serverQueries = (List) server.get("queries");
240                if (serverQueries != null) queryCount += serverQueries.size();
241            }
242        }
243
244        if (queryCount >= option.getLargeQueryCountThreshold()) {
245            return delegated(DetectionType.MANIFEST_QUERY_COUNT,
246                    "Manifest query count " + queryCount + " >= " + option.getLargeQueryCountThreshold());
247        }
248        return notDelegated();
249    }
250
251    private static LargeFileDetectionResult notDelegated() {
252        return new LargeFileDetectionResult(false, null, DetectionType.NONE);
253    }
254
255    private static LargeFileDetectionResult delegated(DetectionType type, String reason) {
256        return new LargeFileDetectionResult(true, reason, type);
257    }
258}