com.datastax.data.dataset.util.DataSetUtility Maven / Gradle / Ivy

Go to download

package com.datastax.data.dataset.util;

import com.datastax.data.dataset.*;
import com.datastax.data.dataset.io.schema.*;
import com.datastax.data.dataset.provider.sql.JDBCDataConnection;
import com.datastax.data.dataset.provider.sql.SQLDataProvider;

import java.io.InputStream;
import java.io.OutputStream;
import java.io.Reader;
import java.io.Writer;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.util.*;


public class DataSetUtility {

    
    public static DataSet setProviderConnection(DataSet ds,
            JDBCDataConnection conn) {
        for (DataTable table : ds.getTables()) {
            if (!(table instanceof DataRelationTable)
                    && table.getDataProvider() != null) {
                table.getDataProvider().setConnection(conn);
            }
        }
        return ds;
    }
    
    
    public static DataSet readDataSet(DataSetSchemaReader reader) throws SchemaReaderException {
       return reader.readDataSet();
    }
    
    
    public static DataSet readDataSet(DataSetSchemaReader reader, String... tableNames) throws SchemaReaderException {
       return reader.readDataSet(tableNames);
    }

    
    public static DataSet readDataSetFromJDBC(JDBCDataConnection jdbcConn) throws SchemaReaderException {
         return readDataSet(new JDBCDataSetSchemaReader(jdbcConn));
    }

    
    public static DataSet readDataSetFromJDBC(JDBCDataConnection jdbcConn, String... tableNames) throws SchemaReaderException {
         return readDataSet(new JDBCDataSetSchemaReader(jdbcConn), tableNames);       
    }

    
    public static DataSet readDataSetFromXml(InputStream is, String... tableNames) throws SchemaReaderException {
       if ( tableNames.length == 0 ) {
         return readDataSet(new XMLDataSetSchemaReader(is));
       } else {
         return readDataSet(new XMLDataSetSchemaReader(is), tableNames);
       }
    }

    
    public static DataSet readDataSetFromXml(Reader reader, String... tableNames) throws SchemaReaderException {
       if ( tableNames.length == 0 ) {
         return readDataSet(new XMLDataSetSchemaReader(reader));
       } else {
         return readDataSet(new XMLDataSetSchemaReader(reader), tableNames);
       }
    }

    
    public static void writeDataSet(DataSetSchemaWriter writer, DataSet ds) throws SchemaWriterException {
      writer.writeDataSet(ds);
    }

    
    public static void writeDataSet(DataSetSchemaWriter writer, DataSet ds, String... tableNames) throws SchemaWriterException {
      writer.writeDataSet(ds, tableNames);
    }

    
    public static void writeDataSetAsXml(OutputStream out, DataSet ds) throws SchemaWriterException {
      XMLDataSetSchemaWriter writer = new XMLDataSetSchemaWriter(out);
      writer.writeDataSet(ds);
    }
    

    
    public static void writeDataSetAsXml(Writer writer, DataSet ds) throws SchemaWriterException {
      XMLDataSetSchemaWriter xwriter = new XMLDataSetSchemaWriter(writer);
      xwriter.writeDataSet(ds);
    }    

    
    public static void createDataTableSchema(JDBCDataConnection jdbcConn,
                                             DataSet ds, String tableName, String statement)
            throws SchemaCreationException {

        assert jdbcConn != null : "Cannot supply a null data connection";
        assert ds != null : "Cannot supply a null DataSource";
        assert tableName != null : "Cannot supply a null table name";
        assert statement != null : "Cannot supply a null statement";
        assert jdbcConn.isConnected() : "The data connection must be open";

        System.out.println(statement);
        DataTable dt = ds.createTable(tableName);

        
        try {
            Connection conn = jdbcConn.getConnection();
            ResultSet set = conn.createStatement().executeQuery(statement);
            ResultSetMetaData rsmd = set.getMetaData();
            DatabaseMetaData md = conn.getMetaData();

            
            
            if (rsmd.getColumnCount() == 0) {
                return;
            }

            String catalogName = rsmd.getCatalogName(1);

            
            
            Set tableNames = new HashSet(5);

            for (int i = 1; i <= rsmd.getColumnCount(); i++) {
                
                
                
                String colCat = rsmd.getCatalogName(i);
                if (!"".equals(colCat) && !catalogName.equals(colCat)) {
                    throw new SchemaCreationException(
                            "Resulting table has source from multiple catalogs - "
                                    + catalogName + " & " + colCat);
                }

                String t = rsmd.getTableName(i);

                
                
                DataColumn col = dt.createColumn(rsmd.getColumnName(i));
                col.setReadOnly(rsmd.isReadOnly(i));
                col.setType(DataSetIOUtility.getType(rsmd.getColumnType(i)));

                ResultSet rs = md.getColumns(catalogName, null, t, col
                        .getName());
                
                
                if (!rs.next()) {
                    continue;
                }

                col.setDefaultValue(rs.getObject(13));
                col.setRequired(rs.getString(18).equals("NO"));

                tableNames.add(t);
            }

            
            
            
            for (String table : tableNames) {
                ResultSet rs = md.getPrimaryKeys(catalogName, null, table);
                while (rs.next()) {
                    DataColumn col = dt.getColumn(rs.getString(4));
                    
                    if (col != null) {
                        col.setKeyColumn(true);
                    }
                }
            }
        } catch (Exception e) {
            
            ds.dropTable(dt);
            e.printStackTrace();
            throw new SchemaCreationException(
                    "Could not create the schema for the table " + tableName, e);
        }
    }

    
    public static DataSet createDataSetSchema(JDBCDataConnection conn,
            String... names) throws SchemaCreationException {

        assert conn != null : "Cannot supply a null data connection";
        assert names != null : "Cannot supply a null list of names";
        assert conn.isConnected() : "The data connection must be open";

        DataSet ds = new DataSet();

        
        
        
        Map> tables = new HashMap>(5);

        for (String name : names) {
            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 tbl : tables.keySet()) {
            Set cols = tables.get(tbl);

            
            
            if (cols == null) {
                createDataTableSchema(conn, ds, tbl, "SELECT * FROM " + tbl);
                        
            } else {
                
                
                assert cols.size() > 0 : "Created an empty column set";

                StringBuilder cmd = new StringBuilder("SELECT ");

                for (String name : cols) {
                    cmd.append(name);
                    cmd.append(",");
                }

                
                cmd.deleteCharAt(cmd.length() - 1);
                cmd.append(" FROM ");
                cmd.append(tbl);
                cmd.append(" LIMIT 0");

                createDataTableSchema(conn, ds, tbl, cmd.toString());
            }

            SQLDataProvider dataProvider = new SQLDataProvider(tbl);
            dataProvider.setConnection(conn);
            ds.getTable(tbl).setDataProvider(dataProvider);
        }

        createRelations(conn, ds);

        return ds;
    }

    
    public static void createRelations(JDBCDataConnection jdbcConn, DataSet ds)
            throws SchemaCreationException {

        assert jdbcConn != null : "Cannot supply a null connection";
        assert ds != null : "Cannot supply a null DataSource";
        assert jdbcConn.isConnected() : "The data connection must be open";

        
        
        
        List dataTables = ds.getTables();

        try {
            Connection conn = jdbcConn.getConnection();

            for (DataTable tbl : dataTables) {
                DatabaseMetaData md = conn.getMetaData();
                
                ResultSet rs = md.getImportedKeys(conn.getCatalog(), null, tbl
                        .getName());

                while (rs.next()) {
                    
                    
                    if (rs.getInt(9) > 1) {
                        continue;
                    }

                    String childTableString = rs.getString(3);
                    DataTable childTable = ds.getTable(rs.getString(3));
                    for (DataTable t : dataTables) {
                        if (t.getName().equalsIgnoreCase(childTableString)) {
                            childTable = t;
                            break;
                        }
                    }

                    
                    if (childTable == null) {
                        continue;
                    }

                    DataColumn childColumn = childTable.getColumn(rs
                            .getString(4));
                    DataColumn parentColumn = tbl.getColumn(rs.getString(8));

                    if (parentColumn != null && childColumn != null
                            && parentColumn != childColumn) {
                        DataRelation rel = ds.createRelation(rs.getString(12));
                        rel.setParentColumn(parentColumn);
                        rel.setChildColumn(childColumn);
                    }
                }
            }
        } catch (Exception e) {
            throw new SchemaCreationException(
                    "Could not create data relations", e);
        }
    }

    public static void main(String[] args) throws Exception {
        JDBCDataConnection jdbcConn = new JDBCDataConnection();

        jdbcConn.setConnected(true);
        DataSet ds = createDataSetSchema(jdbcConn, "Addresses", "Customers",
                "Bookings", "BookingLocations");
        ds.loadAndWait();
        jdbcConn.setConnected(false);

        System.out.println(DataSetUtils.getXmlSchema(ds));
        System.out.println(ds.writeXml());
    }
}