001package gudusoft.gsqlparser.ir.builder.postgresql;
002
003import gudusoft.gsqlparser.*;
004import gudusoft.gsqlparser.analyzer.v2.AnalyzerV2Config;
005import gudusoft.gsqlparser.ir.bound.BoundArgument;
006import gudusoft.gsqlparser.ir.bound.BoundObjectRef;
007import gudusoft.gsqlparser.ir.bound.BoundParameterSymbol;
008import gudusoft.gsqlparser.ir.bound.BoundProgram;
009import gudusoft.gsqlparser.ir.bound.BoundRoutineRef;
010import gudusoft.gsqlparser.ir.bound.BoundRoutineSymbol;
011import gudusoft.gsqlparser.ir.bound.BoundScope;
012import gudusoft.gsqlparser.ir.bound.BoundTypeRef;
013import gudusoft.gsqlparser.ir.bound.BoundVariableSymbol;
014import gudusoft.gsqlparser.ir.bound.EBindingStatus;
015import gudusoft.gsqlparser.ir.bound.EObjectRefKind;
016import gudusoft.gsqlparser.ir.bound.ERoutineKind;
017import gudusoft.gsqlparser.ir.bound.EScopeKind;
018import gudusoft.gsqlparser.ir.bound.ETypeCategory;
019import gudusoft.gsqlparser.ir.builder.common.AbstractSymbolCollector;
020import gudusoft.gsqlparser.ir.common.Confidence;
021import gudusoft.gsqlparser.ir.common.Evidence;
022import gudusoft.gsqlparser.ir.common.EvidenceKind;
023import gudusoft.gsqlparser.ir.common.SourceAnchor;
024import gudusoft.gsqlparser.nodes.*;
025import gudusoft.gsqlparser.stmt.*;
026import gudusoft.gsqlparser.stmt.TDoExecuteBlockStmt;
027import gudusoft.gsqlparser.stmt.TExecuteSqlStatement;
028
029import java.util.*;
030
031/**
032 * PostgreSQL PL/pgSQL symbol collector.
033 * <p>
034 * Extracts routines, variables, table references, and call references
035 * from PostgreSQL PL/pgSQL AST nodes. Handles:
036 * <ul>
037 *   <li>CREATE FUNCTION / CREATE PROCEDURE</li>
038 *   <li>CREATE TRIGGER</li>
039 *   <li>DO blocks (anonymous blocks)</li>
040 *   <li>PERFORM (SELECT that discards result)</li>
041 *   <li>RETURN QUERY / RETURN NEXT</li>
042 *   <li>EXCEPTION WHEN handlers</li>
043 *   <li>EXECUTE dynamic SQL</li>
044 *   <li>FOR ... IN query loops</li>
045 * </ul>
046 */
047public class PostgresqlSymbolCollector extends AbstractSymbolCollector {
048
049    private final List<String> searchPath;
050    private final Set<TFunctionCall> visitedFunctionCalls = new HashSet<TFunctionCall>();
051    private final Deque<Boolean> commonBlockPushedRoutine = new ArrayDeque<Boolean>();
052
053    public PostgresqlSymbolCollector(BoundProgram program, BoundScope globalScope,
054                                     AnalyzerV2Config config, String fileId,
055                                     List<String> searchPath) {
056        super(program, globalScope, config, fileId);
057        this.searchPath = searchPath != null ? searchPath
058                : Arrays.asList("public");
059    }
060
061    @Override
062    protected boolean isExternalDependency(String name, List<String> parts) {
063        return PostgresqlExternalDepClassifier.isExternal(name);
064    }
065
066    @Override
067    protected String routineIdFor(Object routineNode) {
068        return null; // Not used — we build IDs inline
069    }
070
071    // ====== CREATE FUNCTION ======
072
073    @Override
074    public void preVisit(TCreateFunctionStmt node) {
075        if (node.dbvendor != EDbVendor.dbvpostgresql) return;
076
077        TObjectName funcName = node.getFunctionName();
078        if (funcName == null) return;
079
080        String rawName = funcName.toString().trim();
081        String schema = null;
082        String name = rawName;
083        int dot = rawName.lastIndexOf('.');
084        if (dot >= 0) {
085            schema = rawName.substring(0, dot);
086            name = rawName.substring(dot + 1);
087        }
088
089        List<BoundParameterSymbol> params = extractParameters(node.getParameterDeclarations());
090
091        BoundScope routineScope = pushScope(EScopeKind.ROUTINE, anchor(node));
092
093        // Determine if this is a function or procedure based on the AST
094        ERoutineKind kind = ERoutineKind.FUNCTION;
095
096        BoundTypeRef returnType = null;
097        if (node.getReturnDataType() != null) {
098            returnType = new BoundTypeRef(node.getReturnDataType().toString().trim(),
099                    ETypeCategory.SCALAR);
100        }
101
102        String namespace = schema != null ? schema.toUpperCase() : null;
103        BoundRoutineSymbol routine = new BoundRoutineSymbol(
104                name.toUpperCase(), namespace,
105                routineScope, anchor(node), params, returnType, kind);
106
107        program.registerRoutine(routine);
108        pushRoutineId(routine.getRoutineId());
109
110        for (BoundParameterSymbol p : params) {
111            routineScope.declareSymbol(p.getParamName().toUpperCase(), p);
112        }
113    }
114
115    @Override
116    public void postVisit(TCreateFunctionStmt node) {
117        if (node.dbvendor != EDbVendor.dbvpostgresql) return;
118        popScope();
119        popRoutineId();
120    }
121
122    // ====== CREATE PROCEDURE ======
123
124    @Override
125    public void preVisit(TCreateProcedureStmt node) {
126        if (node.dbvendor != EDbVendor.dbvpostgresql) return;
127
128        TObjectName procName = node.getStoredProcedureName();
129        if (procName == null) return;
130
131        String rawName = procName.toString().trim();
132        String schema = null;
133        String name = rawName;
134        int dot = rawName.lastIndexOf('.');
135        if (dot >= 0) {
136            schema = rawName.substring(0, dot);
137            name = rawName.substring(dot + 1);
138        }
139
140        List<BoundParameterSymbol> params = extractParameters(node.getParameterDeclarations());
141
142        BoundScope routineScope = pushScope(EScopeKind.ROUTINE, anchor(node));
143
144        String namespace = schema != null ? schema.toUpperCase() : null;
145        BoundRoutineSymbol routine = new BoundRoutineSymbol(
146                name.toUpperCase(), namespace,
147                routineScope, anchor(node), params, null, ERoutineKind.PROCEDURE);
148
149        program.registerRoutine(routine);
150        pushRoutineId(routine.getRoutineId());
151
152        for (BoundParameterSymbol p : params) {
153            routineScope.declareSymbol(p.getParamName().toUpperCase(), p);
154        }
155    }
156
157    @Override
158    public void postVisit(TCreateProcedureStmt node) {
159        if (node.dbvendor != EDbVendor.dbvpostgresql) return;
160        popScope();
161        popRoutineId();
162    }
163
164    // ====== CREATE TRIGGER ======
165
166    @Override
167    public void preVisit(TCreateTriggerStmt node) {
168        if (node.dbvendor != EDbVendor.dbvpostgresql) return;
169
170        TObjectName triggerName = node.getTriggerName();
171        if (triggerName == null) return;
172
173        String rawName = triggerName.toString().trim();
174        String name = rawName;
175        int dot = rawName.lastIndexOf('.');
176        if (dot >= 0) {
177            name = rawName.substring(dot + 1);
178        }
179
180        BoundScope routineScope = pushScope(EScopeKind.ROUTINE, anchor(node));
181
182        BoundRoutineSymbol routine = new BoundRoutineSymbol(
183                name.toUpperCase(), null,
184                routineScope, anchor(node),
185                Collections.<BoundParameterSymbol>emptyList(), null, ERoutineKind.TRIGGER);
186
187        program.registerRoutine(routine);
188        pushRoutineId(routine.getRoutineId());
189
190        // Track trigger table
191        TTable onTable = node.getOnTable();
192        if (onTable != null) {
193            collectTableRef(onTable);
194        }
195
196        // Register OLD/NEW pseudo-records
197        routineScope.declareSymbol("OLD",
198                new BoundVariableSymbol("OLD", routineScope, null, null, false));
199        routineScope.declareSymbol("NEW",
200                new BoundVariableSymbol("NEW", routineScope, null, null, false));
201    }
202
203    @Override
204    public void postVisit(TCreateTriggerStmt node) {
205        if (node.dbvendor != EDbVendor.dbvpostgresql) return;
206        popScope();
207        popRoutineId();
208    }
209
210    // ====== Anonymous Block (DO block via TDoExecuteBlockStmt) ======
211
212    @Override
213    public void preVisit(TDoExecuteBlockStmt node) {
214        BoundScope blockScope = pushScope(EScopeKind.BLOCK, anchor(node));
215
216        BoundRoutineSymbol routine = new BoundRoutineSymbol(
217                "<anonymous>", null, blockScope, anchor(node),
218                Collections.<BoundParameterSymbol>emptyList(), null, ERoutineKind.ANONYMOUS_BLOCK);
219        program.registerRoutine(routine);
220        pushRoutineId(routine.getRoutineId());
221    }
222
223    @Override
224    public void postVisit(TDoExecuteBlockStmt node) {
225        popScope();
226        popRoutineId();
227    }
228
229    // ====== Anonymous Block (nested TCommonBlock) ======
230
231    @Override
232    public void preVisit(TCommonBlock node) {
233        BoundScope blockScope = pushScope(EScopeKind.BLOCK, anchor(node));
234
235        if (currentRoutineId() == null) {
236            BoundRoutineSymbol routine = new BoundRoutineSymbol(
237                    "<anonymous>", null, blockScope, anchor(node),
238                    Collections.<BoundParameterSymbol>emptyList(), null, ERoutineKind.ANONYMOUS_BLOCK);
239            program.registerRoutine(routine);
240            pushRoutineId(routine.getRoutineId());
241            commonBlockPushedRoutine.push(Boolean.TRUE);
242        } else {
243            commonBlockPushedRoutine.push(Boolean.FALSE);
244        }
245    }
246
247    @Override
248    public void postVisit(TCommonBlock node) {
249        popScope();
250        Boolean pushed = commonBlockPushedRoutine.isEmpty() ? Boolean.FALSE : commonBlockPushedRoutine.pop();
251        if (pushed) {
252            popRoutineId();
253        }
254    }
255
256    // ====== Variable Declaration ======
257
258    @Override
259    public void preVisit(TVarDeclStmt node) {
260        if (node.getElementName() == null) return;
261        String name = node.getElementName().toString().trim();
262        BoundTypeRef typeRef = null;
263        if (node.getDataType() != null) {
264            typeRef = new BoundTypeRef(node.getDataType().toString().trim(), ETypeCategory.SCALAR);
265        }
266
267        boolean isConstant = (node.getWhatDeclared() == TVarDeclStmt.whatDeclared_constant);
268
269        BoundVariableSymbol varSym = new BoundVariableSymbol(
270                name, currentScope(), anchor(node), typeRef, isConstant);
271        currentScope().declareSymbol(name.toUpperCase(), varSym);
272    }
273
274    // ====== Cursor Declaration ======
275
276    @Override
277    public void preVisit(TCursorDeclStmt node) {
278        if (node.getCursorName() == null) return;
279        String name = node.getCursorName().toString().trim();
280        BoundTypeRef cursorType = new BoundTypeRef("CURSOR", ETypeCategory.REF_CURSOR);
281
282        BoundVariableSymbol cursorSym = new BoundVariableSymbol(
283                name, currentScope(), anchor(node), cursorType, false);
284        currentScope().declareSymbol(name.toUpperCase(), cursorSym);
285
286        if (node.getSubquery() != null) {
287            extractTableRefsFromStatement(node.getSubquery());
288        }
289    }
290
291    // ====== Function Calls in Expressions ======
292
293    @Override
294    public void preVisit(TFunctionCall funcCall) {
295        if (visitedFunctionCalls.contains(funcCall)) return;
296        visitedFunctionCalls.add(funcCall);
297
298        if (funcCall.getFunctionName() == null) return;
299        String funcName = funcCall.getFunctionName().toString().trim();
300        if (funcName.isEmpty()) return;
301
302        // Skip built-in functions
303        if (PostgresqlExternalDepClassifier.isExternal(funcName)) {
304            return;
305        }
306
307        List<String> nameParts = splitName(funcName);
308        List<BoundArgument> args = new ArrayList<BoundArgument>();
309        if (funcCall.getArgs() != null) {
310            for (int i = 0; i < funcCall.getArgs().size(); i++) {
311                TExpression argExpr = funcCall.getArgs().getExpression(i);
312                args.add(new BoundArgument(null, argExpr.toString(),
313                        gudusoft.gsqlparser.ir.bound.EParameterMode.IN, SourceAnchor.from(argExpr)));
314            }
315        }
316
317        BoundRoutineRef ref = new BoundRoutineRef(
318                funcName, nameParts,
319                EBindingStatus.UNRESOLVED_SOFT, null, null, args,
320                new Evidence(EvidenceKind.STATIC_RESOLVED, "Function call in expression"),
321                Confidence.HIGH);
322        ref.setSourceAnchor(anchor(funcCall));
323        if (currentRoutineId() != null) {
324            ref.setProperty("owningRoutineId", currentRoutineId());
325            ref.setProperty("argCount", args.size());
326        }
327        program.addRoutineRef(ref);
328    }
329
330    // ====== CALL Statement ======
331
332    @Override
333    public void preVisit(TCallStatement node) {
334        TObjectName routineName = node.getRoutineName();
335        if (routineName == null) return;
336
337        String callText = routineName.toString().trim();
338        List<String> nameParts = splitName(callText);
339        List<BoundArgument> args = new ArrayList<BoundArgument>();
340        if (node.getArgs() != null) {
341            for (int i = 0; i < node.getArgs().size(); i++) {
342                TExpression argExpr = node.getArgs().getExpression(i);
343                args.add(new BoundArgument(null, argExpr.toString(),
344                        gudusoft.gsqlparser.ir.bound.EParameterMode.IN, SourceAnchor.from(argExpr)));
345            }
346        }
347
348        BoundRoutineRef ref = new BoundRoutineRef(
349                callText, nameParts,
350                EBindingStatus.UNRESOLVED_SOFT, null, null, args,
351                new Evidence(EvidenceKind.STATIC_RESOLVED, "CALL statement"),
352                Confidence.HIGH);
353        ref.setSourceAnchor(anchor(node));
354        if (currentRoutineId() != null) {
355            ref.setProperty("owningRoutineId", currentRoutineId());
356            ref.setProperty("argCount", args.size());
357        }
358        program.addRoutineRef(ref);
359    }
360
361    // ====== EXECUTE dynamic SQL (PostgreSQL uses TExecuteSqlStatement) ======
362
363    @Override
364    public void preVisit(TExecuteSqlStatement node) {
365        if (node.dbvendor != EDbVendor.dbvpostgresql) return;
366
367        String dynText = node.toString().trim();
368        Confidence confidence;
369        String evidenceMsg;
370
371        // Check if the SQL text is a literal
372        String sqlText = node.getSqlText();
373        boolean isLiteral = sqlText != null && !sqlText.isEmpty();
374
375        if (isLiteral) {
376            confidence = Confidence.MEDIUM;
377            evidenceMsg = "Dynamic SQL from string literal";
378        } else {
379            confidence = Confidence.LOW;
380            evidenceMsg = "Dynamic SQL from variable/expression";
381        }
382
383        BoundRoutineRef ref = new BoundRoutineRef(
384                "EXECUTE", Collections.singletonList("EXECUTE"),
385                EBindingStatus.UNRESOLVED_SOFT, null, null,
386                Collections.<BoundArgument>emptyList(),
387                new Evidence(EvidenceKind.DYNAMIC_SQL_TEMPLATE, evidenceMsg),
388                confidence);
389        ref.setSourceAnchor(anchor(node));
390        if (currentRoutineId() != null) {
391            ref.setProperty("owningRoutineId", currentRoutineId());
392        }
393        ref.setProperty("externalDependency", true);
394        ref.setProperty("externalType", "DYNAMIC_SQL");
395        program.addRoutineRef(ref);
396    }
397
398    // ====== SQL Statement Table Collection ======
399
400    @Override
401    public void preVisit(TSelectSqlStatement node) {
402        extractTableRefsFromStatement(node);
403    }
404
405    @Override
406    public void preVisit(TInsertSqlStatement node) {
407        extractTableRefsFromStatement(node);
408    }
409
410    @Override
411    public void preVisit(TUpdateSqlStatement node) {
412        extractTableRefsFromStatement(node);
413    }
414
415    @Override
416    public void preVisit(TDeleteSqlStatement node) {
417        extractTableRefsFromStatement(node);
418    }
419
420    // ====== Table Reference Collection ======
421
422    private void collectTableRef(TTable table) {
423        if (table == null || table.getTableName() == null) return;
424
425        String tableName = table.getTableName().toString().trim();
426        if (tableName.isEmpty()) return;
427
428        EObjectRefKind kind = EObjectRefKind.TABLE;
429        if (table.getSubquery() != null) {
430            kind = EObjectRefKind.DERIVED_TABLE;
431        }
432
433        createObjectRef(tableName, splitName(tableName), kind, anchor(table));
434    }
435
436    private void extractTableRefsFromStatement(TCustomSqlStatement stmt) {
437        if (stmt == null || stmt.tables == null) return;
438        for (int i = 0; i < stmt.tables.size(); i++) {
439            TTable table = stmt.tables.getTable(i);
440            collectTableRef(table);
441        }
442    }
443
444    // ====== Parameter Extraction ======
445
446    private List<BoundParameterSymbol> extractParameters(TParameterDeclarationList paramList) {
447        if (paramList == null || paramList.size() == 0) {
448            return Collections.emptyList();
449        }
450
451        List<BoundParameterSymbol> params = new ArrayList<BoundParameterSymbol>();
452        for (int i = 0; i < paramList.size(); i++) {
453            TParameterDeclaration param = paramList.getParameterDeclarationItem(i);
454            if (param == null || param.getParameterName() == null) continue;
455
456            String name = param.getParameterName().toString().trim();
457            gudusoft.gsqlparser.ir.bound.EParameterMode mode =
458                    gudusoft.gsqlparser.ir.bound.EParameterMode.IN;
459            gudusoft.gsqlparser.EParameterMode pdMode = param.getParameterMode();
460            if (pdMode == gudusoft.gsqlparser.EParameterMode.output
461                    || pdMode == gudusoft.gsqlparser.EParameterMode.out
462                    || pdMode == gudusoft.gsqlparser.EParameterMode.inout) {
463                mode = gudusoft.gsqlparser.ir.bound.EParameterMode.IN_OUT;
464            } else if (pdMode == gudusoft.gsqlparser.EParameterMode.out) {
465                mode = gudusoft.gsqlparser.ir.bound.EParameterMode.OUT;
466            }
467
468            BoundTypeRef typeRef = null;
469            if (param.getDataType() != null) {
470                typeRef = new BoundTypeRef(param.getDataType().toString().trim(),
471                        ETypeCategory.SCALAR);
472            }
473
474            params.add(new BoundParameterSymbol(name, null, anchor(param), typeRef, mode));
475        }
476        return params;
477    }
478}