intellij-code-from-ds/src/com/rthoni/intellij/codefromds/dbo/DataSourceSelection.java

package com.rthoni.intellij.codefromds.dbo;

import com.intellij.database.model.DasTable;
import com.intellij.database.psi.DbDataSource;
import com.intellij.database.util.DasUtil;

import java.util.List;
import java.util.Vector;
import java.util.stream.Collectors;

/**
 * Created by robin on 11/15/16.
 */
public class DataSourceSelection {

    public DataSourceSelection() {
    }

    public DataSourceSelection(DbDataSource source) {
        _source = source;
        final List<? extends DasTable> tables = DasUtil.getTables(source).toList();
        _tables = new Vector<>();
        for (DasTable table : tables) {
            TableSelection tableSelection = new TableSelection();
            tableSelection.setTable(table);
            tableSelection.setColumns(DasUtil.getColumns(table).toList().stream().map(ColumnSelection::new).collect(Collectors.toList()));
            _tables.add(tableSelection);
        }
    }

    private DbDataSource _source;

    private List<TableSelection> _tables;

    public DbDataSource getSource() {
        return _source;
    }

    public void setSource(DbDataSource source) {
        _source = source;
    }

    public List<TableSelection> getTables() {
        return _tables;
    }

    public void setTables(List<TableSelection> tables) {
        _tables = tables;
    }
}