com.datastax.data.dataset.io.schema.JDBCDataSetSchemaReader Maven / Gradle / Ivy
The newest version!
package com.datastax.data.dataset.io.schema;
import com.datastax.data.dataset.DataColumn;
import com.datastax.data.dataset.DataRelation;
import com.datastax.data.dataset.DataSet;
import com.datastax.data.dataset.DataTable;
import com.datastax.data.dataset.provider.sql.JDBCDataConnection;
import com.datastax.data.dataset.provider.sql.SQLDataProvider;
import java.sql.*;
import java.util.*;
public class JDBCDataSetSchemaReader implements DataSetSchemaReader {
private JDBCDataConnection jdbcConn;
private String catalog;
private String schema;
public JDBCDataSetSchemaReader(JDBCDataConnection jdbcConn) {
this(jdbcConn, null, null);
}
public JDBCDataSetSchemaReader(JDBCDataConnection jdbcConn, String catalog, String schema) {
this.jdbcConn = jdbcConn;
this.catalog = catalog;
this.schema = schema;
}
public DataSet readDataSet() throws SchemaReaderException {
boolean connectState = jdbcConn.isConnected();
if ( !connectState ) jdbcConn.setConnected(true);
DataSet dataSet = null;
try {
Connection conn = jdbcConn.getConnection();
DatabaseMetaData dbMetaData = conn.getMetaData();
String[] typeArray = {"TABLE"};
ResultSet tableResultSet = dbMetaData.getTables(null, null, "%", typeArray);
List tables = DataSetIOUtility.extractColumn(tableResultSet, "TABLE_NAME");
tableResultSet.close();
String[] array = (String[])tables.toArray(new String[tables.size()]);
dataSet = readDataSet(array);
} catch (SchemaReaderException e) {
throw e;
} catch (Exception e) {
throw new SchemaReaderException("Could not read schemas to create DataSource", e);
} finally {
jdbcConn.setConnected(connectState);
}
return dataSet;
}
public DataSet readDataSet(String... tableNames) throws SchemaReaderException {
boolean connectState = jdbcConn.isConnected();
if ( !connectState ) jdbcConn.setConnected(true);
DataSet dataSet = new DataSet();
try {
addTables(dataSet, tableNames);
addRelations(dataSet, tableNames);
} catch (SchemaReaderException e) {
throw e;
} catch (Exception e) {
throw new SchemaReaderException("Could not read schemas to create DataTables", e);
} finally {
jdbcConn.setConnected(connectState);
}
return dataSet;
}
public List addTables(DataSet dataSet, String... tableNames) throws SchemaReaderException {
List loadedTables = new ArrayList(tableNames.length);
boolean connectState = jdbcConn.isConnected();
if ( !connectState ) jdbcConn.setConnected(true);
try {
Connection conn = jdbcConn.getConnection();
DatabaseMetaData dbMetaData = conn.getMetaData();
Map> tables = new HashMap>(5);
for (String name : tableNames) {
int dotIdx = name.indexOf(".");
if (dotIdx != -1) {
String table = name.substring(0, dotIdx);
String col = name.substring(dotIdx + 1);
if (!tables.containsKey(table)) {
tables.put(table, new HashSet());
}
if (tables.get(table) != null) {
tables.get(table).add(col);
}
} else {
tables.put(name, null);
}
}
for (String tableName : tables.keySet()) {
if ( dataSet.getTable(tableName) != null ) {
System.out.println("SKIPPING TABLE " + tableName + " ALREADY IN DATASET");
continue;
}
Set cols = tables.get(tableName);
List colNames = ( cols == null ? null : new ArrayList(cols));
String tableSelectCmd = getTableSelectNoRowsCommand(tableName, colNames);
DataTable dataTable = readDataTable(dataSet, tableName, tableSelectCmd);
SQLDataProvider dataProvider = new SQLDataProvider(tableName);
dataProvider.setConnection(jdbcConn);
dataTable.setDataProvider(dataProvider);
loadedTables.add(tableName);
}
} catch (SchemaReaderException e) {
throw e;
} catch (Exception e) {
throw new SchemaReaderException("Could not read schemas to create DataTables", e);
} finally {
jdbcConn.setConnected(connectState);
}
return loadedTables;
}
public List addRelations(DataSet dataSet, String... tableNames) throws SchemaReaderException {
List loadedRelations = new ArrayList();
boolean connectState = jdbcConn.isConnected();
if ( !connectState ) jdbcConn.setConnected(true);
Connection conn = jdbcConn.getConnection();
try {
List dataTables = dataSet.getTables();
for ( String tableName : tableNames) {
DataTable dataTable = dataSet.getTable(tableName);
DatabaseMetaData md = conn.getMetaData();
ResultSet rs = md.getImportedKeys(conn.getCatalog(), null, dataTable.getName());
while (rs.next()) {
if (rs.getInt("KEY_SEQ") > 1) {
continue;
}
DataTable parentTable = null;
String pkTableName = rs.getString("PKTABLE_NAME");
for (DataTable t : dataTables) {
if (t.getName().equalsIgnoreCase(pkTableName)) {
parentTable = t;
break;
}
}
if (parentTable == null) {
continue;
}
DataColumn parentColumn = parentTable.getColumn(rs.getString("PKCOLUMN_NAME"));
DataColumn childColumn = dataTable.getColumn(rs.getString("FKCOLUMN_NAME"));
if (parentColumn != null && childColumn != null && parentColumn != childColumn) {
String name = rs.getString("FK_NAME");
if ( dataSet.getRelation(name) != null ) dataSet.dropRelation(name);
DataRelation rel = dataSet.createRelation(rs.getString("FK_NAME"));
rel.setParentColumn(parentColumn);
rel.setChildColumn(childColumn);
loadedRelations.add(rel.getName());
}
}
}
} catch (Exception e) {
throw new SchemaReaderException("Could not create data relations", e);
} finally {
jdbcConn.setConnected(connectState);
}
return loadedRelations;
}
protected String getTableSelectNoRowsCommand(String tableName, List columnNames) {
String comma = ", \n";
StringBuilder builder = new StringBuilder();
builder.append("SELECT ");
if ( columnNames == null ) {
builder.append("*\n");
} else {
for ( String col : columnNames ) {
builder.append(col + comma);
}
builder.delete(builder.length() - comma.length(), builder.length() - 1);
}
builder.append("FROM " + tableName + "\n");
builder.append("WHERE 1 = 0");
return builder.toString();
}
private DataTable readDataTable(DataSet dataSet, String tableName, String tableSelectCmd) throws SchemaReaderException {
boolean connectState = jdbcConn.isConnected();
if ( !connectState ) jdbcConn.setConnected(true);
Connection conn = jdbcConn.getConnection();
DataTable dataTable = null;
try {
DatabaseMetaData dbMetaData = conn.getMetaData();
Statement stmt = conn.createStatement();
ResultSetMetaData colRSMD = stmt.executeQuery(tableSelectCmd).getMetaData();
dataTable = dataSet.createTable(tableName);
int colCnt = colRSMD.getColumnCount();
String catalog = colRSMD.getCatalogName(1);
String schema = colRSMD.getSchemaName(1);;
for ( int i=1; i <= colCnt; i++ ) {
String colName = colRSMD.getColumnName(i);
DataColumn newColumn = dataTable.createColumn(colName);
newColumn.setRequired(colRSMD.isNullable(i) == ResultSetMetaData.columnNullable);
newColumn.setReadOnly(colRSMD.isReadOnly(i));
newColumn.setType(DataSetIOUtility.getType(colRSMD.getColumnType(i)));
String columnPattern = "%";
ResultSet colRS = dbMetaData.getColumns(catalog, schema, tableName, columnPattern );
if ( !colRS.next()) continue;
newColumn.setDefaultValue(colRS.getObject("COLUMN_DEF"));
colRS.close();
}
ResultSet rs = dbMetaData.getPrimaryKeys(catalog, schema, tableName);
while (rs.next()) {
DataColumn col = dataTable.getColumn(rs.getString("COLUMN_NAME"));
if (col != null) {
col.setKeyColumn(true);
}
}
SQLDataProvider dataProvider = new SQLDataProvider(tableName);
dataProvider.setConnection(jdbcConn);
dataTable.setDataProvider(dataProvider);
} catch (Exception e) {
dataSet.dropTable(tableName);
throw new SchemaReaderException("Could not create DataTable for table named " + tableName, e);
} finally {
jdbcConn.setConnected(connectState);
}
return dataTable;
}
}