added(feat): excludeColumns option parameter

Author: c0dearm

README.md

@@ -133,6 +133,8 @@ com.spotify.dbeam.options.JdbcExportPipelineOptions:
   --nullableArrayItems=<Boolean>
     Default: false
     Controls whether array items should be nullable, ignored if arrayMode is 'bytes'. 
+  --excludeColumns=<String>
+    A comma-separated list of columns to be excluded from the export.
 ```
 
 #### Input Avro schema file

dbeam-core/src/main/java/com/spotify/dbeam/args/JdbcExportArgs.java

@@ -22,9 +22,11 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableSet;
 import java.io.Serializable;
 import java.sql.Connection;
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.Optional;
 import org.apache.avro.Schema;
 
@@ -49,6 +51,17 @@ public abstract class JdbcExportArgs implements Serializable {
 
   public abstract Optional<Schema> inputAvroSchema();
 
+  public abstract Optional<ImmutableSet<String>> excludedColumns();
+
+  public static Optional<ImmutableSet<String>> parseExcludedColumns(
+      final Optional<String> rawExcludedColumns) {
+    return rawExcludedColumns.map(
+        s ->
+            Arrays.stream(s.split(","))
+                .map(String::trim)
+                .collect(ImmutableSet.toImmutableSet()));
+  }
+
   @AutoValue.Builder
   abstract static class Builder {
 
@@ -68,6 +81,8 @@ abstract static class Builder {
 
     abstract Builder setInputAvroSchema(Optional<Schema> inputAvroSchema);
 
+    abstract Builder setExcludedColumns(Optional<ImmutableSet<String>> excludedColumns);
+
     abstract JdbcExportArgs build();
   }
 
@@ -82,6 +97,7 @@ static JdbcExportArgs create(
         Optional.empty(),
         false,
         Duration.ofDays(7),
+        Optional.empty(),
         Optional.empty());
   }
 
@@ -93,7 +109,8 @@ public static JdbcExportArgs create(
       final Optional<String> avroDoc,
       final Boolean useAvroLogicalTypes,
       final Duration exportTimeout,
-      final Optional<Schema> inputAvroSchema) {
+      final Optional<Schema> inputAvroSchema,
+      final Optional<ImmutableSet<String>> excludedColumns) {
     return new AutoValue_JdbcExportArgs.Builder()
         .setJdbcAvroOptions(jdbcAvroArgs)
         .setQueryBuilderArgs(queryBuilderArgs)
@@ -103,6 +120,7 @@ public static JdbcExportArgs create(
         .setUseAvroLogicalTypes(useAvroLogicalTypes)
         .setExportTimeout(exportTimeout)
         .setInputAvroSchema(inputAvroSchema)
+        .setExcludedColumns(excludedColumns)
         .build();
   }
 

dbeam-core/src/main/java/com/spotify/dbeam/args/QueryBuilder.java

@@ -21,6 +21,7 @@
 package com.spotify.dbeam.args;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import java.io.Serializable;
 import java.util.List;
 import java.util.Optional;
@@ -117,18 +118,24 @@ public int hashCode() {
   private final QueryBase base;
   private final List<String> whereConditions;
   private final Optional<String> limitStr;
+  private final Optional<ImmutableSet<String>> excludedColumns;
 
   private QueryBuilder(final QueryBase base) {
     this.base = base;
     this.limitStr = Optional.empty();
     this.whereConditions = ImmutableList.of();
+    this.excludedColumns = Optional.empty();
   }
 
   private QueryBuilder(
-      final QueryBase base, final List<String> whereConditions, final Optional<String> limitStr) {
+      final QueryBase base,
+      final List<String> whereConditions,
+      final Optional<String> limitStr,
+      final Optional<ImmutableSet<String>> excludedColumns) {
     this.base = base;
     this.whereConditions = whereConditions;
     this.limitStr = limitStr;
+    this.excludedColumns = excludedColumns;
   }
 
   public static QueryBuilder fromTablename(final String tableName) {
@@ -148,7 +155,22 @@ public QueryBuilder withPartitionCondition(
                 Stream.of(
                     createSqlPartitionCondition(partitionColumn, startPointIncl, endPointExcl)))
             .collect(Collectors.toList()),
-        this.limitStr);
+        this.limitStr,
+        this.excludedColumns);
+  }
+
+  public QueryBuilder withExcludedColumns(final Optional<ImmutableSet<String>> excludedColumns) {
+    if (excludedColumns.isPresent() && this.base instanceof UserQueryBase) {
+      UserQueryBase userQueryBase = (UserQueryBase) this.base;
+      String newSqlQuery = rebuildSelectClause(userQueryBase.userSqlQuery, excludedColumns.get());
+      return new QueryBuilder(
+          new UserQueryBase(newSqlQuery, userQueryBase.selectClause),
+          this.whereConditions,
+          this.limitStr,
+          excludedColumns);
+    } else {
+      return new QueryBuilder(this.base, this.whereConditions, this.limitStr, excludedColumns);
+    }
   }
 
   private static String createSqlPartitionCondition(
@@ -171,7 +193,8 @@ public QueryBuilder withParallelizationCondition(
                     createSqlSplitCondition(
                         partitionColumn, startPointIncl, endPoint, isEndPointExcl)))
             .collect(Collectors.toList()),
-        this.limitStr);
+        this.limitStr,
+        this.excludedColumns);
   }
 
   private static String createSqlSplitCondition(
@@ -205,9 +228,38 @@ private static String removeTrailingSymbols(String sqlQuery) {
     return sqlQuery.replaceAll(regex, "$1");
   }
 
+  private static String rebuildSelectClause(
+      String sqlQuery, ImmutableSet<String> excludedColumns) {
+    String lowerCaseQuery = sqlQuery.toLowerCase();
+    int selectIdx = lowerCaseQuery.indexOf("select");
+    int fromIdx = lowerCaseQuery.indexOf("from");
+
+    if (selectIdx == -1 || fromIdx == -1 || selectIdx > fromIdx) {
+      // Cannot parse, return original query
+      return sqlQuery;
+    }
+
+    String selectClause = sqlQuery.substring(selectIdx + "select".length(), fromIdx).trim();
+    String[] columns = selectClause.split(",");
+    List<String> newColumns =
+        Stream.of(columns)
+            .map(String::trim)
+            .filter(column -> !excludedColumns.contains(column))
+            .collect(Collectors.toList());
+
+    if (newColumns.isEmpty()) {
+      return "SELECT * " + sqlQuery.substring(fromIdx);
+    } else {
+      return "SELECT " + String.join(", ", newColumns) + " " + sqlQuery.substring(fromIdx);
+    }
+  }
+
   public QueryBuilder withLimit(long limit) {
     return new QueryBuilder(
-        this.base, this.whereConditions, Optional.of(String.format(" LIMIT %d", limit)));
+        this.base,
+        this.whereConditions,
+        Optional.of(String.format(" LIMIT %d", limit)),
+        this.excludedColumns);
   }
 
   @Override
@@ -248,6 +300,11 @@ public QueryBuilder generateQueryToGetLimitsOfSplitColumn(
             "SELECT MIN(%s) as %s, MAX(%s) as %s",
             splitColumn, minSplitColumnName, splitColumn, maxSplitColumnName);
 
-    return new QueryBuilder(base.withSelect(selectMinMax), this.whereConditions, this.limitStr);
+    return new QueryBuilder(
+        base.withSelect(selectMinMax),
+        this.whereConditions,
+        this.limitStr,
+        this.excludedColumns);
   }
 }
+

dbeam-core/src/main/java/com/spotify/dbeam/args/QueryBuilderArgs.java

@@ -25,6 +25,7 @@
 import static com.spotify.dbeam.args.ParallelQueryBuilder.queriesForBounds;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import java.io.Serializable;
 import java.sql.Connection;
@@ -57,6 +58,8 @@ public abstract class QueryBuilderArgs implements Serializable {
 
   public abstract Optional<Integer> queryParallelism();
 
+  public abstract Optional<ImmutableSet<String>> excludedColumns();
+
   public abstract Builder builder();
 
   @AutoValue.Builder
@@ -86,6 +89,8 @@ public abstract static class Builder {
 
     public abstract Builder setQueryParallelism(Optional<Integer> queryParallelism);
 
+    public abstract Builder setExcludedColumns(Optional<ImmutableSet<String>> excludedColumns);
+
     public abstract QueryBuilderArgs build();
   }
 
@@ -122,6 +127,9 @@ public String sqlQueryWithLimitOne() {
    */
   public List<String> buildQueries(final Connection connection) throws SQLException {
     QueryBuilder queryBuilder = this.baseSqlQuery();
+    if (this.excludedColumns().isPresent()) {
+      queryBuilder = queryBuilder.withExcludedColumns(this.excludedColumns());
+    }
     if (this.partitionColumn().isPresent() && this.partition().isPresent()) {
       queryBuilder =
           configurePartitionCondition(

dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroSchema.java

@@ -50,6 +50,7 @@
 import static java.sql.Types.VARBINARY;
 import static java.sql.Types.VARCHAR;
 
+import com.google.common.collect.ImmutableSet;
 import com.spotify.dbeam.args.QueryBuilderArgs;
 import com.spotify.dbeam.options.ArrayHandlingMode;
 import java.sql.Array;
@@ -95,9 +96,14 @@ public static Schema createSchemaByReadingOneRow(
               avroDoc,
               useLogicalTypes,
               arrayMode,
-              nullableArrayItems);
-      LOGGER.info("Schema created successfully. useLogicalTypes={}, arrayMode={}, "
-                  + "Generated schema: {}", useLogicalTypes, arrayMode, schema.toString());
+              nullableArrayItems,
+              queryBuilderArgs.excludedColumns());
+      LOGGER.info(
+          "Schema created successfully. useLogicalTypes={}, arrayMode={}, "
+              + "Generated schema: {}",
+          useLogicalTypes,
+          arrayMode,
+          schema.toString());
       return schema;
     }
   }
@@ -110,7 +116,8 @@ public static Schema createAvroSchema(
       final String avroDoc,
       final boolean useLogicalTypes,
       final String arrayMode,
-      final boolean nullableArrayItems)
+      final boolean nullableArrayItems,
+      final Optional<ImmutableSet<String>> excludedColumns)
       throws SQLException {
 
     final ResultSetMetaData meta = resultSet.getMetaData();
@@ -124,13 +131,18 @@ public static Schema createAvroSchema(
             .prop("tableName", tableName)
             .prop("connectionUrl", connectionUrl)
             .fields();
-    return createAvroFields(resultSet, builder, useLogicalTypes, arrayMode, nullableArrayItems)
+    return createAvroFields(
+            resultSet,
+            builder,
+            useLogicalTypes,
+            arrayMode,
+            nullableArrayItems,
+            excludedColumns.orElse(ImmutableSet.of()))
         .endRecord();
   }
 
   static String getDatabaseTableName(final ResultSetMetaData meta) throws SQLException {
     final String defaultTableName = "no_table_name";
-
     for (int i = 1; i <= meta.getColumnCount(); i++) {
       String metaTableName = meta.getTableName(i);
       if (metaTableName != null && !metaTableName.isEmpty()) {
@@ -145,20 +157,24 @@ private static SchemaBuilder.FieldAssembler<Schema> createAvroFields(
       final SchemaBuilder.FieldAssembler<Schema> builder,
       final boolean useLogicalTypes,
       final String arrayMode,
-      final boolean nullableArrayItems)
+      final boolean nullableArrayItems,
+      final ImmutableSet<String> excludedColumns)
       throws SQLException {
 
     ResultSetMetaData meta = resultSet.getMetaData();
 
     for (int i = 1; i <= meta.getColumnCount(); i++) {
-
       final String columnName;
       if (meta.getColumnName(i).isEmpty()) {
         columnName = meta.getColumnLabel(i);
       } else {
         columnName = meta.getColumnName(i);
       }
 
+      if (excludedColumns.contains(columnName)) {
+        continue;
+      }
+      
       final int columnType = meta.getColumnType(i);
       final String typeName = JDBCType.valueOf(columnType).getName();
       final String columnClassName = meta.getColumnClassName(i);

dbeam-core/src/main/java/com/spotify/dbeam/options/JdbcExportArgsFactory.java

@@ -81,7 +81,9 @@ public static JdbcExportArgs fromPipelineOptions(final PipelineOptions options)
         Optional.ofNullable(exportOptions.getAvroDoc()),
         exportOptions.isUseAvroLogicalTypes(),
         Duration.parse(exportOptions.getExportTimeout()),
-        BeamJdbcAvroSchema.parseOptionalInputAvroSchemaFile(exportOptions.getAvroSchemaFilePath()));
+        BeamJdbcAvroSchema.parseOptionalInputAvroSchemaFile(exportOptions.getAvroSchemaFilePath()),
+        JdbcExportArgs.parseExcludedColumns(
+            Optional.ofNullable(exportOptions.getExcludeColumns())));
   }
 
   public static QueryBuilderArgs createQueryArgs(final JdbcExportPipelineOptions options)
@@ -129,6 +131,9 @@ public static QueryBuilderArgs createQueryArgs(final JdbcExportPipelineOptions o
         .setPartitionPeriod(partitionPeriod)
         .setSplitColumn(splitColumn)
         .setQueryParallelism(queryParallelism)
+        .setExcludedColumns(
+            JdbcExportArgs.parseExcludedColumns(
+                Optional.ofNullable(options.getExcludeColumns())))
         .build();
   }
 

dbeam-core/src/main/java/com/spotify/dbeam/options/JdbcExportPipelineOptions.java

@@ -149,4 +149,9 @@ public interface JdbcExportPipelineOptions extends DBeamPipelineOptions {
   Long getMinRows();
 
   void setMinRows(Long value);
+
+  @Description("A comma-separated list of columns to be excluded from the export.")
+  String getExcludeColumns();
+
+  void setExcludeColumns(String value);
 }

dbeam-core/src/test/java/com/spotify/dbeam/args/JdbcExportOptionsTest.java

@@ -368,4 +368,16 @@ public void shouldFailOnNegativeQueryParallelism() throws IOException, ClassNotF
         "--connectionUrl=jdbc:postgresql://some_db --table=some_table "
             + "--password=secret --queryParallelism=-5 --splitColumn=id");
   }
+
+  @Test
+  public void shouldParseExcludedColumns() throws IOException, ClassNotFoundException {
+    final JdbcExportArgs options =
+        optionsFromArgs(
+            "--connectionUrl=jdbc:postgresql://some_db --table=some_table "
+                + "--password=secret --excludeColumns=col1,col2");
+
+    Assert.assertEquals(
+        Optional.of(com.google.common.collect.ImmutableSet.of("col1", "col2")),
+        options.excludedColumns());
+  }
 }

dbeam-core/src/test/java/com/spotify/dbeam/args/QueryBuilderTest.java

@@ -210,6 +210,34 @@ public void testItGeneratesQueryForLimits() {
     Assert.assertEquals(expected, actual);
   }
 
+
+  @Test
+  public void testTableNameWithExcludedColumnsShouldStillUseSelectStar() {
+    final com.google.common.collect.ImmutableSet<String> excludedColumns =
+        com.google.common.collect.ImmutableSet.of("col1");
+    final QueryBuilder wrapper =
+        QueryBuilder.fromTablename("some_table")
+            .withExcludedColumns(java.util.Optional.of(excludedColumns));
+
+    final String expected = "SELECT * FROM some_table WHERE 1=1";
+
+    Assert.assertEquals(expected, wrapper.build());
+  }
+
+  @Test
+  public void testRawSqlWithExcludedColumnShouldRemoveColumns() {
+    final com.google.common.collect.ImmutableSet<String> excludedColumns =
+        com.google.common.collect.ImmutableSet.of("col1");
+    final QueryBuilder wrapper =
+        QueryBuilder.fromSqlQuery("SELECT col1, col2 FROM some_table")
+            .withExcludedColumns(java.util.Optional.of(excludedColumns));
+
+    final String expected = 
+        "SELECT * FROM (SELECT col2 FROM some_table) as user_sql_query WHERE 1=1";
+
+    Assert.assertEquals(expected, wrapper.build());
+  }
+
   private void execAndCompare(String rawInput, String expected) {
     final String actual = QueryBuilder.fromSqlQuery(rawInput).build();