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.oracle.labs.mlrg.olcut.config.Config;
020import com.oracle.labs.mlrg.olcut.config.Configurable;
021import com.oracle.labs.mlrg.olcut.config.PropertyException;
022import com.oracle.labs.mlrg.olcut.provenance.ConfiguredObjectProvenance;
023import com.oracle.labs.mlrg.olcut.provenance.Provenancable;
024import com.oracle.labs.mlrg.olcut.provenance.impl.ConfiguredObjectProvenanceImpl;
025
026import java.sql.Connection;
027import java.sql.DriverManager;
028import java.sql.ResultSet;
029import java.sql.SQLException;
030import java.sql.Statement;
031import java.util.HashMap;
032import java.util.Map;
033import java.util.Properties;
034
035/**
036 *
037 * N.B. This class accepts raw SQL strings and executes them directly via JDBC. It DOES NOT perform
038 * any SQL escaping or other injection prevention. It is the user's responsibility to ensure that SQL passed to this
039 * class performs as desired.
040 *
041 * SQL database configuration. If you specify the {@linkplain SQLDBConfig#host}, {@linkplain SQLDBConfig#port}, and
042 * {@linkplain SQLDBConfig#db} strings and use {@code oracle.jdbc.OracleDriver} as your JDBC Driver, then this will
043 * automatically generate a connectionString, otherwise it must be specified manually and host, port, and db fields
044 * can be omitted.
045 *
046 * {@link java.sql.DriverManager}'s default logic will be used to determine which {@link java.sql.Driver} to use for
047 * a given connection string.
048 */
049public class SQLDBConfig implements Configurable, Provenancable<ConfiguredObjectProvenance> {
050
051    @Config(description="Connection string, including host, port and db.")
052    private String connectionString;
053    @Config(description="Database username.",redact=true)
054    private String username;
055    @Config(description="Database password.",redact=true)
056    private String password;
057
058    @Config(description="Properties to pass to java.sql.DriverManager, username and password will be removed and populated to their fields. If specified both on the map and in the fields, the fields will be used")
059    private Map<String, String> propMap = new HashMap<>();
060
061    @Config(description="Hostname of the database machine.")
062    private String host;
063    @Config(description="Port number.")
064    private String port;
065    @Config(description="Database name.")
066    private String db;
067
068    @Config(description="Size of batches to fetch from DB for queries")
069    private int fetchSize = 1000;
070
071    private SQLDBConfig() {}
072
073    public SQLDBConfig(String connectionString, String username, String password, Map<String, String> properties) {
074        this(connectionString, properties);
075        this.username = username;
076        this.password = password;
077    }
078
079    public SQLDBConfig(String host, String port, String db, String username, String password, Map<String, String> properties) {
080        this(makeConnectionString(host, port, db), properties);
081        this.host = host;
082        this.port = port;
083        this.db = db;
084        this.username = username;
085        this.password = password;
086    }
087
088    public SQLDBConfig(String connectionString, Map<String, String> properties) {
089        this.connectionString = connectionString;
090        this.propMap = properties;
091    }
092
093    private static String makeConnectionString(String host, String port, String db) {
094        return "jdbc:oracle:thin:@" + host + ":" + port + "/" + db;
095    }
096
097    /**
098     * Used by the OLCUT configuration system, and should not be called by external code.
099     */
100    @Override
101    public void postConfig() {
102
103        if(propMap.containsKey("user")) {
104            if(username == null) {
105                username = propMap.remove("user");
106            } else {
107                propMap.remove("user");
108            }
109        }
110        if(propMap.containsKey("password")) {
111            if(password == null) {
112                password = propMap.remove("password");
113            } else {
114                propMap.remove("password");
115            }
116        }
117        if(connectionString == null) {
118            if(host != null && port != null && db != null) {
119                connectionString = makeConnectionString(host, port, db);
120            } else {
121                throw new PropertyException(SQLDBConfig.class.getName(), "connectionString", "All of host, port, and db must be specified if connectionString is null");
122            }
123        }
124    }
125
126    /**
127     * Constructs a connection based on the object fields.
128     * @return A connection to the database.
129     * @throws SQLException If the connection failed.
130     */
131    public Connection getConnection() throws SQLException {
132
133        Properties props = new Properties();
134        props.putAll(propMap);
135        if(username != null && password != null) {
136            props.put("user", username);
137            props.put("password", password);
138        }
139        return DriverManager.getConnection(connectionString, props);
140    }
141
142    /**
143     * Constructs a statement based on the object fields. Uses fetchSize to determine fetch size and sets defaults
144     * for querying data.
145     *
146     * @return A statement object for querying the database.
147     * @throws SQLException If the connection failed.
148     */
149    public Statement getStatement() throws SQLException {
150        Statement stmt = getConnection().createStatement();
151        stmt.setFetchSize(fetchSize);
152        stmt.setFetchDirection(ResultSet.FETCH_FORWARD);
153        return stmt;
154    }
155
156    @Override
157    public String toString() {
158        if (connectionString != null) {
159            return "SQLDBConfig(connectionString="+connectionString+")";
160        } else {
161            return "SQLDBConfig(host="+host+",port="+port+",db="+db+")";
162        }
163    }
164
165    @Override
166    public ConfiguredObjectProvenance getProvenance() {
167        return new ConfiguredObjectProvenanceImpl(this,"SQL-DB-Config");
168    }
169}