001package gudusoft.gsqlparser.resolver2.context;
002
003import gudusoft.gsqlparser.*;
004import gudusoft.gsqlparser.nodes.TDummy;
005import gudusoft.gsqlparser.nodes.TObjectName;
006import gudusoft.gsqlparser.sqlenv.IdentifierService;
007import gudusoft.gsqlparser.sqlenv.IdentifierProfile;
008import gudusoft.gsqlparser.sqlenv.TSQLEnv;
009import gudusoft.gsqlparser.stmt.TSetSchemaStmt;
010import gudusoft.gsqlparser.stmt.TUseDatabase;
011import gudusoft.gsqlparser.stmt.TUseStmt;
012import java.util.List;
013
014/** Tracks SQL-lexical database/schema defaults in statement order. */
015public class DatabaseContextTracker {
016    private String currentDatabase;
017    private String currentSchema;
018    private String currentSchemaSpelling;
019    private boolean databaseChanged;
020    private boolean schemaChanged;
021
022    public DatabaseContextTracker() { }
023
024    public DatabaseContextTracker(String database, String schema) {
025        currentDatabase = database;
026        currentSchema = schema;
027        currentSchemaSpelling = schema;
028    }
029
030    public void processStatements(TStatementList statements) {
031        if (statements != null) {
032            for (int i = 0; i < statements.size(); i++) processStatement(statements.get(i));
033        }
034    }
035
036    public void processStatement(TCustomSqlStatement stmt) {
037        if (stmt == null) return;
038        ESqlStatementType type = stmt.sqlstatementtype;
039        if (type != ESqlStatementType.sstUseDatabase && type != ESqlStatementType.sstUse
040                && type != ESqlStatementType.sstUseSchema && type != ESqlStatementType.sstUseCatalog
041                && type != ESqlStatementType.sstSetSchema) return;
042
043        boolean schema = type == ESqlStatementType.sstUseSchema || type == ESqlStatementType.sstSetSchema;
044        String name = null;
045        if (stmt instanceof TUseDatabase) {
046            TUseDatabase use = (TUseDatabase) stmt;
047            schema |= use.isSchema() || use.getSchemaName() != null;
048            TObjectName object = use.getSchemaName() != null ? use.getSchemaName() : use.getDatabaseName();
049            if (object != null) name = object.toString();
050        } else if (stmt instanceof TUseStmt) {
051            TUseStmt use = (TUseStmt) stmt;
052            schema |= use.getDbObjectType() == EDbObjectType.schema;
053            if (use.getDbObjectName() != null) name = use.getDbObjectName().toString();
054        } else if (stmt instanceof TSetSchemaStmt) {
055            if (((TSetSchemaStmt) stmt).getSchemaName() != null)
056                name = ((TSetSchemaStmt) stmt).getSchemaName().toString();
057        } else if (stmt.rootNode instanceof TDummy && ((TDummy) stmt.rootNode).node1 instanceof TObjectName) {
058            TDummy node = (TDummy) stmt.rootNode;
059            schema |= node.objectType == EDbObjectType.schema;
060            name = node.node1.toString();
061        } else {
062            // Some command-only dialect entries have no AST. Read lexer tokens, preserving
063            // delimited identifiers (including whitespace, dots, and escaped delimiters).
064            StringBuilder spelling = new StringBuilder();
065            int head = 0;
066            for (int i = 0; i < stmt.sourcetokenlist.size(); i++) {
067                TSourceToken token = stmt.sourcetokenlist.get(i);
068                if (!token.issolidtoken() || token.tokentype == ETokenType.ttsemicolon) continue;
069                if (head++ == 0) continue; // USE / SET
070                if (head == 2 && token.tokentype == ETokenType.ttkeyword) {
071                    if ("SCHEMA".equalsIgnoreCase(token.toString())) { // non-identifier-compare: command keyword
072                        schema = true;
073                        continue;
074                    }
075                    if ("DATABASE".equalsIgnoreCase(token.toString()) // non-identifier-compare: command keyword
076                            || "CATALOG".equalsIgnoreCase(token.toString())) continue; // non-identifier-compare: command keyword
077                }
078                spelling.append(token.toString());
079            }
080            if (spelling.length() > 0) name = spelling.toString();
081        }
082        if (name == null || name.isEmpty()) return;
083        IdentifierService identifiers = new IdentifierService(IdentifierProfile.forVendor(
084                stmt.dbvendor, IdentifierProfile.VendorFlags.defaults()), null);
085        List<String> parts = identifiers.parseQualifiedName(name);
086        // Snowflake USE db.schema and USE SCHEMA db.schema both set the two defaults.
087        if (schema || (stmt.dbvendor == EDbVendor.dbvsnowflake && parts.size() == 2)) {
088            currentSchema = parts.get(parts.size() - 1);
089            currentSchemaSpelling = name;
090            schemaChanged = true;
091            if (parts.size() > 1) {
092                currentDatabase = parts.get(parts.size() - 2);
093                databaseChanged = true;
094            }
095        } else {
096            currentDatabase = name;
097            databaseChanged = true;
098            if (stmt.dbvendor == EDbVendor.dbvsnowflake) {
099                currentSchema = null;
100                currentSchemaSpelling = null;
101                schemaChanged = true;
102            }
103        }
104    }
105
106    public void applyDefaults(TSQLEnv env) {
107        if (env == null) return;
108        if (databaseChanged) env.setDefaultCatalogName(currentDatabase);
109        if (schemaChanged) env.setDefaultSchemaName(currentSchema);
110    }
111
112    public boolean hasContext() { return databaseChanged || schemaChanged; }
113    public String getCurrentDatabase() { return currentDatabase; }
114    /** Original SQL spelling, including a catalog qualifier when supplied. */
115    public String getCurrentSchema() { return currentSchemaSpelling; }
116    /** Schema component for hierarchical catalog lookup. */
117    public String getCurrentSchemaComponent() { return currentSchema; }
118    public void reset() {
119        currentDatabase = null;
120        currentSchema = null;
121        currentSchemaSpelling = null;
122        databaseChanged = false;
123        schemaChanged = false;
124    }
125
126    @Override public String toString() {
127        return "DatabaseContextTracker{database='" + currentDatabase + "', schema='" + currentSchema + "'}";
128    }
129}