Source code

001package gudusoft.gsqlparser.stmt.snowflake;
002
003
004import gudusoft.gsqlparser.EDbVendor;
005import gudusoft.gsqlparser.ESqlClause;
006import gudusoft.gsqlparser.ESqlStatementType;
007import gudusoft.gsqlparser.TCustomSqlStatement;
008import gudusoft.gsqlparser.ETableSource;
009import gudusoft.gsqlparser.nodes.TCreateTableOption;
010import gudusoft.gsqlparser.nodes.TObjectName;
011import gudusoft.gsqlparser.nodes.TObjectNameList;
012import gudusoft.gsqlparser.nodes.TParseTreeVisitor;
013import gudusoft.gsqlparser.nodes.TPathSqlNode;
014import gudusoft.gsqlparser.nodes.TStageLocation;
015import gudusoft.gsqlparser.nodes.TTable;
016import gudusoft.gsqlparser.nodes.snowflake.TCopyIntoNode;
017import gudusoft.gsqlparser.nodes.snowflake.TStageReference;
018import gudusoft.gsqlparser.stmt.TSelectSqlStatement;
019
020import java.util.ArrayList;
021
022public class TSnowflakeCopyIntoStmt extends TCustomSqlStatement {
023
024    private ArrayList<String> fileList = new ArrayList<>();
025
026
027    public ArrayList<String> getFileList() {
028        return fileList;
029    }
030
031    private String fileFormatName = null;
032    private String fileFormatType = null;
033
034    public void setFileFormatName(String fileFormatName) {
035        this.fileFormatName = fileFormatName;
036    }
037
038    public void setFileFormatType(String fileFormatType) {
039        this.fileFormatType = fileFormatType;
040    }
041
042    public String getFileFormatName() {
043        return fileFormatName;
044    }
045
046    public String getFileFormatType() {
047        return fileFormatType;
048    }
049
050    public static int COPY_INTO_TABLE = 0;
051    public static int COPY_INTO_LOCATION = 1;
052
053    private String regex_pattern;
054
055    public void setRegex_pattern(String regex_pattern) {
056        this.regex_pattern = regex_pattern;
057    }
058
059    public String getRegex_pattern() {
060        return regex_pattern;
061    }
062
063    private int copyIntoType = TSnowflakeCopyIntoStmt.COPY_INTO_TABLE;
064
065    public int getCopyIntoType() {
066        return copyIntoType;
067    }
068
069    private TStageLocation stageLocation;
070
071    public TStageLocation getStageLocation() {
072        return stageLocation;
073    }
074
075    private TObjectName tableName;
076    private TSelectSqlStatement subQuery;
077    public TSnowflakeCopyIntoStmt(EDbVendor dbvendor) {
078        super(dbvendor);
079        sqlstatementtype = ESqlStatementType.sstCopyInto;
080    }
081
082    public void setTableName(TObjectName tableName) {
083        this.tableName = tableName;
084    }
085
086    public void setSubQuery(TSelectSqlStatement subQuery) {
087        this.subQuery = subQuery;
088    }
089
090    public TObjectName getTableName() {
091
092        return tableName;
093    }
094
095    public TSelectSqlStatement getSubQuery() {
096        return subQuery;
097    }
098
099    private TObjectName fromSourceLocation;
100    public TObjectName getFromSourceLocation() {
101        return fromSourceLocation;
102    }
103
104    private TObjectNameList tableColumnList;
105
106    public void setTableColumnList(TObjectNameList tableColumnList) {
107        this.tableColumnList = tableColumnList;
108    }
109
110    public TObjectNameList getTableColumnList() {
111        return tableColumnList;
112    }
113
114
115    public int doParseStatement(TCustomSqlStatement psql) {
116        if (rootNode == null) return -1;
117        super.doParseStatement(psql);
118        TCopyIntoNode node = (TCopyIntoNode)(rootNode);
119        tableName = node.getTableName();
120        this.copyIntoType = node.getCopyIntoType();
121        this.stageLocation = node.getStageLocation();
122
123        if (node.getSubquery() != null){
124            subQuery = new TSelectSqlStatement(this.dbvendor);
125            subQuery.rootNode = node.getSubquery();
126            subQuery.doParseStatement(this);
127
128            // When the stage location is inside the subquery's FROM clause
129            // (e.g., COPY INTO t FROM (SELECT ... FROM @stage/path)),
130            // propagate it to the COPY INTO statement for API accessibility.
131            if (this.stageLocation == null && subQuery.getTables() != null) {
132                for (int i = 0; i < subQuery.getTables().size(); i++) {
133                    TTable table = subQuery.getTables().getTable(i);
134                    if (table.getTableType() == ETableSource.stageReference
135                            && table.getStageReference() != null) {
136                        TStageReference ref = table.getStageReference();
137                        this.stageLocation = new TStageLocation();
138                        this.stageLocation.setStage(true);
139                        if (ref.getStageName() != null) {
140                            String name = ref.getStageName().toString();
141                            if (name.startsWith("~")) {
142                                this.stageLocation.init(TStageLocation.EStageLocationType.internalUser);
143                            } else if (name.startsWith("%")) {
144                                this.stageLocation.init(TStageLocation.EStageLocationType.internalTable);
145                                TObjectName tblName = new TObjectName();
146                                tblName.setString(name.substring(1));
147                                this.stageLocation.setTableName(tblName);
148                            } else {
149                                this.stageLocation.init(TStageLocation.EStageLocationType.internalNamed, ref.getStageName());
150                            }
151                        }
152                        if (ref.getStagePath() != null) {
153                            this.stageLocation.setPath(ref.getStagePath());
154                        }
155                        break;
156                    }
157                }
158            }
159        }
160
161        ArrayList<TCreateTableOption> tableOptions = node.tableOptions;
162        if (tableOptions != null){
163            for(int i=0;i<tableOptions.size();i++){
164                tableOptions.get(i).doParse(this, ESqlClause.unknown);
165            }
166        }
167
168        fromSourceLocation = node.getFromSourceLocation();
169        tableColumnList = node.getTableColumnList();
170        if (copyIntoType == COPY_INTO_TABLE && tableName != null && subQuery == null) {
171            setTargetTable(new TTable(tableName));
172            if (tableColumnList != null) {
173                tableColumnList.setObjectType(TObjectName.ttobjColumn);
174                tableColumnList.doParse(this, ESqlClause.unknown);
175            }
176        }
177
178        // Mantis #4425: when the FROM clause uses a quoted string literal
179        // (e.g. FROM '@db.schema.stage/path' or FROM 's3://bucket/path'),
180        // the grammar stores it as fromSourceLocation rather than as a
181        // structured stage location, leaving stageLocation null.
182        // Parse the literal and populate stageLocation so the public API
183        // is consistent with the unquoted form.
184        if (this.stageLocation == null && fromSourceLocation != null) {
185            populateStageLocationFromLiteral(fromSourceLocation);
186        }
187
188        return 0;
189    }
190
191    private void populateStageLocationFromLiteral(TObjectName source) {
192        String raw = source.toString();
193        if (raw == null || raw.isEmpty()) return;
194
195        // Strip surrounding single quotes, if any.
196        String content = raw;
197        if (content.length() >= 2 && content.startsWith("'") && content.endsWith("'")) {
198            content = content.substring(1, content.length() - 1);
199        }
200        if (content.isEmpty()) return;
201
202        if (content.startsWith("@")) {
203            populateInternalStageLocation(content);
204        } else if (looksLikeExternalUrl(content)) {
205            TStageLocation loc = new TStageLocation();
206            loc.init(TStageLocation.EStageLocationType.location);
207            loc.setStage(false);
208            TObjectName extLoc = new TObjectName();
209            extLoc.setString(raw);
210            loc.setExternalLocation(extLoc);
211            this.stageLocation = loc;
212        }
213    }
214
215    private void populateInternalStageLocation(String content) {
216        // content begins with '@'
217        String afterAt = content.substring(1);
218
219        // Split into name part and path part at the first '/'.
220        String namePart = afterAt;
221        String pathPart = null;
222        int slashIdx = afterAt.indexOf('/');
223        if (slashIdx >= 0) {
224            namePart = afterAt.substring(0, slashIdx);
225            pathPart = afterAt.substring(slashIdx);
226        }
227
228        TStageLocation loc = new TStageLocation();
229        loc.setStage(true);
230
231        if (namePart.isEmpty() || namePart.equals("~") || namePart.startsWith("~")) {
232            // @~ or @~/path -> user stage
233            loc.init(TStageLocation.EStageLocationType.internalUser);
234        } else if (namePart.startsWith("%")) {
235            // @%table or @%table/path -> table stage
236            loc.init(TStageLocation.EStageLocationType.internalTable);
237            TObjectName tblName = new TObjectName();
238            tblName.setString(namePart.substring(1));
239            loc.setTableName(tblName);
240        } else {
241            // @stage, @schema.stage, or @db.schema.stage (with optional path)
242            TObjectName stageName = new TObjectName();
243            stageName.setString(namePart);
244            loc.init(TStageLocation.EStageLocationType.internalNamed, stageName);
245        }
246
247        if (pathPart != null && !pathPart.isEmpty()) {
248            TPathSqlNode path = new TPathSqlNode();
249            path.setString(pathPart);
250            loc.setPath(path);
251        }
252
253        this.stageLocation = loc;
254    }
255
256    private boolean looksLikeExternalUrl(String content) {
257        String lower = content.toLowerCase(java.util.Locale.ROOT);
258        return lower.startsWith("s3://")
259            || lower.startsWith("azure://")
260            || lower.startsWith("gcs://")
261            || lower.startsWith("https://")
262            || lower.startsWith("http://")
263            || lower.startsWith("file://");
264    }
265
266    public void accept(TParseTreeVisitor v){
267        v.preVisit(this);
268        v.postVisit(this);
269    }
270
271    public void acceptChildren(TParseTreeVisitor v){
272        v.preVisit(this);
273        if (this.subQuery != null){
274            this.subQuery.acceptChildren(v);
275        }
276        v.postVisit(this);
277    }
278
279}