001package gudusoft.gsqlparser.nodes.functions;
002
003import gudusoft.gsqlparser.nodes.TParseTreeNode;
004import gudusoft.gsqlparser.nodes.TParseTreeVisitor;
005import gudusoft.gsqlparser.ESqlClause;
006import gudusoft.gsqlparser.TCustomSqlStatement;
007import gudusoft.gsqlparser.nodes.TObjectName;
008import gudusoft.gsqlparser.nodes.TResultColumnList;
009import gudusoft.gsqlparser.nodes.TSelectSqlNode;
010import gudusoft.gsqlparser.stmt.TSelectSqlStatement;
011
012
013public class TPredictArgument extends TParseTreeNode {
014    private TObjectName modelReference;  // For RW_MODEL model_reference
015
016    private TResultColumnList structColumns; // For RW_STRUCT '(' target_list ')'
017    private EArgumentType argumentType;   // To track which type of argument this is
018    private TObjectName tableReference;    // For RW_TABLE qualified_name
019    private TSelectSqlNode selectSqlNode;  // For select_with_parens
020    private TSelectSqlStatement selectSqlStatement;
021
022    public enum EArgumentType {
023        MODEL,      // RW_MODEL model_reference
024        INPUT_DATA, // input_data
025        STRUCT      // RW_STRUCT '(' target_list ')'
026    }
027
028    public TSelectSqlStatement getSelectSqlStatement() {
029        return selectSqlStatement;
030    }
031
032    public TObjectName getTableReference() {
033        return tableReference;
034    }
035
036    public void init(Object arg1, Object arg2){
037        this.argumentType = (EArgumentType) arg1;
038        switch (this.argumentType) {
039            case MODEL:
040                this.modelReference = (TObjectName) arg2;
041                break;
042            case INPUT_DATA:
043                if (arg2 instanceof TObjectName) {
044                    this.tableReference = (TObjectName) arg2;
045                } else {
046                    this.selectSqlNode = (TSelectSqlNode) arg2;
047                }
048                break;
049            case STRUCT:
050                this.structColumns = (TResultColumnList) arg2;
051                break;
052        }
053    }
054
055     public void doParse(TCustomSqlStatement psql, ESqlClause plocation){
056        if (this.argumentType == EArgumentType.INPUT_DATA) {
057            if (this.tableReference != null) {
058                this.tableReference.doParse(psql, plocation);
059            } else {
060                selectSqlStatement = new TSelectSqlStatement(this.dbvendor);
061                selectSqlStatement.rootNode = this.selectSqlNode;
062                selectSqlStatement.doParseStatement(psql);
063            }
064        }
065    }
066    // Getters and setters
067    public TObjectName getModelReference() {
068        return modelReference;
069    }
070
071    public void setModelReference(TObjectName modelReference) {
072        this.modelReference = modelReference;
073        this.argumentType = EArgumentType.MODEL;
074    }
075
076
077    public TResultColumnList getStructColumns() {
078        return structColumns;
079    }
080
081    public void setStructColumns(TResultColumnList structColumns) {
082        this.structColumns = structColumns;
083        this.argumentType = EArgumentType.STRUCT;
084    }
085
086    public EArgumentType getArgumentType() {
087        return argumentType;
088    } 
089
090    public void accept(TParseTreeVisitor v){
091        v.preVisit(this);
092        v.postVisit(this);
093    }
094
095    public void acceptChildren(TParseTreeVisitor v) {
096        v.preVisit(this);
097        v.postVisit(this);
098    }    
099}