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 com.rthoni.intellij.codefromds.business.Helper;
import org.json.JSONArray;
import org.json.JSONObject;

import java.util.HashMap;
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(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 HashMap<String, Object> toMap()
    {
        HashMap<String, Object> map = new HashMap<>();

        map.put("source", _source == null ? null : _source.getName());
        List<Object> tables = new Vector<>();
        for (TableSelection table : _tables) {
            tables.add(table.toMap());
        }
        map.put("tables", tables);

        return map;
    }

    public void fromJson(JSONObject json)
    {
        JSONArray array = json.getJSONArray("tables");
        for (TableSelection table : _tables) {
            JSONObject obj = Helper.findTableInJson(array, table.getTable().getName());
            if (obj != null) {
                table.fromJson(obj);
            } else {
                for (ColumnSelection column : table.getColumns()) {
                    column.setSelected(false);
                }
            }
        }
    }

    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;
    }
}