简柏林 / thingskit · Commits

Commit 99444395393f0f2966e94de9a67017ee6be742f9

Authored by Igor Kulikov
1 parent 063a1093

Improve CSV data dump.

Inline Side-by-side
Showing 4 changed files with 60 additions and 19 deletions
... ... @@ -169,7 +169,7 @@ public class CassandraDatabaseUpgradeService implements DatabaseUpgradeService {
169 169 Path dashboardsDump = CassandraDbHelper.dumpCfIfExists(ks, cluster.getSession(), DASHBOARD,
170 170 new String[]{ID, TENANT_ID, CUSTOMER_ID, TITLE, SEARCH_TEXT, ASSIGNED_CUSTOMERS, CONFIGURATION},
171 171 new String[]{"", "", "", "", "", "", ""},
172   - "tb-dashboards");
  172 + "tb-dashboards", true);
173 173 log.info("Dashboards dumped.");
174 174
175 175
... ... @@ -181,7 +181,7 @@ public class CassandraDatabaseUpgradeService implements DatabaseUpgradeService {
181 181 log.info("Restoring dashboards ...");
182 182 if (dashboardsDump != null) {
183 183 CassandraDbHelper.loadCf(ks, cluster.getSession(), DASHBOARD,
184   - new String[]{ID, TENANT_ID, TITLE, SEARCH_TEXT, CONFIGURATION}, dashboardsDump);
  184 + new String[]{ID, TENANT_ID, TITLE, SEARCH_TEXT, CONFIGURATION}, dashboardsDump, true);
185 185 DatabaseHelper.upgradeTo40_assignDashboards(dashboardsDump, dashboardService, false);
186 186 Files.deleteIfExists(dashboardsDump);
187 187 }
... ...
... ... @@ -79,7 +79,7 @@ public class SqlDatabaseUpgradeService implements DatabaseUpgradeService {
79 79 Path dashboardsDump = SqlDbHelper.dumpTableIfExists(conn, DASHBOARD,
80 80 new String[]{ID, TENANT_ID, CUSTOMER_ID, TITLE, SEARCH_TEXT, ASSIGNED_CUSTOMERS, CONFIGURATION},
81 81 new String[]{"", "", "", "", "", "", ""},
82   - "tb-dashboards");
  82 + "tb-dashboards", true);
83 83 log.info("Dashboards dumped.");
84 84
85 85 log.info("Updating schema ...");
... ... @@ -91,7 +91,7 @@ public class SqlDatabaseUpgradeService implements DatabaseUpgradeService {
91 91 log.info("Restoring dashboards ...");
92 92 if (dashboardsDump != null) {
93 93 SqlDbHelper.loadTable(conn, DASHBOARD,
94   - new String[]{ID, TENANT_ID, TITLE, SEARCH_TEXT, CONFIGURATION}, dashboardsDump);
  94 + new String[]{ID, TENANT_ID, TITLE, SEARCH_TEXT, CONFIGURATION}, dashboardsDump, true);
95 95 DatabaseHelper.upgradeTo40_assignDashboards(dashboardsDump, dashboardService, true);
96 96 Files.deleteIfExists(dashboardsDump);
97 97 }
... ...
... ... @@ -17,6 +17,7 @@
17 17 package org.thingsboard.server.service.install.cql;
18 18
19 19 import com.datastax.driver.core.*;
  20 +import org.apache.commons.csv.CSVFormat;
20 21 import org.apache.commons.csv.CSVParser;
21 22 import org.apache.commons.csv.CSVPrinter;
22 23 import org.apache.commons.csv.CSVRecord;
... ... @@ -33,10 +34,19 @@ public class CassandraDbHelper {
33 34
34 35 public static Path dumpCfIfExists(KeyspaceMetadata ks, Session session, String cfName,
35 36 String[] columns, String[] defaultValues, String dumpPrefix) throws Exception {
  37 + return dumpCfIfExists(ks, session, cfName, columns, defaultValues, dumpPrefix, false);
  38 + }
  39 +
  40 + public static Path dumpCfIfExists(KeyspaceMetadata ks, Session session, String cfName,
  41 + String[] columns, String[] defaultValues, String dumpPrefix, boolean printHeader) throws Exception {
36 42 if (ks.getTable(cfName) != null) {
37 43 Path dumpFile = Files.createTempFile(dumpPrefix, null);
38 44 Files.deleteIfExists(dumpFile);
39   - try (CSVPrinter csvPrinter = new CSVPrinter(Files.newBufferedWriter(dumpFile), CSV_DUMP_FORMAT)) {
  45 + CSVFormat csvFormat = CSV_DUMP_FORMAT;
  46 + if (printHeader) {
  47 + csvFormat = csvFormat.withHeader(columns);
  48 + }
  49 + try (CSVPrinter csvPrinter = new CSVPrinter(Files.newBufferedWriter(dumpFile), csvFormat)) {
40 50 Statement stmt = new SimpleStatement("SELECT * FROM " + cfName);
41 51 stmt.setFetchSize(1000);
42 52 ResultSet rs = session.execute(stmt);
... ... @@ -74,9 +84,19 @@ public class CassandraDbHelper {
74 84 }
75 85
76 86 public static void loadCf(KeyspaceMetadata ks, Session session, String cfName, String[] columns, Path sourceFile) throws Exception {
  87 + loadCf(ks, session, cfName, columns, sourceFile, false);
  88 + }
  89 +
  90 + public static void loadCf(KeyspaceMetadata ks, Session session, String cfName, String[] columns, Path sourceFile, boolean parseHeader) throws Exception {
77 91 TableMetadata tableMetadata = ks.getTable(cfName);
78 92 PreparedStatement prepared = session.prepare(createInsertStatement(cfName, columns));
79   - try (CSVParser csvParser = new CSVParser(Files.newBufferedReader(sourceFile), CSV_DUMP_FORMAT.withHeader(columns))) {
  93 + CSVFormat csvFormat = CSV_DUMP_FORMAT;
  94 + if (parseHeader) {
  95 + csvFormat = csvFormat.withFirstRecordAsHeader();
  96 + } else {
  97 + csvFormat = CSV_DUMP_FORMAT.withHeader(columns);
  98 + }
  99 + try (CSVParser csvParser = new CSVParser(Files.newBufferedReader(sourceFile), csvFormat)) {
80 100 csvParser.forEach(record -> {
81 101 BoundStatement boundStatement = prepared.bind();
82 102 for (String column : columns) {
... ...
... ... @@ -16,6 +16,7 @@
16 16 package org.thingsboard.server.service.install.sql;
17 17
18 18 import lombok.extern.slf4j.Slf4j;
  19 +import org.apache.commons.csv.CSVFormat;
19 20 import org.apache.commons.csv.CSVParser;
20 21 import org.apache.commons.csv.CSVPrinter;
21 22 import org.apache.commons.csv.CSVRecord;
... ... @@ -38,6 +39,11 @@ public class SqlDbHelper {
38 39
39 40 public static Path dumpTableIfExists(Connection conn, String tableName,
40 41 String[] columns, String[] defaultValues, String dumpPrefix) throws Exception {
  42 + return dumpTableIfExists(conn, tableName, columns, defaultValues, dumpPrefix, false);
  43 + }
  44 +
  45 + public static Path dumpTableIfExists(Connection conn, String tableName,
  46 + String[] columns, String[] defaultValues, String dumpPrefix, boolean printHeader) throws Exception {
41 47
42 48 DatabaseMetaData metaData = conn.getMetaData();
43 49 ResultSet res = metaData.getTables(null, null, tableName,
... ... @@ -46,7 +52,11 @@ public class SqlDbHelper {
46 52 res.close();
47 53 Path dumpFile = Files.createTempFile(dumpPrefix, null);
48 54 Files.deleteIfExists(dumpFile);
49   - try (CSVPrinter csvPrinter = new CSVPrinter(Files.newBufferedWriter(dumpFile), CSV_DUMP_FORMAT)) {
  55 + CSVFormat csvFormat = CSV_DUMP_FORMAT;
  56 + if (printHeader) {
  57 + csvFormat = csvFormat.withHeader(columns);
  58 + }
  59 + try (CSVPrinter csvPrinter = new CSVPrinter(Files.newBufferedWriter(dumpFile), csvFormat)) {
50 60 try (PreparedStatement stmt = conn.prepareStatement("SELECT * FROM " + tableName)) {
51 61 try (ResultSet tableRes = stmt.executeQuery()) {
52 62 ResultSetMetaData resMetaData = tableRes.getMetaData();
... ... @@ -68,19 +78,30 @@ public class SqlDbHelper {
68 78 }
69 79
70 80 public static void loadTable(Connection conn, String tableName, String[] columns, Path sourceFile) throws Exception {
71   - PreparedStatement prepared = conn.prepareStatement(createInsertStatement(tableName, columns));
72   - prepared.getParameterMetaData();
73   - try (CSVParser csvParser = new CSVParser(Files.newBufferedReader(sourceFile), CSV_DUMP_FORMAT.withHeader(columns))) {
74   - csvParser.forEach(record -> {
75   - try {
76   - for (int i=0;i<columns.length;i++) {
77   - setColumnValue(i, columns[i], record, prepared);
  81 + loadTable(conn, tableName, columns, sourceFile, false);
  82 + }
  83 +
  84 + public static void loadTable(Connection conn, String tableName, String[] columns, Path sourceFile, boolean parseHeader) throws Exception {
  85 + CSVFormat csvFormat = CSV_DUMP_FORMAT;
  86 + if (parseHeader) {
  87 + csvFormat = csvFormat.withFirstRecordAsHeader();
  88 + } else {
  89 + csvFormat = CSV_DUMP_FORMAT.withHeader(columns);
  90 + }
  91 + try (PreparedStatement prepared = conn.prepareStatement(createInsertStatement(tableName, columns))) {
  92 + prepared.getParameterMetaData();
  93 + try (CSVParser csvParser = new CSVParser(Files.newBufferedReader(sourceFile), csvFormat)) {
  94 + csvParser.forEach(record -> {
  95 + try {
  96 + for (int i = 0; i < columns.length; i++) {
  97 + setColumnValue(i, columns[i], record, prepared);
  98 + }
  99 + prepared.execute();
  100 + } catch (SQLException e) {
  101 + log.error("Unable to load table record!", e);
78 102 }
79   - prepared.execute();
80   - } catch (SQLException e) {
81   - log.error("Unable to load table record!", e);
82   - }
83   - });
  103 + });
  104 + }
84 105 }
85 106 }
86 107
... ...