简柏林 / 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,7 +169,7 @@ public class CassandraDatabaseUpgradeService implements DatabaseUpgradeService {
169 Path dashboardsDump = CassandraDbHelper.dumpCfIfExists(ks, cluster.getSession(), DASHBOARD, 169 Path dashboardsDump = CassandraDbHelper.dumpCfIfExists(ks, cluster.getSession(), DASHBOARD,
170 new String[]{ID, TENANT_ID, CUSTOMER_ID, TITLE, SEARCH_TEXT, ASSIGNED_CUSTOMERS, CONFIGURATION}, 170 new String[]{ID, TENANT_ID, CUSTOMER_ID, TITLE, SEARCH_TEXT, ASSIGNED_CUSTOMERS, CONFIGURATION},
171 new String[]{"", "", "", "", "", "", ""}, 171 new String[]{"", "", "", "", "", "", ""},
172 - "tb-dashboards"); 172 + "tb-dashboards", true);
173 log.info("Dashboards dumped."); 173 log.info("Dashboards dumped.");
174 174
175 175
@@ -181,7 +181,7 @@ public class CassandraDatabaseUpgradeService implements DatabaseUpgradeService { @@ -181,7 +181,7 @@ public class CassandraDatabaseUpgradeService implements DatabaseUpgradeService {
181 log.info("Restoring dashboards ..."); 181 log.info("Restoring dashboards ...");
182 if (dashboardsDump != null) { 182 if (dashboardsDump != null) {
183 CassandraDbHelper.loadCf(ks, cluster.getSession(), DASHBOARD, 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 DatabaseHelper.upgradeTo40_assignDashboards(dashboardsDump, dashboardService, false); 185 DatabaseHelper.upgradeTo40_assignDashboards(dashboardsDump, dashboardService, false);
186 Files.deleteIfExists(dashboardsDump); 186 Files.deleteIfExists(dashboardsDump);
187 } 187 }
@@ -79,7 +79,7 @@ public class SqlDatabaseUpgradeService implements DatabaseUpgradeService { @@ -79,7 +79,7 @@ public class SqlDatabaseUpgradeService implements DatabaseUpgradeService {
79 Path dashboardsDump = SqlDbHelper.dumpTableIfExists(conn, DASHBOARD, 79 Path dashboardsDump = SqlDbHelper.dumpTableIfExists(conn, DASHBOARD,
80 new String[]{ID, TENANT_ID, CUSTOMER_ID, TITLE, SEARCH_TEXT, ASSIGNED_CUSTOMERS, CONFIGURATION}, 80 new String[]{ID, TENANT_ID, CUSTOMER_ID, TITLE, SEARCH_TEXT, ASSIGNED_CUSTOMERS, CONFIGURATION},
81 new String[]{"", "", "", "", "", "", ""}, 81 new String[]{"", "", "", "", "", "", ""},
82 - "tb-dashboards"); 82 + "tb-dashboards", true);
83 log.info("Dashboards dumped."); 83 log.info("Dashboards dumped.");
84 84
85 log.info("Updating schema ..."); 85 log.info("Updating schema ...");
@@ -91,7 +91,7 @@ public class SqlDatabaseUpgradeService implements DatabaseUpgradeService { @@ -91,7 +91,7 @@ public class SqlDatabaseUpgradeService implements DatabaseUpgradeService {
91 log.info("Restoring dashboards ..."); 91 log.info("Restoring dashboards ...");
92 if (dashboardsDump != null) { 92 if (dashboardsDump != null) {
93 SqlDbHelper.loadTable(conn, DASHBOARD, 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 DatabaseHelper.upgradeTo40_assignDashboards(dashboardsDump, dashboardService, true); 95 DatabaseHelper.upgradeTo40_assignDashboards(dashboardsDump, dashboardService, true);
96 Files.deleteIfExists(dashboardsDump); 96 Files.deleteIfExists(dashboardsDump);
97 } 97 }
@@ -17,6 +17,7 @@ @@ -17,6 +17,7 @@
17 package org.thingsboard.server.service.install.cql; 17 package org.thingsboard.server.service.install.cql;
18 18
19 import com.datastax.driver.core.*; 19 import com.datastax.driver.core.*;
  20 +import org.apache.commons.csv.CSVFormat;
20 import org.apache.commons.csv.CSVParser; 21 import org.apache.commons.csv.CSVParser;
21 import org.apache.commons.csv.CSVPrinter; 22 import org.apache.commons.csv.CSVPrinter;
22 import org.apache.commons.csv.CSVRecord; 23 import org.apache.commons.csv.CSVRecord;
@@ -33,10 +34,19 @@ public class CassandraDbHelper { @@ -33,10 +34,19 @@ public class CassandraDbHelper {
33 34
34 public static Path dumpCfIfExists(KeyspaceMetadata ks, Session session, String cfName, 35 public static Path dumpCfIfExists(KeyspaceMetadata ks, Session session, String cfName,
35 String[] columns, String[] defaultValues, String dumpPrefix) throws Exception { 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 if (ks.getTable(cfName) != null) { 42 if (ks.getTable(cfName) != null) {
37 Path dumpFile = Files.createTempFile(dumpPrefix, null); 43 Path dumpFile = Files.createTempFile(dumpPrefix, null);
38 Files.deleteIfExists(dumpFile); 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 Statement stmt = new SimpleStatement("SELECT * FROM " + cfName); 50 Statement stmt = new SimpleStatement("SELECT * FROM " + cfName);
41 stmt.setFetchSize(1000); 51 stmt.setFetchSize(1000);
42 ResultSet rs = session.execute(stmt); 52 ResultSet rs = session.execute(stmt);
@@ -74,9 +84,19 @@ public class CassandraDbHelper { @@ -74,9 +84,19 @@ public class CassandraDbHelper {
74 } 84 }
75 85
76 public static void loadCf(KeyspaceMetadata ks, Session session, String cfName, String[] columns, Path sourceFile) throws Exception { 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 TableMetadata tableMetadata = ks.getTable(cfName); 91 TableMetadata tableMetadata = ks.getTable(cfName);
78 PreparedStatement prepared = session.prepare(createInsertStatement(cfName, columns)); 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 csvParser.forEach(record -> { 100 csvParser.forEach(record -> {
81 BoundStatement boundStatement = prepared.bind(); 101 BoundStatement boundStatement = prepared.bind();
82 for (String column : columns) { 102 for (String column : columns) {
@@ -16,6 +16,7 @@ @@ -16,6 +16,7 @@
16 package org.thingsboard.server.service.install.sql; 16 package org.thingsboard.server.service.install.sql;
17 17
18 import lombok.extern.slf4j.Slf4j; 18 import lombok.extern.slf4j.Slf4j;
  19 +import org.apache.commons.csv.CSVFormat;
19 import org.apache.commons.csv.CSVParser; 20 import org.apache.commons.csv.CSVParser;
20 import org.apache.commons.csv.CSVPrinter; 21 import org.apache.commons.csv.CSVPrinter;
21 import org.apache.commons.csv.CSVRecord; 22 import org.apache.commons.csv.CSVRecord;
@@ -38,6 +39,11 @@ public class SqlDbHelper { @@ -38,6 +39,11 @@ public class SqlDbHelper {
38 39
39 public static Path dumpTableIfExists(Connection conn, String tableName, 40 public static Path dumpTableIfExists(Connection conn, String tableName,
40 String[] columns, String[] defaultValues, String dumpPrefix) throws Exception { 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 DatabaseMetaData metaData = conn.getMetaData(); 48 DatabaseMetaData metaData = conn.getMetaData();
43 ResultSet res = metaData.getTables(null, null, tableName, 49 ResultSet res = metaData.getTables(null, null, tableName,
@@ -46,7 +52,11 @@ public class SqlDbHelper { @@ -46,7 +52,11 @@ public class SqlDbHelper {
46 res.close(); 52 res.close();
47 Path dumpFile = Files.createTempFile(dumpPrefix, null); 53 Path dumpFile = Files.createTempFile(dumpPrefix, null);
48 Files.deleteIfExists(dumpFile); 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 try (PreparedStatement stmt = conn.prepareStatement("SELECT * FROM " + tableName)) { 60 try (PreparedStatement stmt = conn.prepareStatement("SELECT * FROM " + tableName)) {
51 try (ResultSet tableRes = stmt.executeQuery()) { 61 try (ResultSet tableRes = stmt.executeQuery()) {
52 ResultSetMetaData resMetaData = tableRes.getMetaData(); 62 ResultSetMetaData resMetaData = tableRes.getMetaData();
@@ -68,19 +78,30 @@ public class SqlDbHelper { @@ -68,19 +78,30 @@ public class SqlDbHelper {
68 } 78 }
69 79
70 public static void loadTable(Connection conn, String tableName, String[] columns, Path sourceFile) throws Exception { 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