001/*
002 * Copyright (c) 2015-2020, Oracle and/or its affiliates. All rights reserved.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *     http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.tribuo.data.sql;
018
019import com.opencsv.CSVParserWriter;
020import com.opencsv.ICSVWriter;
021import com.opencsv.RFC4180Parser;
022import com.oracle.labs.mlrg.olcut.config.ConfigurationManager;
023import com.oracle.labs.mlrg.olcut.config.Option;
024import com.oracle.labs.mlrg.olcut.config.Options;
025import com.oracle.labs.mlrg.olcut.config.UsageException;
026import com.oracle.labs.mlrg.olcut.util.LabsLogFormatter;
027
028import java.io.BufferedReader;
029import java.io.BufferedWriter;
030import java.io.IOException;
031import java.io.InputStreamReader;
032import java.io.OutputStreamWriter;
033import java.nio.charset.StandardCharsets;
034import java.nio.file.Files;
035import java.nio.file.Path;
036import java.sql.Connection;
037import java.sql.DriverManager;
038import java.sql.ResultSet;
039import java.sql.SQLException;
040import java.sql.Statement;
041import java.util.logging.Level;
042import java.util.logging.Logger;
043
044/**
045 * Read an SQL query in on the standard input, write a CSV file containing the
046 * results to the standard output.
047 * <p>
048 * N.B. This class accepts raw SQL strings and executes them directly via JDBC. It DOES NOT perform
049 * any SQL escaping or other injection prevention. It is the user's responsibility to ensure that SQL passed to this
050 * class performs as desired.
051 */
052public class SQLToCSV {
053
054    public static class SQLToCSVOptions implements Options {
055        @Option(charName='n', longName="connection", usage="Connection string to the SQL database")
056        public String connString;
057        @Option(charName='p', longName="password", usage="Password for the SQL database")
058        public String password;
059        @Option(charName='u', longName="username", usage="Username for the SQL database")
060        public String username;
061        @Option(charName='i', longName="input-sql", usage="SQL File to run as a query, defaults to stdin")
062        public Path inputPath;
063        @Option(charName='o', longName="output-csv", usage="File to write query results as CSV, defaults to stdout")
064        public Path outputPath;
065        @Option(longName="db-config", usage="Name of the DBConfig to use")
066        public SQLDBConfig dbConfig;
067
068    }
069
070    private static final Logger logger = Logger.getLogger(SQLToCSV.class.getName());
071
072    /**
073     * Reads an SQL query from the standard input and writes the results of the
074     * query to the standard output.
075     *
076     * @param args Single arg is the JDBC connection string.
077     */
078    public static void main(String[] args) {
079
080        LabsLogFormatter.setAllLogFormatters();
081
082        SQLToCSVOptions opts = new SQLToCSVOptions();
083        ConfigurationManager cm;
084        try {
085            cm = new ConfigurationManager(args,opts);
086        } catch (UsageException e) {
087            logger.info(e.getUsage());
088            System.exit(1);
089        }
090
091        if (opts.dbConfig == null) {
092            if (opts.connString == null) {
093                logger.log(Level.SEVERE, "Must specify connection string with -n");
094                System.exit(1);
095            }
096
097            if (opts.username != null || opts.password != null) {
098                if (opts.username == null || opts.password == null) {
099                    logger.log(Level.SEVERE, "Must specify both of user and password with -u, -p if one is specified!");
100                    System.exit(1);
101                }
102            }
103        } else if(opts.username != null || opts.password != null || opts.connString != null) {
104            logger.warning("dbConfig provided but username/password/connstring also provided. Options from -u, -p, -n being ignored");
105        }
106
107
108        String query;
109        try (BufferedReader br = opts.inputPath != null ? Files.newBufferedReader(opts.inputPath) : new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8))) {
110            StringBuilder qsb = new StringBuilder();
111            String l;
112            while ((l = br.readLine()) != null) {
113                qsb.append(l);
114                qsb.append("\n");
115            }
116            query = qsb.toString().trim();
117        } catch (IOException ex) {
118            logger.log(Level.SEVERE, "Error reading query: " + ex);
119            System.exit(1);
120            return;
121        }
122
123        if (query.isEmpty()) {
124            logger.log(Level.SEVERE, "Query is empty string");
125            System.exit(1);
126        }
127
128        Connection conn = null;
129        try {
130            if (opts.dbConfig != null) {
131                conn = opts.dbConfig.getConnection();
132            } else if (opts.username != null) {
133                conn = DriverManager.getConnection(opts.connString, opts.username, opts.password);
134            } else {
135                conn = DriverManager.getConnection(opts.connString);
136            }
137        } catch (SQLException ex) {
138            logger.log(Level.SEVERE, "Can't connect to database: " + opts.connString, ex);
139            System.exit(1);
140        }
141
142
143        try (Statement stmt = conn.createStatement()){
144            stmt.setFetchSize(1000);
145            stmt.setFetchDirection(ResultSet.FETCH_FORWARD);
146
147            ResultSet results;
148            try {
149                results = stmt.executeQuery(query);
150            } catch (SQLException ex) {
151                logger.log(Level.SEVERE, "Error running query", ex);
152                try {
153                    conn.close();
154                } catch (SQLException ex1) {
155                    logger.log(Level.SEVERE, "Failed to close connection", ex1);
156                }
157                return;
158            }
159
160            try(ICSVWriter writer = new CSVParserWriter(opts.outputPath != null ?
161                    Files.newBufferedWriter(opts.outputPath) :
162                    new BufferedWriter(new OutputStreamWriter(System.out, StandardCharsets.UTF_8), 1024 * 1024), new RFC4180Parser(), "\n")) {
163                writer.writeAll(results, true);
164            } catch (IOException ex) {
165                logger.log(Level.SEVERE, "Error writing CSV", ex);
166                System.exit(1);
167            } catch (SQLException ex) {
168                logger.log(Level.SEVERE, "Error retrieving results", ex);
169                System.exit(1);
170            }
171        } catch (SQLException ex) {
172            logger.log(Level.SEVERE, "Couldn't create statement", ex);
173            try {
174                conn.close();
175            } catch (SQLException ex1) {
176                logger.log(Level.SEVERE, "Failed to close connection", ex1);
177            }
178            System.exit(1);
179            return;
180        }
181
182        try {
183            conn.close();
184        } catch (SQLException ex1) {
185            logger.log(Level.SEVERE, "Failed to close connection", ex1);
186        }
187
188    }
189
190}