001package gudusoft.gsqlparser.sqlenv.parser.sharded;
002
003import gudusoft.gsqlparser.EDbVendor;
004import gudusoft.gsqlparser.sqlenv.TSQLEnv;
005import gudusoft.gsqlparser.sqlenv.TSQLCatalog;
006import gudusoft.gsqlparser.sqlenv.TSQLSchema;
007import gudusoft.gsqlparser.sqlenv.TSQLTable;
008import gudusoft.gsqlparser.util.SQLUtil;
009import gudusoft.gsqlparser.util.json.JSON;
010
011import java.io.File;
012import java.io.IOException;
013import java.nio.charset.StandardCharsets;
014import java.nio.file.Files;
015import java.nio.file.Paths;
016import java.util.ArrayList;
017import java.util.List;
018import java.util.Map;
019import java.util.logging.Level;
020import java.util.logging.Logger;
021
022/**
023 * SQLEnv implementation for sqlflow-sharded format.
024 * <p>
025 * Parses catalog.json files from the sharded export format, loading metadata
026 * including tables, views, columns, and synonyms. Supports different database
027 * topologies based on the manifest configuration.
028 * <p>
029 * The sharded format structure:
030 * <pre>
031 * manifest.json
032 * catalog/&lt;shardId&gt;.catalog.json
033 * source/&lt;shardId&gt;.source.jsonl
034 * index/&lt;shardId&gt;.source.idx
035 * </pre>
036 */
037@SuppressWarnings("rawtypes")
038public class SqlflowShardedSQLEnv extends TSQLEnv {
039
040    private Map server;
041    private String baseDir;
042    private Boolean init = false;
043
044    /**
045     * Constructor with explicit base directory for catalog files.
046     *
047     * @param defaultServer   default server name
048     * @param defaultDatabase default database name
049     * @param defaultSchema   default schema name
050     * @param dbVendor        database vendor
051     * @param server          server configuration from manifest
052     * @param baseDir         base directory containing catalog/, source/, index/ subdirectories
053     */
054    public SqlflowShardedSQLEnv(String defaultServer, String defaultDatabase, String defaultSchema,
055                                 EDbVendor dbVendor, Map server, String baseDir) {
056        super(dbVendor);
057        initDefaults(defaultServer, defaultDatabase, defaultSchema);
058        this.server = server;
059        this.baseDir = baseDir;
060        initSQLEnv();
061    }
062
063    private void initDefaults(String defaultServer, String defaultDatabase, String defaultSchema) {
064        if (!SQLUtil.isEmpty(defaultServer) && !defaultServer.equals(TSQLEnv.DEFAULT_SERVER_NAME)) {
065            setDefaultServerName(defaultServer);
066        }
067        if (!SQLUtil.isEmpty(defaultDatabase) && !defaultDatabase.equals(TSQLEnv.DEFAULT_DB_NAME)) {
068            setDefaultCatalogName(defaultDatabase);
069        }
070        if (!SQLUtil.isEmpty(defaultSchema) && !defaultSchema.equals(TSQLEnv.DEFAULT_SCHEMA_NAME)) {
071            setDefaultSchemaName(defaultSchema);
072        }
073    }
074
075    @Override
076    public void initSQLEnv() {
077        synchronized (init) {
078            if (init)
079                return;
080
081            if (server == null) {
082                init = true;
083                return;
084            }
085
086            try {
087                // Parse dbVendor from server config
088                String dbVendor = (String) server.get("dbVendor");
089                EDbVendor vendor = getDBVendor();
090                if (dbVendor != null) {
091                    vendor = EDbVendor.valueOf(dbVendor);
092                }
093
094                // Set server name from config
095                if (!SQLUtil.isEmpty((String) server.get("name"))) {
096                    this.setDefaultServerName((String) server.get("name"));
097                }
098
099                boolean supportsCatalogs = TSQLEnv.supportCatalog(vendor);
100                boolean supportsSchemas = TSQLEnv.supportSchema(vendor);
101
102                // Get nameLevels for topology determination
103                List<String> nameLevels = getNameLevels(server);
104
105                // Resolve base directory
106                String resolvedBaseDir = resolveBaseDir();
107
108                // List for pending synonyms (resolved after all tables are loaded)
109                List<Object[]> pendingSynonyms = new ArrayList<>();
110
111                // Process based on topology
112                if (supportsCatalogs && supportsSchemas) {
113                    // Catalog + Schema topology (e.g., SQL Server, PostgreSQL)
114                    processDatabases(server, resolvedBaseDir, pendingSynonyms);
115                } else if (supportsCatalogs) {
116                    // Catalog-only topology (e.g., MySQL)
117                    processDatabasesAsCatalog(server, resolvedBaseDir, pendingSynonyms);
118                } else if (supportsSchemas) {
119                    // Schema-only topology (e.g., Oracle)
120                    processSchemas(server, resolvedBaseDir, pendingSynonyms);
121                }
122
123                // Resolve pending synonyms
124                for (Object[] pending : pendingSynonyms) {
125                    CatalogMetadataHelper.appendSynonym((TSQLSchema) pending[0], (Map) pending[1]);
126                }
127
128            } catch (Exception e) {
129                Logger.getLogger(SqlflowShardedSQLEnv.class.getName())
130                        .log(Level.WARNING, "Parse sharded catalog failed", e);
131            }
132
133            init = true;
134        }
135    }
136
137    /**
138     * Get nameLevels from server config, defaulting based on vendor capabilities.
139     */
140    private List<String> getNameLevels(Map server) {
141        List<String> nameLevels = (List<String>) server.get("nameLevels");
142        if (nameLevels != null && !nameLevels.isEmpty()) {
143            return nameLevels;
144        }
145
146        // Default based on supportsCatalogs/supportsSchemas
147        Boolean supportsCatalogs = (Boolean) server.get("supportsCatalogs");
148        Boolean supportsSchemas = (Boolean) server.get("supportsSchemas");
149
150        List<String> defaultLevels = new ArrayList<>();
151        if (supportsCatalogs != null && supportsCatalogs) {
152            defaultLevels.add("catalog");
153        }
154        if (supportsSchemas != null && supportsSchemas) {
155            defaultLevels.add("schema");
156        }
157
158        return defaultLevels;
159    }
160
161    /**
162     * Get base directory.
163     */
164    private String resolveBaseDir() {
165        return baseDir;
166    }
167
168    /**
169     * Process databases in catalog+schema topology.
170     */
171    private void processDatabases(Map server, String baseDir, List<Object[]> pendingSynonyms) {
172        List databases = (List) server.get("databases");
173        if (databases == null) {
174            return;
175        }
176
177        for (int i = 0; i < databases.size(); i++) {
178            Map jsonDatabase = (Map) databases.get(i);
179            String databaseName = (String) jsonDatabase.get("name");
180
181            // Get catalog path
182            Map catalogInfo = (Map) jsonDatabase.get("catalog");
183            if (catalogInfo == null) {
184                continue;
185            }
186
187            String catalogPath = (String) catalogInfo.get("path");
188            if (SQLUtil.isEmpty(catalogPath)) {
189                continue;
190            }
191
192            // Resolve full path to catalog file
193            String fullCatalogPath = resolvePath(baseDir, catalogPath);
194            Map catalogJson = loadCatalogJson(fullCatalogPath);
195            if (catalogJson == null) {
196                continue;
197            }
198
199            // Process schemas within catalog
200            List schemas = (List) catalogJson.get("schemas");
201            if (schemas == null) {
202                // Fallback: database might be directly in catalog root (nameLevels = ["database"])
203                // In this case, create a default schema to hold tables/views
204                TSQLCatalog sqlCatalog = getSQLCatalog(databaseName, true);
205                TSQLSchema sqlSchema = sqlCatalog.getSchema(TSQLEnv.DEFAULT_SCHEMA_NAME, true);
206                appendTables(catalogJson, sqlSchema);
207                appendViews(catalogJson, sqlSchema);
208                CatalogMetadataHelper.appendPackages(catalogJson, sqlSchema);
209                CatalogMetadataHelper.appendProcedures(catalogJson, sqlSchema, null);
210                collectSynonyms(catalogJson, sqlSchema, pendingSynonyms);
211            } else {
212                TSQLCatalog sqlCatalog = getSQLCatalog(databaseName, true);
213                for (int j = 0; j < schemas.size(); j++) {
214                    Map jsonSchema = (Map) schemas.get(j);
215                    String schemaName = (String) jsonSchema.get("name");
216                    TSQLSchema sqlSchema = sqlCatalog.getSchema(schemaName, true);
217                    appendTables(jsonSchema, sqlSchema);
218                    appendViews(jsonSchema, sqlSchema);
219                    CatalogMetadataHelper.appendPackages(jsonSchema, sqlSchema);
220                    CatalogMetadataHelper.appendProcedures(jsonSchema, sqlSchema, null);
221                    collectSynonyms(jsonSchema, sqlSchema, pendingSynonyms);
222                }
223            }
224        }
225    }
226
227    /**
228     * Process databases in catalog-only topology (no schema level).
229     */
230    private void processDatabasesAsCatalog(Map server, String baseDir, List<Object[]> pendingSynonyms) {
231        List databases = (List) server.get("databases");
232        if (databases == null) {
233            return;
234        }
235
236        for (int i = 0; i < databases.size(); i++) {
237            Map jsonDatabase = (Map) databases.get(i);
238            String databaseName = (String) jsonDatabase.get("name");
239
240            Map catalogInfo = (Map) jsonDatabase.get("catalog");
241            if (catalogInfo == null) {
242                continue;
243            }
244
245            String catalogPath = (String) catalogInfo.get("path");
246            if (SQLUtil.isEmpty(catalogPath)) {
247                continue;
248            }
249
250            String fullCatalogPath = resolvePath(baseDir, catalogPath);
251            Map catalogJson = loadCatalogJson(fullCatalogPath);
252            if (catalogJson == null) {
253                continue;
254            }
255
256            // In catalog-only mode, each database is treated as a catalog
257            // Tables/views are at the catalog level (not in schemas)
258            TSQLCatalog sqlCatalog = getSQLCatalog(databaseName, true);
259            TSQLSchema sqlSchema = sqlCatalog.getSchema(TSQLEnv.DEFAULT_SCHEMA_NAME, true);
260
261            // Tables/views might be at root level or in schemas
262            List tables = (List) catalogJson.get("tables");
263            List views = (List) catalogJson.get("views");
264
265            if (tables != null) {
266                for (Object table : tables) {
267                    appendTableObject((Map) table, sqlSchema, false);
268                }
269            }
270            if (views != null) {
271                for (Object view : views) {
272                    appendTableObject((Map) view, sqlSchema, true);
273                }
274            }
275            CatalogMetadataHelper.appendPackages(catalogJson, sqlSchema);
276            CatalogMetadataHelper.appendProcedures(catalogJson, sqlSchema, null);
277        }
278    }
279
280    /**
281     * Process schemas in schema-only topology (e.g., Oracle).
282     */
283    private void processSchemas(Map server, String baseDir, List<Object[]> pendingSynonyms) {
284        List schemas = (List) server.get("schemas");
285        if (schemas == null) {
286            return;
287        }
288
289        for (int i = 0; i < schemas.size(); i++) {
290            Map jsonSchema = (Map) schemas.get(i);
291            String schemaName = (String) jsonSchema.get("name");
292
293            // Get catalog path
294            Map catalogInfo = (Map) jsonSchema.get("catalog");
295            String fullCatalogPath = null;
296            Map catalogJson = null;
297
298            if (catalogInfo != null) {
299                String catalogPath = (String) catalogInfo.get("path");
300                if (!SQLUtil.isEmpty(catalogPath)) {
301                    fullCatalogPath = resolvePath(baseDir, catalogPath);
302                    catalogJson = loadCatalogJson(fullCatalogPath);
303                }
304            }
305
306            // Create synthetic catalog for schema-only vendors
307            TSQLCatalog sqlCatalog = getSQLCatalog(TSQLEnv.DEFAULT_DB_NAME, true);
308            TSQLSchema sqlSchema = sqlCatalog.getSchema(schemaName, true);
309
310            if (catalogJson != null) {
311                appendTables(catalogJson, sqlSchema);
312                appendViews(catalogJson, sqlSchema);
313                CatalogMetadataHelper.appendPackages(catalogJson, sqlSchema);
314                CatalogMetadataHelper.appendProcedures(catalogJson, sqlSchema, null);
315                collectSynonyms(catalogJson, sqlSchema, pendingSynonyms);
316            } else {
317                // Process directly from schema config
318                appendTables(jsonSchema, sqlSchema);
319                appendViews(jsonSchema, sqlSchema);
320                CatalogMetadataHelper.appendPackages(jsonSchema, sqlSchema);
321                CatalogMetadataHelper.appendProcedures(jsonSchema, sqlSchema, null);
322                collectSynonyms(jsonSchema, sqlSchema, pendingSynonyms);
323            }
324        }
325    }
326
327    /**
328     * Load and parse a catalog JSON file.
329     */
330    private Map loadCatalogJson(String path) {
331        try {
332            File file = new File(path);
333            if (!file.exists()) {
334                Logger.getLogger(SqlflowShardedSQLEnv.class.getName())
335                        .warning("Catalog file not found: " + path);
336                return null;
337            }
338
339            String content = new String(Files.readAllBytes(file.toPath()), StandardCharsets.UTF_8);
340            Object parsed = JSON.parseObject(content);
341            if (parsed instanceof Map) {
342                return (Map) parsed;
343            }
344        } catch (IOException e) {
345            Logger.getLogger(SqlflowShardedSQLEnv.class.getName())
346                    .log(Level.WARNING, "Failed to read catalog file: " + path, e);
347        }
348        return null;
349    }
350
351    /**
352     * Resolve relative path against base directory.
353     */
354    private String resolvePath(String baseDir, String relativePath) {
355        if (baseDir == null) {
356            return relativePath;
357        }
358
359        // Handle both forward slash and backslash
360        String normalizedBase = baseDir.replace('\\', '/');
361        String normalizedRelative = relativePath.replace('\\', '/');
362
363        // If relativePath is absolute, return as-is
364        if (new File(relativePath).isAbsolute()) {
365            return relativePath;
366        }
367
368        return Paths.get(normalizedBase, normalizedRelative).toString();
369    }
370
371    /**
372     * Append tables from a schema/catalog JSON object.
373     */
374    private void appendTables(Map jsonSchema, TSQLSchema sqlSchema) {
375        List tables = (List) jsonSchema.get("tables");
376        if (tables == null) {
377            return;
378        }
379
380        for (int i = 0; i < tables.size(); i++) {
381            Map jsonTable = (Map) tables.get(i);
382            appendTableObject(jsonTable, sqlSchema, false);
383        }
384    }
385
386    /**
387     * Append views from a schema/catalog JSON object.
388     */
389    private void appendViews(Map jsonSchema, TSQLSchema sqlSchema) {
390        List views = (List) jsonSchema.get("views");
391        if (views == null) {
392            return;
393        }
394
395        for (int i = 0; i < views.size(); i++) {
396            Map jsonTable = (Map) views.get(i);
397            appendTableObject(jsonTable, sqlSchema, true);
398        }
399    }
400
401    /**
402     * Append a single table or view object.
403     */
404    private void appendTableObject(Map jsonTable, TSQLSchema sqlSchema, boolean isView) {
405        String tableName = (String) jsonTable.get("name");
406        if (SQLUtil.isEmpty(tableName)) {
407            return;
408        }
409
410        TSQLTable sqlTable = sqlSchema.createTable(tableName, isView ? 3 : 0);
411        if (isView) {
412            sqlTable.setView(true);
413        }
414
415        // Set table type if present
416        String type = (String) jsonTable.get("type");
417        if (type != null && type.toLowerCase().indexOf("view") != -1) {
418            sqlTable.setView(true);
419        }
420
421        // Add columns
422        List columns = (List) jsonTable.get("columns");
423        if (columns != null) {
424            for (int j = 0; j < columns.size(); j++) {
425                Map jsonColumn = (Map) columns.get(j);
426                String columnName = (String) jsonColumn.get("name");
427                if (!SQLUtil.isEmpty(columnName)) {
428                    sqlTable.addColumn(columnName);
429                }
430            }
431        }
432    }
433
434    /**
435     * Collect synonyms for later resolution.
436     */
437    private void collectSynonyms(Map jsonSchema, TSQLSchema sqlSchema, List<Object[]> pendingSynonyms) {
438        List synonyms = (List) jsonSchema.get("synonyms");
439        if (synonyms == null) {
440            return;
441        }
442
443        for (int i = 0; i < synonyms.size(); i++) {
444            Map jsonSynonym = (Map) synonyms.get(i);
445            if (jsonSynonym == null) {
446                continue;
447            }
448            pendingSynonyms.add(new Object[]{sqlSchema, jsonSynonym});
449        }
450    }
451}