diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml index c20bf249..3eea9fc0 100644 --- a/.github/workflows/maven.yml +++ b/.github/workflows/maven.yml @@ -21,7 +21,7 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - java_version: ['11', '17', '21'] + java_version: ['17', '21', '25'] os: [ubuntu-latest] steps: - uses: actions/checkout@v4 @@ -44,27 +44,31 @@ jobs: uses: codecov/codecov-action@v4 with: token: ${{ secrets.CODECOV_TOKEN }} - files: ./dbeam-core/target/site/jacoco/jacoco.xml + files: ./dbeam-core/target/site/jacoco/jacoco.xml,./dbeam-parquet/target/site/jacoco/jacoco.xml flags: unittests # optional name: codecov-umbrella # optional fail_ci_if_error: false # optional (default = false) verbose: true # optional (default = false) - - name: Install avro tools - if: matrix.java_version == '17' + - name: Install avro and parquet tools + if: matrix.java_version == '21' run: | - wget https://dlcdn.apache.org/avro/avro-1.11.3/java/avro-tools-1.11.3.jar + wget https://dlcdn.apache.org/avro/avro-1.11.5/java/avro-tools-1.11.5.jar mkdir -p /opt/avro - mv avro-tools-1.11.3.jar /opt/avro/ - chmod +x /opt/avro/avro-tools-1.11.3.jar - echo "alias avro-tools='java -jar /opt/avro/avro-tools-1.11.3.jar'" > ~/.bashrc + mv avro-tools-1.11.5.jar /opt/avro/ + chmod +x /opt/avro/avro-tools-1.11.5.jar + echo "alias avro-tools='java -jar /opt/avro/avro-tools-1.11.5.jar'" > ~/.bashrc + wget -q https://repo1.maven.org/maven2/org/apache/parquet/parquet-cli/1.14.4/parquet-cli-1.14.4-runtime.jar + mkdir -p /opt/parquet + mv parquet-cli-1.14.4-runtime.jar /opt/parquet/ + echo "alias parquet-tools='java -jar /opt/parquet/parquet-cli-1.14.4-runtime.jar'" >> ~/.bashrc - name: End to end tests shell: bash - if: matrix.java_version == '17' + if: matrix.java_version == '21' run: ./e2e/e2e.sh - run: mkdir staging && cp /tmp/debeam_e2e.log ./staging - if: matrix.java_version == '17' + if: matrix.java_version == '21' - uses: actions/upload-artifact@v4 - if: github.event_name == 'push' && (github.ref == 'refs/heads/master' || github.ref == 'refs/heads/main') && matrix.java_version == '17' + if: github.event_name == 'push' && (github.ref == 'refs/heads/master' || github.ref == 'refs/heads/main') && matrix.java_version == '21' with: name: Package path: staging @@ -86,12 +90,12 @@ jobs: uses: actions/cache@v4 with: path: ~/.m2 - key: Linux-java11-m2-${{ hashFiles('**/pom.xml') }} - restore-keys: Linux-java11-m2 - - name: Set up JDK 11 + key: Linux-java21-m2-${{ hashFiles('**/pom.xml') }} + restore-keys: Linux-java21-m2 + - name: Set up JDK 21 uses: actions/setup-java@v4 with: - java-version: '11' + java-version: '21' distribution: 'adopt' # use Corretto once supported https://github.com/actions/setup-java/issues/68 - name: Deploy env: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 2ec3169f..d41c40e4 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -23,12 +23,12 @@ jobs: uses: actions/cache@v4 with: path: ~/.m2 - key: Linux-java11-m2-${{ hashFiles('**/pom.xml') }} - restore-keys: Linux-java11-m2 - - name: Set up JDK 11 + key: Linux-java17-m2-${{ hashFiles('**/pom.xml') }} + restore-keys: Linux-java17-m2 + - name: Set up JDK 17 uses: actions/setup-java@v4 with: - java-version: '11' + java-version: '17' distribution: 'adopt' # setup-java generates a settings.xml pointing deployments to Sonatype # See https://github.com/actions/setup-java/blob/v3.11.0/docs/advanced-usage.md#publishing-using-apache-maven diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000..0f7d5295 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,42 @@ +# Changelog + +## [Unreleased] + +### Added + +- **dbeam-parquet module**: New module for exporting SQL databases directly to Parquet files, + complementing the existing Avro export in dbeam-core. + - `JdbcParquetJob`: Main entry point for Parquet exports, reuses all shared infrastructure + from dbeam-core (JDBC connectors, query building, partitioning, parallel queries, Beam runners). + - `JdbcParquetSchema`: Converts JDBC ResultSetMetaData to Parquet MessageType schema with + full SQL type mapping including logical types (TIMESTAMP_MILLIS, UUID). + - `JdbcParquetWriteSupport`: Writes ResultSet rows directly to Parquet RecordConsumer + without intermediate objects. + - `ChannelOutputFile`: OutputFile implementation backed by WritableByteChannel, bypassing + Hadoop for file I/O. + - Parquet compression codecs: snappy, gzip, zstd, lz4, uncompressed via `--parquetCodec` + option (also auto-maps from `--avroCodec`). + - SQL ARRAY columns use Parquet 3-level LIST type with typed elements (INT32, INT64, FLOAT, + DOUBLE, BOOLEAN, STRING) inferred from column type name. + - Parquet file footer includes `parquet.avro.schema` key with the Avro schema JSON, generated + via JdbcAvroSchema for full compatibility with Spark, Hive, and BigQuery. + - `_AVRO_SCHEMA.avsc` file saved alongside `_PARQUET_SCHEMA.json` in output directory. + - `--parquetSchemaFilePath`: Input schema override (Parquet MessageType text format). + - `--rowGroupSize` (default 128MB) and `--pageSize` (default 1MB) options. + - `BeamJdbcParquetSchema`: Exposes schema creation timing as Beam metric. + - `PsqlParquetJob`: PostgreSQL-specific job with replication lag check. + - `BenchJdbcParquetJob`: Benchmarking job for multi-execution performance testing. + - End-to-end test support in `e2e/e2e.sh` with `parquet-tools` validation. + - Comprehensive test suite (90+ tests) covering schema conversion, record roundtrips, + typed arrays, null handling, mock-based PostgreSQL type tests, and footer metadata. + +### Changed + +- Parent POM: Shade plugin config parameterized via `${dbeam.mainClass}` property, + eliminating duplication between dbeam-core and dbeam-parquet. +- `PsqlReplicationCheck.validateOptions()` made public for cross-module access. +- dbeam-core: Added test-jar packaging for test helper sharing. +- CI: Codecov upload includes dbeam-parquet coverage. Parquet-tools installed for e2e validation. +- e2e: PostgreSQL upgraded from 16 to 18. Parquet test suite runs alongside Avro suite. +- Documentation: README updated with Parquet usage examples, library dependencies, and feature list. + New `docs/parquet-type-conversion.md` with full type mapping table. diff --git a/README.md b/README.md index 8d659edd..5f2d9b47 100644 --- a/README.md +++ b/README.md @@ -17,28 +17,38 @@ This tool is runnable locally, or on any other backend supported by Apache Beam, ## Overview DBeam is a tool that reads all the data from single SQL database table, -converts the data into [Avro](https://avro.apache.org/) and stores it into +converts the data into [Avro](https://avro.apache.org/) or [Parquet](https://parquet.apache.org/) and stores it into appointed location, usually in GCS. It runs as a single threaded [Apache Beam](https://beam.apache.org/) pipeline. DBeam requires the database credentials, the database table name to read, and the output location to store the extracted data into. DBeam first makes a single select into the target table with limit one to infer the table schema. After the schema is created the job will be launched which -simply streams the table contents via JDBC into target location as Avro. +simply streams the table contents via JDBC into target location as Avro or Parquet. -[Generated Avro Schema Type Conversion Details](docs/type-conversion.md) +Type conversion details: [Avro](docs/type-conversion.md) | [Parquet](docs/parquet-type-conversion.md) ## dbeam-core package features - Supports both PostgreSQL, MySQL, MariaDB, and H2 JDBC connectors - Supports [Google CloudSQL](https://cloud.google.com/sql/) managed databases -- Currently outputs only to Avro format +- Outputs to Avro format - Reads database from an external password file (`--passwordFile`) or an external [KMS](https://cloud.google.com/kms/) encrypted password file (`--passwordFileKmsEncrypted`) - Can filter only records of the current day with the `--partitionColumn` parameter - Check and fail on too old partition dates. Snapshot dumps are not filtered by a given date/partition, when running for a too old partition, the job fails to avoid new data in old partitions. (can be disabled with `--skipPartitionCheck`) - Implemented as [Apache Beam SDK](https://beam.apache.org/) pipeline, supporting any of its [runners](https://beam.apache.org/documentation/runners/capability-matrix/) (tested with `DirectRunner` and `DataflowRunner`) +## dbeam-parquet package features + +- Outputs to Parquet format (columnar, optimized for analytical queries) +- Reuses all dbeam-core infrastructure (JDBC connectors, CloudSQL, partitioning, parallel queries) +- Supports Parquet compression codecs: snappy, gzip, zstd, lz4, uncompressed +- No full Hadoop dependency — uses a minimal `hadoop-common` for the Parquet API surface +- Separate `--parquetCodec` option or automatic mapping from `--avroCodec` +- Optional input schema file via `--parquetSchemaFilePath` (Parquet MessageType text format) +- PostgreSQL replication check via `PsqlParquetJob` + ### DBeam export parameters ``` @@ -171,15 +181,19 @@ Building and testing can be achieved with `mvn`: mvn verify ``` -In order to create a jar with all dependencies under `./dbeam-core/target/dbeam-core-shaded.jar` run the following: +In order to create jars with all dependencies run the following: ```sh mvn clean package -Ppack ``` +This produces: +- `./dbeam-core/target/dbeam-core-shaded.jar` — Avro export +- `./dbeam-parquet/target/dbeam-parquet-shaded.jar` — Parquet export + ## Usage examples -Using Java from the command line: +### Avro export ```sh java -cp ./dbeam-core/target/dbeam-core-shaded.jar \ @@ -191,7 +205,20 @@ java -cp ./dbeam-core/target/dbeam-core-shaded.jar \ --table=my_table ``` -For CloudSQL: +### Parquet export + +```sh +java -cp ./dbeam-parquet/target/dbeam-parquet-shaded.jar \ + com.spotify.dbeam.parquet.JdbcParquetJob \ + --output=gs://my-testing-bucket-name/ \ + --username=my_database_username \ + --password=secret \ + --connectionUrl=jdbc:postgresql://some.database.uri.example.org:5432/my_database \ + --table=my_table \ + --parquetCodec=snappy +``` + +For CloudSQL (Avro): ```sh java -cp ./dbeam-core/target/dbeam-core-shaded.jar \ @@ -256,11 +283,19 @@ When using Google Cloud, [IAM authentication](https://github.com/GoogleCloudPlat To include DBeam library in a mvn project add the following dependency in `pom.xml`: ```xml + com.spotify dbeam-core ${dbeam.version} + + + + com.spotify + dbeam-parquet + ${dbeam.version} + ``` @@ -268,7 +303,8 @@ To include DBeam library in a SBT project add the following dependency in `build ```sbt libraryDependencies ++= Seq( - "com.spotify" % "dbeam-core" % dbeamVersion + "com.spotify" % "dbeam-core" % dbeamVersion, // Avro + "com.spotify" % "dbeam-parquet" % dbeamVersion // Parquet ) ``` diff --git a/dbeam-bom/pom.xml b/dbeam-bom/pom.xml index 4e7d364e..f4435743 100644 --- a/dbeam-bom/pom.xml +++ b/dbeam-bom/pom.xml @@ -21,6 +21,11 @@ dbeam-core ${project.version} + + com.spotify + dbeam-parquet + ${project.version} + @@ -47,7 +52,7 @@ org.codehaus.mojo build-helper-maven-plugin - 3.6.0 + 3.6.1 attach-effective-pom diff --git a/dbeam-core/pom.xml b/dbeam-core/pom.xml index 6e2cd66e..c5fff8e3 100644 --- a/dbeam-core/pom.xml +++ b/dbeam-core/pom.xml @@ -13,6 +13,10 @@ DBeam Core Top level DBeam core implementation + + com.spotify.dbeam.jobs.JdbcAvroJob + + @@ -93,8 +97,8 @@ - junit - junit + org.junit.jupiter + junit-jupiter test @@ -124,4 +128,20 @@ + + + + org.apache.maven.plugins + maven-jar-plugin + + + + test-jar + + + + + + + diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/args/JdbcAvroArgs.java b/dbeam-core/src/main/java/com/spotify/dbeam/args/JdbcAvroArgs.java index f3d04aa0..92e3c841 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/args/JdbcAvroArgs.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/args/JdbcAvroArgs.java @@ -103,8 +103,8 @@ public static JdbcAvroArgs create( } public static JdbcAvroArgs create(final JdbcConnectionArgs jdbcConnectionArgs) { - return create(jdbcConnectionArgs, 10000, "deflate6", Collections.emptyList(), - "typed_first_row", false); + return create( + jdbcConnectionArgs, 10000, "deflate6", Collections.emptyList(), "typed_first_row", false); } public interface StatementPreparator extends Serializable { diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroIO.java b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroIO.java index 92e69b78..74d327a8 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroIO.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroIO.java @@ -25,7 +25,6 @@ import com.google.common.collect.ImmutableMap; import com.google.common.io.CountingOutputStream; import com.spotify.dbeam.args.JdbcAvroArgs; -import com.spotify.dbeam.args.JdbcExportArgs; import java.nio.channels.Channels; import java.nio.channels.WritableByteChannel; import java.sql.Connection; @@ -137,7 +136,7 @@ private static class JdbcAvroWriter extends FileBasedSink.Writer { private final JdbcAvroArgs jdbcAvroArgs; private DataFileWriter dataFileWriter; private Connection connection; - private JdbcAvroMetering metering; + private JdbcMetering metering; private CountingOutputStream countingOutputStream; JdbcAvroWriter( @@ -147,7 +146,7 @@ private static class JdbcAvroWriter extends FileBasedSink.Writer { super(writeOperation, MimeTypes.BINARY); this.dynamicDestinations = dynamicDestinations; this.jdbcAvroArgs = jdbcAvroArgs; - this.metering = JdbcAvroMetering.create(); + this.metering = JdbcMetering.create("jdbcavroio"); } public Void getDestination() { @@ -181,8 +180,7 @@ private ResultSet executeQuery(final String query) throws Exception { jdbcAvroArgs.statementPreparator().setParameters(statement); } - if (jdbcAvroArgs.preCommand() != null - && !jdbcAvroArgs.preCommand().isEmpty()) { + if (jdbcAvroArgs.preCommand() != null && !jdbcAvroArgs.preCommand().isEmpty()) { final Statement stmt = connection.createStatement(); for (String command : jdbcAvroArgs.preCommand()) { stmt.execute(command); @@ -205,8 +203,9 @@ public void write(final String query) throws Exception { LOGGER.info("jdbcavroio : Starting write..."); try (ResultSet resultSet = executeQuery(query)) { metering.startWriteMeter(); - final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(resultSet, - this.jdbcAvroArgs.arrayMode(), this.jdbcAvroArgs.nullableArrayItems()); + final JdbcAvroRecordConverter converter = + JdbcAvroRecordConverter.create( + resultSet, this.jdbcAvroArgs.arrayMode(), this.jdbcAvroArgs.nullableArrayItems()); while (resultSet.next()) { dataFileWriter.appendEncoded(converter.convertResultSetIntoAvroBytes()); this.metering.incrementRecordCount(); diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroMetering.java b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroMetering.java index cb317775..c01752df 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroMetering.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroMetering.java @@ -20,89 +20,17 @@ package com.spotify.dbeam.avro; -import org.apache.beam.sdk.metrics.Counter; -import org.apache.beam.sdk.metrics.Gauge; -import org.apache.beam.sdk.metrics.Metrics; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -public class JdbcAvroMetering { - - public static final String RECORD_COUNT_METRIC_NAME = "recordCount"; - private final int countReportEvery; - private final int logEvery; - private final Logger logger = LoggerFactory.getLogger(JdbcAvroMetering.class); - private Counter recordCount = - Metrics.counter(this.getClass().getCanonicalName(), RECORD_COUNT_METRIC_NAME); - private Counter executeQueryElapsedMs = - Metrics.counter(this.getClass().getCanonicalName(), "executeQueryElapsedMs"); - private Counter writeElapsedMs = - Metrics.counter(this.getClass().getCanonicalName(), "writeElapsedMs"); - private Gauge msPerMillionRows = - Metrics.gauge(this.getClass().getCanonicalName(), "msPerMillionRows"); - private Gauge rowsPerMinute = Metrics.gauge(this.getClass().getCanonicalName(), "rowsPerMinute"); - private Counter bytesWritten = - Metrics.counter(this.getClass().getCanonicalName(), "bytesWritten"); - private long rowCount = 0; - private long writeIterateStartTime; +/** + * @deprecated Use {@link JdbcMetering} directly instead. + */ +@Deprecated +public class JdbcAvroMetering extends JdbcMetering { public JdbcAvroMetering(int countReportEvery, int logEvery) { - this.countReportEvery = countReportEvery; - this.logEvery = logEvery; + super(countReportEvery, logEvery, "jdbcavroio"); } public static JdbcAvroMetering create() { return new JdbcAvroMetering(100000, 100000); } - - /** - * Increment and report counters to Beam SDK and logs. To avoid slowing down the writes, counts - * are reported every x 1000s of rows. This exposes the job progress. - */ - public void incrementRecordCount() { - this.rowCount++; - if ((this.rowCount % countReportEvery) == 0) { - this.recordCount.inc(countReportEvery); - final long elapsedNano = System.nanoTime() - this.writeIterateStartTime; - final long msPerMillionRows = elapsedNano / rowCount; - final long rowsPerMinute = (60 * 1000000000L) * rowCount / elapsedNano; - this.msPerMillionRows.set(msPerMillionRows); - this.rowsPerMinute.set(rowsPerMinute); - if ((this.rowCount % logEvery) == 0) { - logger.info( - String.format( - "jdbcavroio : Fetched # %08d rows at %08d rows per minute and %08d ms per M rows", - rowCount, rowsPerMinute, msPerMillionRows)); - } - } - } - - public void exposeWriteElapsed() { - long elapsedMs = (System.nanoTime() - this.writeIterateStartTime) / 1000000L; - logger.info("jdbcavroio : Read {} rows, took {} seconds", rowCount, elapsedMs / 1000.0); - this.writeElapsedMs.inc(elapsedMs); - if (rowCount > 0) { - this.recordCount.inc((this.rowCount % countReportEvery)); - this.msPerMillionRows.set(1000000L * elapsedMs / rowCount); - if (elapsedMs != 0) { - this.rowsPerMinute.set((60 * 1000L) * rowCount / elapsedMs); - } - } - } - - public long startWriteMeter() { - long startTs = System.nanoTime(); - this.writeIterateStartTime = startTs; - this.rowCount = 0; - return startTs; - } - - public void exposeExecuteQueryMs(final long elapsedMs) { - logger.info("jdbcavroio : Execute query took {} seconds", elapsedMs / 1000.0); - this.executeQueryElapsedMs.inc(elapsedMs); - } - - public void exposeWrittenBytes(final long count) { - this.bytesWritten.inc(count); - } } diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroRecord.java b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroRecord.java index 082cb489..7abffd58 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroRecord.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroRecord.java @@ -51,10 +51,8 @@ private static ByteBuffer nullableBytes(final byte[] bts) { } } - static SqlFunction computeMapping(final ResultSetMetaData meta, - final int column, - final String arrayMode) - throws SQLException { + static SqlFunction computeMapping( + final ResultSetMetaData meta, final int column, final String arrayMode) throws SQLException { switch (meta.getColumnType(column)) { case VARCHAR: case CHAR: diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroRecordConverter.java b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroRecordConverter.java index 576d2b4a..9c96f593 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroRecordConverter.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroRecordConverter.java @@ -49,9 +49,8 @@ public JdbcAvroRecordConverter( this.nullableArrayItems = nullableArrayItems; } - public static JdbcAvroRecordConverter create(final ResultSet resultSet, - final String arrayMode, - final boolean nullableArrayItems) + public static JdbcAvroRecordConverter create( + final ResultSet resultSet, final String arrayMode, final boolean nullableArrayItems) throws SQLException { return new JdbcAvroRecordConverter( computeAllMappings(resultSet, arrayMode), @@ -62,8 +61,7 @@ public static JdbcAvroRecordConverter create(final ResultSet resultSet, @SuppressWarnings("unchecked") static JdbcAvroRecord.SqlFunction[] computeAllMappings( - final ResultSet resultSet, final String arrayMode) - throws SQLException { + final ResultSet resultSet, final String arrayMode) throws SQLException { final ResultSetMetaData meta = resultSet.getMetaData(); final int columnCount = meta.getColumnCount(); @@ -151,8 +149,8 @@ private void writeValue(Object value, String column, BinaryEncoder binaryEncoder } else { if (arrayItem == null) { throw new RuntimeException( - String.format("Array item is null in column '%s', use --nullableArrayItems", - column)); + String.format( + "Array item is null in column '%s', use --nullableArrayItems", column)); } writeValue(arrayItem, column, binaryEncoder); @@ -162,8 +160,8 @@ private void writeValue(Object value, String column, BinaryEncoder binaryEncoder binaryEncoder.writeArrayEnd(); } else { throw new RuntimeException( - String.format("Value of type %s in column '%s' is not supported", value.getClass(), - column)); + String.format( + "Value of type %s in column '%s' is not supported", value.getClass(), column)); } } } diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroSchema.java b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroSchema.java index 9d2eefee..39f3cec9 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroSchema.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcAvroSchema.java @@ -96,8 +96,12 @@ public static Schema createSchemaByReadingOneRow( useLogicalTypes, arrayMode, nullableArrayItems); - LOGGER.info("Schema created successfully. useLogicalTypes={}, arrayMode={}, " - + "Generated schema: {}", useLogicalTypes, arrayMode, schema.toString()); + LOGGER.info( + "Schema created successfully. useLogicalTypes={}, arrayMode={}, " + + "Generated schema: {}", + useLogicalTypes, + arrayMode, + schema.toString()); return schema; } } @@ -181,9 +185,11 @@ private static SchemaBuilder.FieldAssembler createAvroFields( fieldSchemaBuilder = field.type().unionOf().nullBuilder().endNull().and(); Array arrayInstance = - resultSet.isFirst() && columnType == ARRAY - && arrayMode.equals(ArrayHandlingMode.TypedMetaFromFirstRow) - ? resultSet.getArray(i) : null; + resultSet.isFirst() + && columnType == ARRAY + && arrayMode.equals(ArrayHandlingMode.TypedMetaFromFirstRow) + ? resultSet.getArray(i) + : null; final SchemaBuilder.UnionAccumulator> schemaFieldAssembler = buildAvroFieldType( @@ -216,17 +222,19 @@ private static SchemaBuilder.FieldAssembler createAvroFields( */ private static SchemaBuilder.UnionAccumulator> buildAvroFieldType( - final String columnName, - final int columnType, - final Array arrayInstance, - final int precision, - final String columnClassName, - final String columnTypeName, - final boolean useLogicalTypes, - final String arrayMode, - final boolean nullableArrayItems, - final SchemaBuilder.BaseTypeBuilder>> field) throws SQLException { + final String columnName, + final int columnType, + final Array arrayInstance, + final int precision, + final String columnClassName, + final String columnTypeName, + final boolean useLogicalTypes, + final String arrayMode, + final boolean nullableArrayItems, + final SchemaBuilder.BaseTypeBuilder< + SchemaBuilder.UnionAccumulator>> + field) + throws SQLException { switch (columnType) { case BIGINT: return field.longType(); @@ -265,19 +273,24 @@ private static SchemaBuilder.FieldAssembler createAvroFields( if (arrayMode.equals(ArrayHandlingMode.TypedMetaPostgres)) { if (!columnTypeName.startsWith("_")) { - throw new RuntimeException("columnName=" + columnName - + " columnTypeName=" + columnTypeName - + " should start with '_'"); + throw new RuntimeException( + "columnName=" + + columnName + + " columnTypeName=" + + columnTypeName + + " should start with '_'"); } - return buildAvroFieldTypeFromPGType(columnName, columnTypeName.substring(1), - useLogicalTypes, buildArrayItems(nullableArrayItems, field)); + return buildAvroFieldTypeFromPGType( + columnName, + columnTypeName.substring(1), + useLogicalTypes, + buildArrayItems(nullableArrayItems, field)); } if (arrayInstance == null) { throw new RuntimeException( - "When inspecting ARRAY column type in '" + columnName - + "' its first value is NULL"); + "When inspecting ARRAY column type in '" + columnName + "' its first value is NULL"); } return buildAvroFieldType( columnName, @@ -324,12 +337,13 @@ private static SchemaBuilder.FieldAssembler createAvroFields( } } - private static SchemaBuilder.BaseTypeBuilder>> + private static SchemaBuilder.BaseTypeBuilder< + SchemaBuilder.UnionAccumulator>> buildArrayItems( - final boolean nullableArrayItems, - final SchemaBuilder.BaseTypeBuilder>> field) { + final boolean nullableArrayItems, + final SchemaBuilder.BaseTypeBuilder< + SchemaBuilder.UnionAccumulator>> + field) { if (nullableArrayItems) { return field.array().items().nullable(); } else { @@ -339,12 +353,12 @@ private static SchemaBuilder.FieldAssembler createAvroFields( private static SchemaBuilder.UnionAccumulator> buildAvroFieldTypeFromPGType( - final String columnName, - final String columnTypeName, - final boolean useLogicalTypes, - final SchemaBuilder.BaseTypeBuilder< - SchemaBuilder.UnionAccumulator>> - field) { + final String columnName, + final String columnTypeName, + final boolean useLogicalTypes, + final SchemaBuilder.BaseTypeBuilder< + SchemaBuilder.UnionAccumulator>> + field) { switch (columnTypeName) { case "uuid": if (useLogicalTypes) { @@ -361,8 +375,14 @@ private static SchemaBuilder.FieldAssembler createAvroFields( case "text": return field.stringType(); default: - throw new RuntimeException("columnName=" + columnName + " Postgres type '" - + columnTypeName + "' is " + "not " + "supported"); + throw new RuntimeException( + "columnName=" + + columnName + + " Postgres type '" + + columnTypeName + + "' is " + + "not " + + "supported"); } } diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcMetering.java b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcMetering.java new file mode 100644 index 00000000..cd337f9a --- /dev/null +++ b/dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcMetering.java @@ -0,0 +1,110 @@ +/*- + * -\-\- + * DBeam Core + * -- + * Copyright (C) 2016 - 2018 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.avro; + +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Gauge; +import org.apache.beam.sdk.metrics.Metrics; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class JdbcMetering { + + public static final String RECORD_COUNT_METRIC_NAME = "recordCount"; + private final int countReportEvery; + private final int logEvery; + private final String formatName; + private final Logger logger = LoggerFactory.getLogger(JdbcMetering.class); + private Counter recordCount = + Metrics.counter(this.getClass().getCanonicalName(), RECORD_COUNT_METRIC_NAME); + private Counter executeQueryElapsedMs = + Metrics.counter(this.getClass().getCanonicalName(), "executeQueryElapsedMs"); + private Counter writeElapsedMs = + Metrics.counter(this.getClass().getCanonicalName(), "writeElapsedMs"); + private Gauge msPerMillionRows = + Metrics.gauge(this.getClass().getCanonicalName(), "msPerMillionRows"); + private Gauge rowsPerMinute = Metrics.gauge(this.getClass().getCanonicalName(), "rowsPerMinute"); + private Counter bytesWritten = + Metrics.counter(this.getClass().getCanonicalName(), "bytesWritten"); + private long rowCount = 0; + private long writeIterateStartTime; + + public JdbcMetering(int countReportEvery, int logEvery, String formatName) { + this.countReportEvery = countReportEvery; + this.logEvery = logEvery; + this.formatName = formatName; + } + + public static JdbcMetering create(String formatName) { + return new JdbcMetering(100000, 100000, formatName); + } + + /** + * Increment and report counters to Beam SDK and logs. To avoid slowing down the writes, counts + * are reported every x 1000s of rows. This exposes the job progress. + */ + public void incrementRecordCount() { + this.rowCount++; + if ((this.rowCount % countReportEvery) == 0) { + this.recordCount.inc(countReportEvery); + final long elapsedNano = System.nanoTime() - this.writeIterateStartTime; + final long msPerMillionRows = elapsedNano / rowCount; + final long rowsPerMinute = (60 * 1000000000L) * rowCount / elapsedNano; + this.msPerMillionRows.set(msPerMillionRows); + this.rowsPerMinute.set(rowsPerMinute); + if ((this.rowCount % logEvery) == 0) { + logger.info( + String.format( + "%s : Fetched # %08d rows at %08d rows per minute and %08d ms per M rows", + formatName, rowCount, rowsPerMinute, msPerMillionRows)); + } + } + } + + public void exposeWriteElapsed() { + long elapsedMs = (System.nanoTime() - this.writeIterateStartTime) / 1000000L; + logger.info("{} : Read {} rows, took {} seconds", formatName, rowCount, elapsedMs / 1000.0); + this.writeElapsedMs.inc(elapsedMs); + if (rowCount > 0) { + this.recordCount.inc((this.rowCount % countReportEvery)); + this.msPerMillionRows.set(1000000L * elapsedMs / rowCount); + if (elapsedMs != 0) { + this.rowsPerMinute.set((60 * 1000L) * rowCount / elapsedMs); + } + } + } + + public long startWriteMeter() { + long startTs = System.nanoTime(); + this.writeIterateStartTime = startTs; + this.rowCount = 0; + return startTs; + } + + public void exposeExecuteQueryMs(final long elapsedMs) { + logger.info("{} : Execute query took {} seconds", formatName, elapsedMs / 1000.0); + this.executeQueryElapsedMs.inc(elapsedMs); + } + + public void exposeWrittenBytes(final long count) { + this.bytesWritten.inc(count); + } +} diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcAvroJob.java b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcAvroJob.java index 27265f64..269a6eca 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcAvroJob.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcAvroJob.java @@ -20,112 +20,18 @@ package com.spotify.dbeam.jobs; -import static com.google.common.collect.Lists.newArrayList; - -import com.google.common.collect.ImmutableMap; -import com.google.common.math.Stats; -import com.spotify.dbeam.beam.MetricsHelper; -import com.spotify.dbeam.options.OutputOptions; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.function.Function; -import java.util.stream.Collector; -import java.util.stream.Collectors; -import java.util.stream.IntStream; -import java.util.stream.Stream; -import org.apache.beam.sdk.PipelineResult; -import org.apache.beam.sdk.options.Default; -import org.apache.beam.sdk.options.Description; -import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; /** Used on e2e test, allows benchmarking with different configuration parameters. */ public class BenchJdbcAvroJob { - public interface BenchJdbcAvroOptions extends PipelineOptions { - @Description("The JDBC connection url to perform the export.") - @Default.Integer(3) - int getExecutions(); - - void setExecutions(int value); - } - - private final PipelineOptions pipelineOptions; - private List> metrics = newArrayList(); - - public BenchJdbcAvroJob(final PipelineOptions pipelineOptions) { - this.pipelineOptions = pipelineOptions; - } - - public static BenchJdbcAvroJob create(final String[] cmdLineArgs) { - PipelineOptionsFactory.register(BenchJdbcAvroOptions.class); - PipelineOptions options = JdbcAvroJob.buildPipelineOptions(cmdLineArgs); - return new BenchJdbcAvroJob(options); - } - - public void run() throws Exception { - int executions = pipelineOptions.as(BenchJdbcAvroOptions.class).getExecutions(); - for (int i = 0; i < executions; i++) { - String output = - String.format("%s/run_%d", pipelineOptions.as(OutputOptions.class).getOutput(), i); - final PipelineResult pipelineResult = JdbcAvroJob.create(pipelineOptions, output).runExport(); - this.metrics.add(MetricsHelper.getMetrics(pipelineResult)); - } - System.out.println("Summary for BenchJdbcAvroJob"); - System.out.println(pipelineOptions.toString()); - System.out.println(tsvMetrics()); - } - - private String tsvMetrics() { - final List columns = - newArrayList( - "recordCount", "writeElapsedMs", "msPerMillionRows", "bytesWritten", "KbWritePerSec"); - final Collector tabJoining = Collectors.joining("\t"); - final Stream lines = - IntStream.range(0, this.metrics.size()) - .mapToObj( - i -> - String.format( - "run_%02d \t%s", - i, - columns.stream() - .map( - c -> - String.format( - "% 10d", - Optional.of(this.metrics.get(i).get(c)).orElse(0L))) - .collect(tabJoining))); - final List stats = - columns.stream() - .map( - c -> - Stats.of( - (Iterable) - this.metrics.stream().map(m -> Optional.of(m.get(c)).orElse(0L)) - ::iterator)) - .collect(Collectors.toList()); - final Map> relevantStats = - ImmutableMap.of( - "max ", Stats::max, - "mean ", Stats::mean, - "min ", Stats::min, - "stddev ", Stats::populationStandardDeviation); - final Stream statsSummary = - relevantStats.entrySet().stream() - .map( - e -> - String.format( - "%s\t%s", - e.getKey(), - stats.stream() - .map(e.getValue()) - .map(v -> String.format("% 10.1f", v)) - .collect(tabJoining))); - return String.format( - "name \t%12s\n%s", - String.join("\t", columns), - Stream.concat(lines, statsSummary).collect(Collectors.joining("\n"))); + public static BenchJdbcJob create(final String[] cmdLineArgs) { + PipelineOptionsFactory.register(BenchJdbcJob.BenchJdbcOptions.class); + return BenchJdbcJob.create( + "BenchJdbcAvroJob", + cmdLineArgs, + JdbcAvroJob::buildPipelineOptions, + (opts, output) -> JdbcAvroJob.create(opts, output).runExport()); } public static void main(String[] cmdLineArgs) { diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java new file mode 100644 index 00000000..b3728981 --- /dev/null +++ b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java @@ -0,0 +1,143 @@ +/*- + * -\-\- + * DBeam Core + * -- + * Copyright (C) 2016 - 2018 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.jobs; + +import static com.google.common.collect.Lists.newArrayList; + +import com.google.common.collect.ImmutableMap; +import com.google.common.math.Stats; +import com.spotify.dbeam.beam.MetricsHelper; +import com.spotify.dbeam.options.OutputOptions; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.function.Function; +import java.util.stream.Collector; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.options.Default; +import org.apache.beam.sdk.options.Description; +import org.apache.beam.sdk.options.PipelineOptions; + +/** Used on e2e test, allows benchmarking with different configuration parameters. */ +public class BenchJdbcJob { + + /** Functional interface for creating a job and running it, given an output path. */ + @FunctionalInterface + public interface JobRunner { + PipelineResult runExport(PipelineOptions pipelineOptions, String output) throws Exception; + } + + public interface BenchJdbcOptions extends PipelineOptions { + @Description("Number of benchmark executions.") + @Default.Integer(3) + int getExecutions(); + + void setExecutions(int value); + } + + private final String jobName; + private final PipelineOptions pipelineOptions; + private final JobRunner jobRunner; + private List> metrics = newArrayList(); + + public BenchJdbcJob( + final String jobName, final PipelineOptions pipelineOptions, final JobRunner jobRunner) { + this.jobName = jobName; + this.pipelineOptions = pipelineOptions; + this.jobRunner = jobRunner; + } + + public static BenchJdbcJob create( + final String jobName, + final String[] cmdLineArgs, + final Function optionsBuilder, + final JobRunner jobRunner) { + PipelineOptions options = optionsBuilder.apply(cmdLineArgs); + return new BenchJdbcJob(jobName, options, jobRunner); + } + + public void run() throws Exception { + int executions = pipelineOptions.as(BenchJdbcOptions.class).getExecutions(); + for (int i = 0; i < executions; i++) { + String output = + String.format("%s/run_%d", pipelineOptions.as(OutputOptions.class).getOutput(), i); + final PipelineResult pipelineResult = jobRunner.runExport(pipelineOptions, output); + this.metrics.add(MetricsHelper.getMetrics(pipelineResult)); + } + System.out.println("Summary for " + jobName); + System.out.println(pipelineOptions.toString()); + System.out.println(tsvMetrics()); + } + + String tsvMetrics() { + final List columns = + newArrayList( + "recordCount", "writeElapsedMs", "msPerMillionRows", "bytesWritten", "KbWritePerSec"); + final Collector tabJoining = Collectors.joining("\t"); + final Stream lines = + IntStream.range(0, this.metrics.size()) + .mapToObj( + i -> + String.format( + "run_%02d \t%s", + i, + columns.stream() + .map( + c -> + String.format( + "% 10d", + Optional.ofNullable(this.metrics.get(i).get(c)).orElse(0L))) + .collect(tabJoining))); + final List stats = + columns.stream() + .map( + c -> + Stats.of( + (Iterable) + this.metrics.stream().map(m -> Optional.ofNullable(m.get(c)).orElse(0L)) + ::iterator)) + .collect(Collectors.toList()); + final Map> relevantStats = + ImmutableMap.of( + "max ", Stats::max, + "mean ", Stats::mean, + "min ", Stats::min, + "stddev ", Stats::populationStandardDeviation); + final Stream statsSummary = + relevantStats.entrySet().stream() + .map( + e -> + String.format( + "%s\t%s", + e.getKey(), + stats.stream() + .map(e.getValue()) + .map(v -> String.format("% 10.1f", v)) + .collect(tabJoining))); + return String.format( + "name \t%12s\n%s", + String.join("\t", columns), + Stream.concat(lines, statsSummary).collect(Collectors.joining("\n"))); + } +} diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/jobs/JdbcAvroJob.java b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/JdbcAvroJob.java index 7d74421a..ad1aba89 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/jobs/JdbcAvroJob.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/JdbcAvroJob.java @@ -24,7 +24,7 @@ import com.spotify.dbeam.args.JdbcExportArgs; import com.spotify.dbeam.avro.BeamJdbcAvroSchema; import com.spotify.dbeam.avro.JdbcAvroIO; -import com.spotify.dbeam.avro.JdbcAvroMetering; +import com.spotify.dbeam.avro.JdbcMetering; import com.spotify.dbeam.beam.BeamHelper; import com.spotify.dbeam.beam.MetricsHelper; import com.spotify.dbeam.options.DBeamPipelineOptions; @@ -172,7 +172,7 @@ private void checkMetrics(PipelineResult pipelineResult) throws FailedValidation if (!this.dataOnly) { BeamHelper.saveMetrics(metrics, output); } - final Long recordCount = metrics.getOrDefault(JdbcAvroMetering.RECORD_COUNT_METRIC_NAME, 0L); + final Long recordCount = metrics.getOrDefault(JdbcMetering.RECORD_COUNT_METRIC_NAME, 0L); if (recordCount < this.minRows) { throw new FailedValidationException( String.format( diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/jobs/PsqlReplicationCheck.java b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/PsqlReplicationCheck.java index 5ebe007d..58f1d9e3 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/jobs/PsqlReplicationCheck.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/PsqlReplicationCheck.java @@ -53,7 +53,7 @@ public static PsqlReplicationCheck create(final JdbcExportArgs jdbcExportArgs) { return new PsqlReplicationCheck(jdbcExportArgs, REPLICATION_QUERY); } - static void validateOptions(final JdbcExportArgs jdbcExportArgs) { + public static void validateOptions(final JdbcExportArgs jdbcExportArgs) { Preconditions.checkArgument( jdbcExportArgs .jdbcAvroOptions() diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/options/ArrayHandlingMode.java b/dbeam-core/src/main/java/com/spotify/dbeam/options/ArrayHandlingMode.java index 79aafa08..35fe37ff 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/options/ArrayHandlingMode.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/options/ArrayHandlingMode.java @@ -32,7 +32,8 @@ public static String validateValue(String value) { List possibleValues = Arrays.asList(Bytes, TypedMetaFromFirstRow, TypedMetaPostgres); if (value == null || !possibleValues.contains(value)) { throw new RuntimeException( - String.format("Invalid value '%s' for array handling mode. Allowed values: %s", + String.format( + "Invalid value '%s' for array handling mode. Allowed values: %s", value, possibleValues)); } diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/options/JdbcExportArgsFactory.java b/dbeam-core/src/main/java/com/spotify/dbeam/options/JdbcExportArgsFactory.java index 7b3c725c..1c4a06fd 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/options/JdbcExportArgsFactory.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/options/JdbcExportArgsFactory.java @@ -70,8 +70,7 @@ public static JdbcExportArgs fromPipelineOptions(final PipelineOptions options) exportOptions.getAvroCodec(), Optional.ofNullable(exportOptions.getPreCommand()).orElse(Collections.emptyList()), ArrayHandlingMode.validateValue(exportOptions.getArrayMode()), - exportOptions.isNullableArrayItems() - ); + exportOptions.isNullableArrayItems()); return JdbcExportArgs.create( jdbcAvroArgs, diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/TestHelper.java b/dbeam-core/src/test/java/com/spotify/dbeam/TestHelper.java index 58985902..c5e4bcfd 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/TestHelper.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/TestHelper.java @@ -39,7 +39,6 @@ import java.util.UUID; import java.util.stream.Collectors; import org.mockito.Mockito; -import org.postgresql.jdbc.PgArray; public class TestHelper { @@ -67,13 +66,18 @@ public static UUID byteBufferToUuid(final ByteBuffer byteBuffer) { return new UUID(high, low); } - public static void mockArrayColumn(ResultSetMetaData meta, ResultSet resultSet, - int columnIdx, String columnName, - String columnTypeName, int arrayType, String arrayTypeName, - Object array1, Object... arrays) + public static void mockArrayColumn( + ResultSetMetaData meta, + ResultSet resultSet, + int columnIdx, + String columnName, + String columnTypeName, + int arrayType, + String arrayTypeName, + Object array1, + Object... arrays) throws SQLException { - mockResultSetMeta(meta, columnIdx, Types.ARRAY, columnName, "java.sql.Array", - columnTypeName); + mockResultSetMeta(meta, columnIdx, Types.ARRAY, columnName, "java.sql.Array", columnTypeName); Array res1; if (array1 == null) { res1 = null; @@ -93,9 +97,13 @@ public static void mockArrayColumn(ResultSetMetaData meta, ResultSet resultSet, when(resultSet.getArray(columnIdx)).thenReturn(res1, resX); } - public static void mockResultSetMeta(ResultSetMetaData meta, int columnIdx, int columnType, - String columnName, - String columnClassName, String columnTypeName) + public static void mockResultSetMeta( + ResultSetMetaData meta, + int columnIdx, + int columnType, + String columnName, + String columnClassName, + String columnTypeName) throws SQLException { when(meta.getColumnType(columnIdx)).thenReturn(columnType); when(meta.getColumnName(columnIdx)).thenReturn(columnName); diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/args/JdbcExportOptionsTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/args/JdbcExportOptionsTest.java index d6a0faf4..d87062a3 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/args/JdbcExportOptionsTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/args/JdbcExportOptionsTest.java @@ -33,21 +33,21 @@ import org.apache.avro.file.CodecFactory; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.junit.AfterClass; -import org.junit.Assert; -import org.junit.BeforeClass; -import org.junit.Test; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; public class JdbcExportOptionsTest { private static File sqlFile; - @BeforeClass + @BeforeAll public static void beforeAll() throws IOException { sqlFile = File.createTempFile("query", ".sql"); sqlFile.deleteOnExit(); } - @AfterClass + @AfterAll public static void afterAll() throws IOException { Files.delete(sqlFile.toPath()); } @@ -63,27 +63,33 @@ JdbcExportArgs optionsFromArgs(String[] cmdLineArgs) throws IOException, ClassNo return JdbcExportArgsFactory.fromPipelineOptions(opts); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailParseOnInvalidArg() throws IOException, ClassNotFoundException { - optionsFromArgs("--foo=bar"); + @Test + public void shouldFailParseOnInvalidArg() { + Assertions.assertThrows(IllegalArgumentException.class, () -> optionsFromArgs("--foo=bar")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnMissingConnectionUrl() throws IOException, ClassNotFoundException { - optionsFromArgs("--table=sometable"); + @Test + public void shouldFailOnMissingConnectionUrl() { + Assertions.assertThrows( + IllegalArgumentException.class, () -> optionsFromArgs("--table=sometable")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnMissingTableAndSqlFile() throws IOException, ClassNotFoundException { - optionsFromArgs("--connectionUrl=jdbc:postgresql://some_db"); + @Test + public void shouldFailOnMissingTableAndSqlFile() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> optionsFromArgs("--connectionUrl=jdbc:postgresql://some_db")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnTableAndSqlFilePresent() throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db --sqlFile=" - + sqlFile.getAbsolutePath() - + " --table=some_table"); + @Test + public void shouldFailOnTableAndSqlFilePresent() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db --sqlFile=" + + sqlFile.getAbsolutePath() + + " --table=some_table")); } @Test @@ -100,7 +106,7 @@ public void shouldNotFailOnMissingTableSqlFile() throws IOException, ClassNotFou QueryBuilderArgs.createFromQuery( com.google.common.io.Files.asCharSource(sqlFile, StandardCharsets.UTF_8).read())); - Assert.assertEquals(expected, actual); + Assertions.assertEquals(expected, actual); } @Test @@ -116,7 +122,7 @@ public void shouldParseWithDefaultsOnConnectionUrlAndTable() .withUsername("dbeam-extractor")), QueryBuilderArgs.create("some_table")); - Assert.assertEquals(expected, actual); + Assertions.assertEquals(expected, actual); } @Test @@ -130,56 +136,77 @@ public void shouldParseWithMySqlConnection() throws IOException, ClassNotFoundEx JdbcConnectionArgs.create("jdbc:mysql://some_db").withUsername("dbeam-extractor")), QueryBuilderArgs.create("some_table")); - Assert.assertEquals(expected, actual); + Assertions.assertEquals(expected, actual); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnInvalidTable() throws IOException, ClassNotFoundException { - optionsFromArgs("--connectionUrl=jdbc:postgresql://some_db --table=some-table-with-dash"); + @Test + public void shouldFailOnInvalidTable() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db --table=some-table-with-dash")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnNonJdbcUrl() throws IOException, ClassNotFoundException { - optionsFromArgs("--connectionUrl=bar --table=sometable"); + @Test + public void shouldFailOnNonJdbcUrl() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> optionsFromArgs("--connectionUrl=bar --table=sometable")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnUnsupportedJdbcUrl() throws IOException, ClassNotFoundException { - optionsFromArgs("--connectionUrl=jdbc:paradox:./foo --table=sometable"); + @Test + public void shouldFailOnUnsupportedJdbcUrl() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> optionsFromArgs("--connectionUrl=jdbc:paradox:./foo --table=sometable")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnMissingPartitionButPresentPartitionColumn() - throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db --table=sometable " + "--partitionColumn=col"); + @Test + public void shouldFailOnMissingPartitionButPresentPartitionColumn() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db --table=sometable " + + "--partitionColumn=col")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnTooOldPartition() throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db --table=sometable " + "--partition=2015-01-01"); + @Test + public void shouldFailOnTooOldPartition() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db --table=sometable " + + "--partition=2015-01-01")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnTooOldPartitionWithConfiguredMinPartitionPeriodMoreThanPartition() - throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db --table=sometable " - + "--partition=2015-01-01 --minPartitionPeriod=2015-01-02"); + @Test + public void shouldFailOnTooOldPartitionWithConfiguredMinPartitionPeriodMoreThanPartition() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db --table=sometable " + + "--partition=2015-01-01 --minPartitionPeriod=2015-01-02")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnTooOldPartitionWithConfiguredMinPartitionPeriodLessThanPartition() - throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db --table=sometable " - + "--partition=2015-01-01 --minPartitionPeriod=2015-01-01"); + @Test + public void shouldFailOnTooOldPartitionWithConfiguredMinPartitionPeriodLessThanPartition() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db --table=sometable " + + "--partition=2015-01-01 --minPartitionPeriod=2015-01-01")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnNonJdbcUrl2() throws IOException, ClassNotFoundException { - optionsFromArgs("--connectionUrl=some:foo:bar --table=sometable"); + @Test + public void shouldFailOnNonJdbcUrl2() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> optionsFromArgs("--connectionUrl=some:foo:bar --table=sometable")); } @Test @@ -197,7 +224,7 @@ public void shouldConfigureUserAndPassword() throws IOException, ClassNotFoundEx .withPassword("somepassword")), QueryBuilderArgs.create("some_table")); - Assert.assertEquals(expected, actual); + Assertions.assertEquals(expected, actual); } @Test @@ -207,7 +234,7 @@ public void shouldConfigureAvroLogicalTypes() throws IOException, ClassNotFoundE "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--password=secret --useAvroLogicalTypes=true"); - Assert.assertTrue(options.useAvroLogicalTypes()); + Assertions.assertTrue(options.useAvroLogicalTypes()); } @Test @@ -217,7 +244,7 @@ public void shouldConfigureAvroDoc() throws IOException, ClassNotFoundException "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--password=secret --avroDoc=somedoc"); - Assert.assertEquals(Optional.of("somedoc"), options.avroDoc()); + Assertions.assertEquals(Optional.of("somedoc"), options.avroDoc()); } @Test @@ -227,7 +254,7 @@ public void shouldConfigureAvroSchemaNamespace() throws IOException, ClassNotFou "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--password=secret --avroSchemaNamespace=ns"); - Assert.assertEquals("ns", options.avroSchemaNamespace()); + Assertions.assertEquals("ns", options.avroSchemaNamespace()); } @Test @@ -238,7 +265,7 @@ public void shouldDefaultOutputDataOnlyToFalse() throws IOException, ClassNotFou "--connectionUrl=jdbc:postgresql://some_db", "--table=some_table", "--password=secret" }); - Assert.assertEquals(false, defaultDataOnlyoptions.as(OutputOptions.class).getDataOnly()); + Assertions.assertEquals(false, defaultDataOnlyoptions.as(OutputOptions.class).getDataOnly()); } @Test @@ -252,7 +279,7 @@ public void shouldConfigureOutputDataOnly() throws IOException, ClassNotFoundExc "--dataOnly" }); - Assert.assertEquals(true, options.as(OutputOptions.class).getDataOnly()); + Assertions.assertEquals(true, options.as(OutputOptions.class).getDataOnly()); } @Test @@ -262,7 +289,7 @@ public void shouldConfigureFetchSize() throws IOException, ClassNotFoundExceptio "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--password=secret --fetchSize=1234"); - Assert.assertEquals(1234, options.jdbcAvroOptions().fetchSize()); + Assertions.assertEquals(1234, options.jdbcAvroOptions().fetchSize()); } @Test @@ -274,7 +301,7 @@ public void shouldSupportMonthlyPartitionPeriod() throws IOException, ClassNotFo "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--password=secret --partitionPeriod=P1M --partition=2050-12"); - Assert.assertEquals(Period.ofMonths(1), options.queryBuilderArgs().partitionPeriod()); + Assertions.assertEquals(Period.ofMonths(1), options.queryBuilderArgs().partitionPeriod()); } @Test @@ -284,8 +311,8 @@ public void shouldConfigureDeflateCodec() throws IOException, ClassNotFoundExcep "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--password=secret --avroCodec=deflate7"); - Assert.assertEquals("deflate7", options.jdbcAvroOptions().avroCodec()); - Assert.assertEquals( + Assertions.assertEquals("deflate7", options.jdbcAvroOptions().avroCodec()); + Assertions.assertEquals( CodecFactory.deflateCodec(7).toString(), options.jdbcAvroOptions().getCodecFactory().toString()); } @@ -297,8 +324,8 @@ public void shouldConfigureZstandardCodec() throws IOException, ClassNotFoundExc "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--password=secret --avroCodec=zstandard9"); - Assert.assertEquals("zstandard9", options.jdbcAvroOptions().avroCodec()); - Assert.assertEquals( + Assertions.assertEquals("zstandard9", options.jdbcAvroOptions().avroCodec()); + Assertions.assertEquals( CodecFactory.zstandardCodec(9).toString(), options.jdbcAvroOptions().getCodecFactory().toString()); } @@ -310,8 +337,8 @@ public void shouldConfigureSnappyCodec() throws IOException, ClassNotFoundExcept "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--password=secret --avroCodec=snappy"); - Assert.assertEquals("snappy", options.jdbcAvroOptions().avroCodec()); - Assert.assertEquals( + Assertions.assertEquals("snappy", options.jdbcAvroOptions().avroCodec()); + Assertions.assertEquals( CodecFactory.snappyCodec().toString(), options.jdbcAvroOptions().getCodecFactory().toString()); } @@ -328,44 +355,58 @@ public void shouldConfiguraPreCommands() throws IOException, ClassNotFoundExcept "--preCommand=set bar=2" }); - Assert.assertEquals("set foo='1'", options.jdbcAvroOptions().preCommand().get(0)); - Assert.assertEquals("set bar=2", options.jdbcAvroOptions().preCommand().get(1)); + Assertions.assertEquals("set foo='1'", options.jdbcAvroOptions().preCommand().get(0)); + Assertions.assertEquals("set bar=2", options.jdbcAvroOptions().preCommand().get(1)); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnInvalidAvroCodec() throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db --table=some_table " - + "--password=secret --avroCodec=lzma"); + @Test + public void shouldFailOnInvalidAvroCodec() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + + "--password=secret --avroCodec=lzma")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnQueryParallelismWithNoSplitColumn() - throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db " - + "--table=some_table --password=secret --queryParallelism=10"); + @Test + public void shouldFailOnQueryParallelismWithNoSplitColumn() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db " + + "--table=some_table --password=secret --queryParallelism=10")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnSplitColumnWithNoQueryParallelism() - throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db " - + "--table=some_table --password=secret --splitColumn=id"); + @Test + public void shouldFailOnSplitColumnWithNoQueryParallelism() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db " + + "--table=some_table --password=secret --splitColumn=id")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnZeroQueryParallelism() throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db " - + "--table=some_table --password=secret --queryParallelism=0 --splitColumn=id"); + @Test + public void shouldFailOnZeroQueryParallelism() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db " + + "--table=some_table --password=secret" + + " --queryParallelism=0 --splitColumn=id")); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnNegativeQueryParallelism() throws IOException, ClassNotFoundException { - optionsFromArgs( - "--connectionUrl=jdbc:postgresql://some_db --table=some_table " - + "--password=secret --queryParallelism=-5 --splitColumn=id"); + @Test + public void shouldFailOnNegativeQueryParallelism() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + optionsFromArgs( + "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + + "--password=secret --queryParallelism=-5 --splitColumn=id")); } } diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/args/ParallelQueryBuilderTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/args/ParallelQueryBuilderTest.java index e21d8423..b8ce49af 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/args/ParallelQueryBuilderTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/args/ParallelQueryBuilderTest.java @@ -26,8 +26,8 @@ import com.google.common.collect.Lists; import java.util.List; import org.hamcrest.Matchers; -import org.junit.Ignore; -import org.junit.Test; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; public class ParallelQueryBuilderTest { @@ -91,7 +91,7 @@ public void shouldBuildSingleQueryWhenParallelismIsOne() { Matchers.is(Lists.newArrayList(format("%s AND sp >= %s AND sp <= %s", QUERY_BASE, 1, 10)))); } - @Ignore // TODO: fix this + @Disabled("TODO: fix this") @Test public void shouldBuildMultipleQueriesWhenQueryingFromTwoRows() { final List actual = ParallelQueryBuilder.queriesForBounds(1, 2, 2, "sp", QUERY_FORMAT); diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/args/QueryBuilderArgsTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/args/QueryBuilderArgsTest.java index 0012770a..06470246 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/args/QueryBuilderArgsTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/args/QueryBuilderArgsTest.java @@ -34,10 +34,10 @@ import java.time.Instant; import java.util.Optional; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.junit.AfterClass; -import org.junit.Assert; -import org.junit.BeforeClass; -import org.junit.Test; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; public class QueryBuilderArgsTest { @@ -46,7 +46,7 @@ public class QueryBuilderArgsTest { private static Connection connection; private static Path coffeesSqlQueryPath; - @BeforeClass + @BeforeAll public static void beforeAll() throws SQLException, ClassNotFoundException, IOException { coffeesSqlQueryPath = TestHelper.createTmpDirPath("jdbc-export-args-test").resolve("coffees_query_1.sql"); @@ -57,19 +57,20 @@ public static void beforeAll() throws SQLException, ClassNotFoundException, IOEx DbTestHelper.createFixtures(CONNECTION_URL); } - @AfterClass + @AfterAll public static void afterAll() throws SQLException { connection.close(); } - @Test(expected = IllegalArgumentException.class) + @Test public void shouldFailOnNullTableName() { - QueryBuilderArgs.create(null); + Assertions.assertThrows(IllegalArgumentException.class, () -> QueryBuilderArgs.create(null)); } - @Test(expected = IllegalArgumentException.class) + @Test public void shouldFailOnInvalidTableName() { - QueryBuilderArgs.create("*invalid#name@!"); + Assertions.assertThrows( + IllegalArgumentException.class, () -> QueryBuilderArgs.create("*invalid#name@!")); } public void shouldNotFailOnTableNameWithDots() { @@ -80,7 +81,7 @@ public void shouldNotFailOnTableNameWithDots() { public void shouldCreateValidSqlQueryFromUserQuery() throws SQLException { final QueryBuilderArgs args = QueryBuilderArgs.createFromQuery("SELECT * FROM some_table"); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList("SELECT * FROM (SELECT * FROM some_table) as user_sql_query WHERE 1=1"), args.buildQueries(null)); } @@ -90,7 +91,7 @@ public void shouldConfigureLimit() throws IOException, SQLException { final QueryBuilderArgs actual = parseOptions("--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--limit=7"); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList("SELECT * FROM some_table WHERE 1=1 LIMIT 7"), actual.buildQueries(null)); } @@ -102,8 +103,8 @@ public void shouldConfigurePartition() throws IOException, SQLException { "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--partition=2027-07-31"); - Assert.assertEquals(Optional.of(Instant.parse("2027-07-31T00:00:00Z")), actual.partition()); - Assert.assertEquals( + Assertions.assertEquals(Optional.of(Instant.parse("2027-07-31T00:00:00Z")), actual.partition()); + Assertions.assertEquals( Lists.newArrayList("SELECT * FROM some_table WHERE 1=1"), actual.buildQueries(null)); } @@ -114,7 +115,7 @@ public void shouldConfigurePartitionForFullIsoString() throws IOException, SQLEx "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--partition=2027-07-31T13:37:59Z"); - Assert.assertEquals(Optional.of(Instant.parse("2027-07-31T13:37:59Z")), actual.partition()); + Assertions.assertEquals(Optional.of(Instant.parse("2027-07-31T13:37:59Z")), actual.partition()); } @Test @@ -124,7 +125,7 @@ public void shouldConfigurePartitionForMonthlySchedule() throws IOException, SQL "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--partition=2027-05"); - Assert.assertEquals(Optional.of(Instant.parse("2027-05-01T00:00:00Z")), actual.partition()); + Assertions.assertEquals(Optional.of(Instant.parse("2027-05-01T00:00:00Z")), actual.partition()); } @Test @@ -134,7 +135,7 @@ public void shouldConfigurePartitionForHourlySchedule() throws IOException, SQLE "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--partition=2027-05-02T23"); - Assert.assertEquals(Optional.of(Instant.parse("2027-05-02T23:00:00Z")), actual.partition()); + Assertions.assertEquals(Optional.of(Instant.parse("2027-05-02T23:00:00Z")), actual.partition()); } @Test @@ -144,7 +145,7 @@ public void shouldConfigurePartitionColumn() throws IOException, SQLException { "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--partition=2027-07-31 --partitionColumn=col"); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList( "SELECT * FROM some_table WHERE 1=1 " + "AND col >= '2027-07-31' AND col < '2027-08-01'"), @@ -158,7 +159,7 @@ public void shouldConfigurePartitionColumnAndLimit() throws IOException, SQLExce "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--partition=2027-07-31 --partitionColumn=col --limit=5"); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList( "SELECT * FROM some_table WHERE 1=1 " + "AND col >= '2027-07-31' AND col < '2027-08-01' LIMIT 5"), @@ -172,7 +173,7 @@ public void shouldConfigurePartitionColumnAndPartitionPeriod() throws IOExceptio "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--partition=2027-07-31 --partitionColumn=col --partitionPeriod=P1M"); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList( "SELECT * FROM some_table WHERE 1=1 " + "AND col >= '2027-07-31' AND col < '2027-08-31'"), @@ -187,7 +188,7 @@ public void shouldConfigurePartitionColumnAndPartitionPeriodForHourly() "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--partition=2027-07-31T00 --partitionColumn=col --partitionPeriod=PT1H"); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList( "SELECT * FROM some_table WHERE 1=1 " + "AND col >= '2027-07-31T00:00:00Z' AND col < '2027-07-31T01:00:00Z'"), @@ -204,7 +205,7 @@ public void shouldConfigureLimitForSqlFile() throws IOException, SQLException { "--connectionUrl=jdbc:postgresql://some_db " + "--sqlFile=%s --limit=7", coffeesSqlQueryPath.toString())); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList( "SELECT * FROM (SELECT * FROM COFFEES WHERE SIZE > 10) as user_sql_query" + " WHERE 1=1 LIMIT 7"), @@ -220,7 +221,7 @@ public void shouldConfigurePartitionColumnAndLimitForSqlFile() throws IOExceptio + "--sqlFile=%s --partition=2027-07-31 --partitionColumn=col --limit=7", coffeesSqlQueryPath.toString())); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList( "SELECT * FROM (SELECT * FROM COFFEES WHERE SIZE > 10) as user_sql_query WHERE 1=1 " + "AND col >= '2027-07-31' AND col < '2027-08-01' LIMIT 7"), @@ -238,7 +239,7 @@ public void shouldConfigurePartitionColumnAndPartitionPeriodForSqlFile() + "--partitionColumn=col --partitionPeriod=P1M", coffeesSqlQueryPath.toString())); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList( "SELECT * FROM (SELECT * FROM COFFEES WHERE SIZE > 10) as user_sql_query WHERE 1=1 " + "AND col >= '2027-07-31' AND col < '2027-08-31'"), @@ -253,7 +254,7 @@ public void shouldCreateParallelQueries() throws IOException, SQLException { "--connectionUrl=jdbc:postgresql://some_db --table=COFFEES " + "--splitColumn=ROWNUM --queryParallelism=5"); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList("SELECT * FROM COFFEES WHERE 1=1" + " AND ROWNUM >= 1 AND ROWNUM <= 2"), actual.buildQueries(connection)); } @@ -267,7 +268,7 @@ public void shouldCreateParallelQueriesWithSqlFile() throws IOException, SQLExce + "--sqlFile=%s --splitColumn=ROWNUM --queryParallelism=5", coffeesSqlQueryPath.toString())); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList( "SELECT * FROM (SELECT * FROM COFFEES WHERE SIZE > 10) as user_sql_query WHERE 1=1" + " AND ROWNUM >= 1 AND ROWNUM <= 2"), @@ -284,7 +285,7 @@ public void shouldCreateParallelQueriesWithPartitionColumn() throws IOException, + "--partitionColumn=col --partitionPeriod=P1M --limit=7", coffeesSqlQueryPath.toString())); - Assert.assertEquals( + Assertions.assertEquals( Lists.newArrayList( "SELECT * FROM (SELECT * FROM COFFEES WHERE SIZE > 10) as user_sql_query WHERE 1=1" + " AND col >= '2027-07-31' AND col < '2027-08-31' LIMIT 7"), diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/args/QueryBuilderTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/args/QueryBuilderTest.java index 5abf0687..e82e88ba 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/args/QueryBuilderTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/args/QueryBuilderTest.java @@ -22,8 +22,8 @@ import java.util.Arrays; import java.util.List; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; public class QueryBuilderTest { @@ -33,7 +33,7 @@ public void testCtorFromTable() { final String expected = "SELECT * FROM abc WHERE 1=1"; - Assert.assertEquals(expected, wrapper.build()); + Assertions.assertEquals(expected, wrapper.build()); } @Test @@ -42,7 +42,7 @@ public void testCtorRawSqlWithoutWhere() { final String expected = "SELECT * FROM (SELECT * FROM t1) as user_sql_query WHERE 1=1"; - Assert.assertEquals(expected, wrapper.build()); + Assertions.assertEquals(expected, wrapper.build()); } @Test @@ -50,7 +50,7 @@ public void testCtorCopyWithConditionNotEquals() { final QueryBuilder q1 = QueryBuilder.fromSqlQuery("SELECT * FROM t1"); final QueryBuilder copy = q1.withPartitionCondition("pary", "20180101", "20180201"); - Assert.assertNotEquals(q1.build(), copy.build()); + Assertions.assertNotEquals(q1.build(), copy.build()); } @Test @@ -58,7 +58,7 @@ public void testCtorCopyWithLimitNotEquals() { final QueryBuilder q1 = QueryBuilder.fromSqlQuery("SELECT * FROM t1"); final QueryBuilder copy = q1.withLimit(3L); - Assert.assertNotEquals(q1.build(), copy.build()); + Assertions.assertNotEquals(q1.build(), copy.build()); } @Test @@ -68,7 +68,7 @@ public void testCtorRawSqlWithWhere() { final String expected = "SELECT * FROM (SELECT * FROM t1 WHERE a > 100) as user_sql_query WHERE 1=1"; - Assert.assertEquals(expected, wrapper.build()); + Assertions.assertEquals(expected, wrapper.build()); } @Test @@ -78,7 +78,7 @@ public void testRawSqlWithLimit() { final String expected = "SELECT * FROM (SELECT * FROM t1) as user_sql_query WHERE 1=1 LIMIT 102"; - Assert.assertEquals(expected, wrapper.build()); + Assertions.assertEquals(expected, wrapper.build()); } @Test @@ -91,7 +91,7 @@ public void testRawSqlwithParallelization() { "SELECT * FROM (SELECT * FROM t1) as user_sql_query" + " WHERE 1=1 AND bucket >= 10 AND bucket < 20"; - Assert.assertEquals(expected, wrapper.build()); + Assertions.assertEquals(expected, wrapper.build()); } @Test @@ -104,7 +104,7 @@ public void testRawSqlWithPartition() { "SELECT * FROM (SELECT * FROM t1) as user_sql_query WHERE 1=1" + " AND birthDate >= '2018-01-01' AND birthDate < '2018-02-01'"; - Assert.assertEquals(expected, wrapper.build()); + Assertions.assertEquals(expected, wrapper.build()); } @Test @@ -117,7 +117,7 @@ public void testRawSqlMultiline() { "SELECT * FROM (SELECT a, b, c FROM t1\n WHERE total > 100\n AND country = 262\n)" + " as user_sql_query WHERE 1=1"; - Assert.assertEquals(expected, wrapper.build()); + Assertions.assertEquals(expected, wrapper.build()); } @Test @@ -131,7 +131,7 @@ public void testRawSqlWithComments() { + "-- We perform initial query here\nSELECT a, b, c FROM t1\n WHERE total > 100)" + " as user_sql_query WHERE 1=1"; - Assert.assertEquals(expected, wrapper.build()); + Assertions.assertEquals(expected, wrapper.build()); } @Test @@ -161,7 +161,7 @@ public void testRawSqlWithCte() { + "GROUP BY date\n" + ") as user_sql_query WHERE 1=1"; - Assert.assertEquals(expected, wrapper.build()); + Assertions.assertEquals(expected, wrapper.build()); } @Test @@ -207,12 +207,12 @@ public void testItGeneratesQueryForLimits() { .withPartitionCondition("partition", "a", "d") .generateQueryToGetLimitsOfSplitColumn("splitCol", "mixy", "maxy") .build(); - Assert.assertEquals(expected, actual); + Assertions.assertEquals(expected, actual); } private void execAndCompare(String rawInput, String expected) { final String actual = QueryBuilder.fromSqlQuery(rawInput).build(); - Assert.assertEquals(expected, actual); + Assertions.assertEquals(expected, actual); } } diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/avro/JdbcAvroRecordTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/avro/JdbcAvroRecordTest.java index f9906186..c07ba782 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/avro/JdbcAvroRecordTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/avro/JdbcAvroRecordTest.java @@ -28,21 +28,16 @@ import com.spotify.dbeam.TestHelper; import com.spotify.dbeam.args.QueryBuilderArgs; import com.spotify.dbeam.options.ArrayHandlingMode; -import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.ByteBuffer; -import java.nio.charset.StandardCharsets; -import java.sql.Array; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Timestamp; -import java.sql.Types; import java.util.ArrayList; import java.util.List; import java.util.Optional; -import java.util.UUID; import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.avro.Schema; @@ -53,22 +48,18 @@ import org.apache.avro.generic.GenericDatumReader; import org.apache.avro.generic.GenericDatumWriter; import org.apache.avro.generic.GenericRecord; -import org.apache.avro.io.DatumReader; -import org.apache.avro.io.Decoder; -import org.apache.avro.io.DecoderFactory; import org.apache.avro.util.Utf8; -import org.junit.Assert; -import org.junit.BeforeClass; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; import org.mockito.Mockito; -import org.postgresql.jdbc.PgArray; public class JdbcAvroRecordTest { private static String CONNECTION_URL = "jdbc:h2:mem:test;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; - @BeforeClass + @BeforeAll public static void beforeAll() throws SQLException, ClassNotFoundException { DbTestHelper.createFixtures(CONNECTION_URL); } @@ -83,16 +74,18 @@ public void shouldCreateSchema() throws ClassNotFoundException, SQLException { "dbeam_generated", Optional.empty(), "Generate schema from JDBC ResultSet from COFFEES jdbc:h2:mem:test", - false, ArrayHandlingMode.TypedMetaFromFirstRow, false); + false, + ArrayHandlingMode.TypedMetaFromFirstRow, + false); - Assert.assertNotNull(actual); - Assert.assertEquals("dbeam_generated", actual.getNamespace()); - Assert.assertEquals("COFFEES", actual.getProp("tableName")); - Assert.assertEquals("jdbc:h2:mem:test", actual.getProp("connectionUrl")); - Assert.assertEquals( + Assertions.assertNotNull(actual); + Assertions.assertEquals("dbeam_generated", actual.getNamespace()); + Assertions.assertEquals("COFFEES", actual.getProp("tableName")); + Assertions.assertEquals("jdbc:h2:mem:test", actual.getProp("connectionUrl")); + Assertions.assertEquals( "Generate schema from JDBC ResultSet from COFFEES jdbc:h2:mem:test", actual.getDoc()); - Assert.assertEquals(fieldCount, actual.getFields().size()); - Assert.assertEquals( + Assertions.assertEquals(fieldCount, actual.getFields().size()); + Assertions.assertEquals( Lists.newArrayList( "COF_NAME", "SUP_ID", @@ -110,45 +103,46 @@ public void shouldCreateSchema() throws ClassNotFoundException, SQLException { "TEXT_ARR"), actual.getFields().stream().map(Schema.Field::name).collect(Collectors.toList())); for (Schema.Field f : actual.getFields()) { - Assert.assertEquals(Schema.Type.UNION, f.schema().getType()); - Assert.assertEquals(2, f.schema().getTypes().size()); - Assert.assertEquals(Schema.Type.NULL, f.schema().getTypes().get(0).getType()); + Assertions.assertEquals(Schema.Type.UNION, f.schema().getType()); + Assertions.assertEquals(2, f.schema().getTypes().size()); + Assertions.assertEquals(Schema.Type.NULL, f.schema().getTypes().get(0).getType()); } - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.STRING, actual.getField("COF_NAME").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.INT, actual.getField("SUP_ID").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.STRING, actual.getField("PRICE").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.FLOAT, actual.getField("TEMPERATURE").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.DOUBLE, actual.getField("SIZE").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.BOOLEAN, actual.getField("IS_ARABIC").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.INT, actual.getField("SALES").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.LONG, actual.getField("TOTAL").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.LONG, actual.getField("CREATED").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.LONG, actual.getField("UPDATED").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.BYTES, actual.getField("UID").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.LONG, actual.getField("ROWNUM").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.ARRAY, actual.getField("INT_ARR").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.INT, actual.getField("INT_ARR").schema().getTypes().get(1).getElementType().getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.ARRAY, actual.getField("TEXT_ARR").schema().getTypes().get(1).getType()); - Assert.assertEquals( + Assertions.assertEquals( Schema.Type.STRING, actual.getField("TEXT_ARR").schema().getTypes().get(1).getElementType().getType()); - Assert.assertNull(actual.getField("UPDATED").schema().getTypes().get(1).getProp("logicalType")); + Assertions.assertNull( + actual.getField("UPDATED").schema().getTypes().get(1).getProp("logicalType")); } @Test @@ -161,10 +155,12 @@ public void shouldCreateSchemaWithLogicalTypes() throws ClassNotFoundException, "dbeam_generated", Optional.empty(), "Generate schema from JDBC ResultSet from COFFEES jdbc:h2:mem:test", - true, ArrayHandlingMode.TypedMetaFromFirstRow, false); + true, + ArrayHandlingMode.TypedMetaFromFirstRow, + false); - Assert.assertEquals(fieldCount, actual.getFields().size()); - Assert.assertEquals( + Assertions.assertEquals(fieldCount, actual.getFields().size()); + Assertions.assertEquals( "timestamp-millis", actual.getField("UPDATED").schema().getTypes().get(1).getProp("logicalType")); } @@ -178,9 +174,11 @@ public void shouldCreateSchemaWithCustomSchemaName() throws ClassNotFoundExcepti "dbeam_generated", Optional.of("CustomSchemaName"), "Generate schema from JDBC ResultSet from COFFEES jdbc:h2:mem:test", - false, ArrayHandlingMode.TypedMetaFromFirstRow, false); + false, + ArrayHandlingMode.TypedMetaFromFirstRow, + false); - Assert.assertEquals("CustomSchemaName", actual.getName()); + Assertions.assertEquals("CustomSchemaName", actual.getName()); } @Test @@ -195,8 +193,7 @@ public void shouldEncodeResultSetToValidAvro() String arrayMode = ArrayHandlingMode.TypedMetaFromFirstRow; final Schema schema = JdbcAvroSchema.createAvroSchema( - rs, "dbeam_generated", "connection", Optional.empty(), "doc", - false, arrayMode, false); + rs, "dbeam_generated", "connection", Optional.empty(), "doc", false, arrayMode, false); final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(rs, arrayMode, false); final DataFileWriter dataFileWriter = new DataFileWriter<>(new GenericDatumWriter<>(schema)); @@ -217,15 +214,15 @@ public void shouldEncodeResultSetToValidAvro() final List records = StreamSupport.stream(dataFileReader.spliterator(), false).collect(Collectors.toList()); - Assert.assertEquals(2, records.size()); + Assertions.assertEquals(2, records.size()); final GenericRecord record = records.stream() .filter(r -> Coffee.COFFEE1.name().equals(r.get(0).toString())) .findFirst() .orElseThrow(() -> new IllegalArgumentException("not found")); - Assert.assertEquals(14, record.getSchema().getFields().size()); - Assert.assertEquals(schema, record.getSchema()); + Assertions.assertEquals(14, record.getSchema().getFields().size()); + Assertions.assertEquals(schema, record.getSchema()); List actualTxtArray = ((GenericData.Array) record.get(13)) .stream().map(x -> x.toString()).collect(Collectors.toList()); @@ -245,7 +242,7 @@ public void shouldEncodeResultSetToValidAvro() (Long) record.get(11), new ArrayList<>((GenericData.Array) record.get(12)), actualTxtArray); - Assert.assertEquals(Coffee.COFFEE1, actual); + Assertions.assertEquals(Coffee.COFFEE1, actual); } @Test @@ -267,7 +264,7 @@ public void shouldCorrectlyEncodeUnsignedIntToAvroLong() throws SQLException { when(resultSet.getLong(columnNum)).thenReturn(valueUnderTest); final Object result = mapping.apply(resultSet); - Assert.assertEquals(Long.class, result.getClass()); - Assert.assertEquals(valueUnderTest, ((Long) result).longValue()); + Assertions.assertEquals(Long.class, result.getClass()); + Assertions.assertEquals(valueUnderTest, ((Long) result).longValue()); } } diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/avro/JdbcAvroSchemaTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/avro/JdbcAvroSchemaTest.java index f981c16f..cfb5d22a 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/avro/JdbcAvroSchemaTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/avro/JdbcAvroSchemaTest.java @@ -29,8 +29,8 @@ import java.sql.Types; import java.util.Optional; import org.apache.avro.Schema; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; import org.mockito.Mockito; public class JdbcAvroSchemaTest { @@ -43,7 +43,7 @@ public void shouldGetDatabaseTableNameFromMetaData() throws SQLException { when(meta.getColumnCount()).thenReturn(1); when(meta.getTableName(1)).thenReturn("test_table"); - Assert.assertEquals("test_table", JdbcAvroSchema.getDatabaseTableName(meta)); + Assertions.assertEquals("test_table", JdbcAvroSchema.getDatabaseTableName(meta)); } @Test @@ -52,7 +52,7 @@ public void shouldDefaultTableNameWhenMetaDataHasEmptyTableName() throws SQLExce when(meta.getColumnCount()).thenReturn(1); when(meta.getTableName(1)).thenReturn(""); - Assert.assertEquals("no_table_name", JdbcAvroSchema.getDatabaseTableName(meta)); + Assertions.assertEquals("no_table_name", JdbcAvroSchema.getDatabaseTableName(meta)); } @Test @@ -61,7 +61,7 @@ public void shouldDefaultTableNameWhenMetaDataHasNullTableName() throws SQLExcep when(meta.getColumnCount()).thenReturn(1); when(meta.getTableName(1)).thenReturn(null); - Assert.assertEquals("no_table_name", JdbcAvroSchema.getDatabaseTableName(meta)); + Assertions.assertEquals("no_table_name", JdbcAvroSchema.getDatabaseTableName(meta)); } @Test @@ -71,7 +71,7 @@ public void shouldGetDatabaseTableNameFromFirstNonNullMetaData() throws SQLExcep when(meta.getTableName(1)).thenReturn(""); when(meta.getTableName(2)).thenReturn("test_table"); - Assert.assertEquals("test_table", JdbcAvroSchema.getDatabaseTableName(meta)); + Assertions.assertEquals("test_table", JdbcAvroSchema.getDatabaseTableName(meta)); } @Test @@ -80,8 +80,8 @@ public void shouldConvertDateSqlTypeWithAvroLogicalType() throws SQLException { final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, true); - Assert.assertEquals(Schema.Type.LONG, fieldSchema.getType()); - Assert.assertEquals("timestamp-millis", fieldSchema.getProp("logicalType")); + Assertions.assertEquals(Schema.Type.LONG, fieldSchema.getType()); + Assertions.assertEquals("timestamp-millis", fieldSchema.getProp("logicalType")); } @Test @@ -90,8 +90,8 @@ public void shouldConvertDateSqlTypeWithoutAvroLogicalType() throws SQLException final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, false); - Assert.assertEquals(Schema.Type.LONG, fieldSchema.getType()); - Assert.assertNull(fieldSchema.getProp("logicalType")); + Assertions.assertEquals(Schema.Type.LONG, fieldSchema.getType()); + Assertions.assertNull(fieldSchema.getProp("logicalType")); } @Test @@ -100,7 +100,7 @@ public void shouldConvertBigIntSqlTypeToLong() throws SQLException { final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, false); - Assert.assertEquals(Schema.Type.LONG, fieldSchema.getType()); + Assertions.assertEquals(Schema.Type.LONG, fieldSchema.getType()); } @Test @@ -109,7 +109,7 @@ public void shouldConvertBitSqlTypeWithNoPrecisionToBoolean() throws SQLExceptio final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, false); - Assert.assertEquals(Schema.Type.BOOLEAN, fieldSchema.getType()); + Assertions.assertEquals(Schema.Type.BOOLEAN, fieldSchema.getType()); } @Test @@ -119,30 +119,34 @@ public void shouldConvertBitSqlTypeWithPrecision2ToBytes() throws SQLException { final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, false); - Assert.assertEquals(Schema.Type.BYTES, fieldSchema.getType()); + Assertions.assertEquals(Schema.Type.BYTES, fieldSchema.getType()); } @Test public void shouldThrowOnNonSupportedTypes() throws SQLException { final ResultSet resultSet = buildMockResultSet(Types.STRUCT); - RuntimeException thrown = Assert.assertThrows(RuntimeException.class, - () -> createAvroSchemaForSingleField(resultSet, false)); - Assert.assertEquals("STRUCT type is not supported", thrown.getMessage()); + RuntimeException thrown = + Assertions.assertThrows( + RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet, false)); + Assertions.assertEquals("STRUCT type is not supported", thrown.getMessage()); final ResultSet resultSet2 = buildMockResultSet(Types.REF); - RuntimeException thrown2 = Assert.assertThrows(RuntimeException.class, - () -> createAvroSchemaForSingleField(resultSet2, false)); - Assert.assertEquals("REF and REF_CURSOR type are not supported", thrown2.getMessage()); + RuntimeException thrown2 = + Assertions.assertThrows( + RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet2, false)); + Assertions.assertEquals("REF and REF_CURSOR type are not supported", thrown2.getMessage()); final ResultSet resultSet3 = buildMockResultSet(Types.REF_CURSOR); - RuntimeException thrown3 = Assert.assertThrows(RuntimeException.class, - () -> createAvroSchemaForSingleField(resultSet3, false)); - Assert.assertEquals("REF and REF_CURSOR type are not supported", thrown3.getMessage()); + RuntimeException thrown3 = + Assertions.assertThrows( + RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet3, false)); + Assertions.assertEquals("REF and REF_CURSOR type are not supported", thrown3.getMessage()); final ResultSet resultSet4 = buildMockResultSet(Types.DATALINK); - RuntimeException thrown4 = Assert.assertThrows(RuntimeException.class, - () -> createAvroSchemaForSingleField(resultSet4, false)); - Assert.assertEquals("DATALINK type is not supported", thrown4.getMessage()); + RuntimeException thrown4 = + Assertions.assertThrows( + RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet4, false)); + Assertions.assertEquals("DATALINK type is not supported", thrown4.getMessage()); } @Test @@ -152,7 +156,7 @@ public void shouldConvertIntegerWithLongColumnClassNameToLong() throws SQLExcept final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, false); - Assert.assertEquals(Schema.Type.LONG, fieldSchema.getType()); + Assertions.assertEquals(Schema.Type.LONG, fieldSchema.getType()); } @Test @@ -161,7 +165,7 @@ public void shouldConvertIntegerSqlTypeToInteger() throws SQLException { final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, false); - Assert.assertEquals(Schema.Type.INT, fieldSchema.getType()); + Assertions.assertEquals(Schema.Type.INT, fieldSchema.getType()); } @Test @@ -170,7 +174,7 @@ public void shouldDefaultConversionToStringType() throws SQLException { final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, false); - Assert.assertEquals(Schema.Type.STRING, fieldSchema.getType()); + Assertions.assertEquals(Schema.Type.STRING, fieldSchema.getType()); } @Test @@ -178,8 +182,8 @@ public void shouldConvertUuidSqlTypeWithAvroLogicalType() throws SQLException { final ResultSet resultSet = buildMockResultSet(Types.OTHER, "uuid"); final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, true); - Assert.assertEquals(Schema.Type.STRING, fieldSchema.getType()); - Assert.assertEquals("uuid", fieldSchema.getProp("logicalType")); + Assertions.assertEquals(Schema.Type.STRING, fieldSchema.getType()); + Assertions.assertEquals("uuid", fieldSchema.getProp("logicalType")); } @Test @@ -187,16 +191,22 @@ public void shouldConvertUuidSqlTypeWithoutAvroLogicalType() throws SQLException final ResultSet resultSet = buildMockResultSet(Types.OTHER, "uuid"); final Schema fieldSchema = createAvroSchemaForSingleField(resultSet, false); - Assert.assertEquals(Schema.Type.STRING, fieldSchema.getType()); - Assert.assertNull(fieldSchema.getProp("logicalType")); + Assertions.assertEquals(Schema.Type.STRING, fieldSchema.getType()); + Assertions.assertNull(fieldSchema.getProp("logicalType")); } private Schema createAvroSchemaForSingleField( final ResultSet resultSet, final boolean useLogicalTypes) throws SQLException { Schema avroSchema = JdbcAvroSchema.createAvroSchema( - resultSet, "namespace1", "url1", Optional.empty(), "doc1", useLogicalTypes, - ArrayHandlingMode.TypedMetaFromFirstRow, false); + resultSet, + "namespace1", + "url1", + Optional.empty(), + "doc1", + useLogicalTypes, + ArrayHandlingMode.TypedMetaFromFirstRow, + false); return avroSchema.getField("column1").schema().getTypes().get(COLUMN_NUM); } @@ -205,8 +215,8 @@ private ResultSet buildMockResultSet(final int inputColumnType) throws SQLExcept return buildMockResultSet(inputColumnType, null); } - private ResultSet buildMockResultSet(final int inputColumnType, - final String columnTypeName) throws SQLException { + private ResultSet buildMockResultSet(final int inputColumnType, final String columnTypeName) + throws SQLException { final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); when(meta.getColumnCount()).thenReturn(COLUMN_NUM); when(meta.getTableName(COLUMN_NUM)).thenReturn("test_table"); diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/avro/PostgresJdbcAvroTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/avro/PostgresJdbcAvroTest.java index 5f1d8239..35502fb5 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/avro/PostgresJdbcAvroTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/avro/PostgresJdbcAvroTest.java @@ -45,14 +45,13 @@ import org.apache.avro.generic.GenericDatumWriter; import org.apache.avro.generic.GenericRecord; import org.apache.avro.util.Utf8; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; import org.mockito.Mockito; public class PostgresJdbcAvroTest { - public static GenericRecord[] bytesToGenericRecords(Schema schema, - ByteBuffer... avroRecordBytes) + public static GenericRecord[] bytesToGenericRecords(Schema schema, ByteBuffer... avroRecordBytes) throws IOException { DataFileWriter dataFileWriter = new DataFileWriter<>(new GenericDatumWriter<>(schema)); @@ -66,16 +65,16 @@ public static GenericRecord[] bytesToGenericRecords(Schema schema, SeekableByteArrayInput inputStream = new SeekableByteArrayInput(avroOutputStream.toByteArray()); List genericRecords = new ArrayList<>(); - DataFileReader dataReader = new DataFileReader<>(inputStream, - new GenericDatumReader<>(schema)); + DataFileReader dataReader = + new DataFileReader<>(inputStream, new GenericDatumReader<>(schema)); while (dataReader.hasNext()) { genericRecords.add(dataReader.next()); } return genericRecords.toArray(new GenericRecord[0]); } - public void assertGenericRecordArrayField(GenericRecord record, String fieldName, - String... expectedItems) { + public void assertGenericRecordArrayField( + GenericRecord record, String fieldName, String... expectedItems) { Utf8[] expectedItemsConverted = new Utf8[expectedItems.length]; for (int i = 0; i < expectedItems.length; i++) { expectedItemsConverted[i] = expectedItems[i] != null ? new Utf8(expectedItems[i]) : null; @@ -83,15 +82,15 @@ public void assertGenericRecordArrayField(GenericRecord record, String fieldName assertGenericRecordArrayField(record, fieldName, (Object[]) expectedItemsConverted); } - public void assertGenericRecordArrayField(GenericRecord record, String fieldName, - Object... expectedItems) { + public void assertGenericRecordArrayField( + GenericRecord record, String fieldName, Object... expectedItems) { final GenericData.Array arrayValue = (GenericData.Array) record.get(fieldName); - Assert.assertEquals(expectedItems.length, arrayValue.size()); + Assertions.assertEquals(expectedItems.length, arrayValue.size()); for (int i = 0; i < expectedItems.length; i++) { - Assert.assertEquals(expectedItems[i], arrayValue.get(i)); + Assertions.assertEquals(expectedItems[i], arrayValue.get(i)); } } @@ -105,20 +104,21 @@ public void shouldEncodeUUIDValues() throws SQLException, IOException { when(resultSet.getMetaData()).thenReturn(meta); final UUID uuidExpected = UUID.randomUUID(); when(resultSet.getObject(1)).thenReturn(uuidExpected); - TestHelper.mockArrayColumn(meta, resultSet, 2, "array_field", "_uuid", Types.OTHER, - "uuid", new UUID[] {uuidExpected}); + TestHelper.mockArrayColumn( + meta, resultSet, 2, "array_field", "_uuid", Types.OTHER, "uuid", new UUID[] {uuidExpected}); when(resultSet.isFirst()).thenReturn(true); String arrayMode = ArrayHandlingMode.TypedMetaFromFirstRow; - final Schema schema = JdbcAvroSchema.createAvroSchema(resultSet, "ns", "conn_url", - Optional.empty(), "doc", true, arrayMode, false); - final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create( - resultSet, arrayMode, false); + final Schema schema = + JdbcAvroSchema.createAvroSchema( + resultSet, "ns", "conn_url", Optional.empty(), "doc", true, arrayMode, false); + final JdbcAvroRecordConverter converter = + JdbcAvroRecordConverter.create(resultSet, arrayMode, false); - GenericRecord actualRecord = bytesToGenericRecords(schema, - converter.convertResultSetIntoAvroBytes())[0]; + GenericRecord actualRecord = + bytesToGenericRecords(schema, converter.convertResultSetIntoAvroBytes())[0]; - Assert.assertEquals(actualRecord.get("uuid_field"), new Utf8(uuidExpected.toString())); + Assertions.assertEquals(actualRecord.get("uuid_field"), new Utf8(uuidExpected.toString())); assertGenericRecordArrayField(actualRecord, "array_field", new Utf8(uuidExpected.toString())); } @@ -127,32 +127,54 @@ public void shouldEncodeStringAndOtherValues() throws SQLException, IOException final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); when(meta.getColumnCount()).thenReturn(5); TestHelper.mockResultSetMeta(meta, 1, Types.VARCHAR, "text_field", "java.lang.String", "text"); - TestHelper.mockResultSetMeta(meta, 2, Types.OTHER, "other_field", "java.util.UUID", - "something_else"); + TestHelper.mockResultSetMeta( + meta, 2, Types.OTHER, "other_field", "java.util.UUID", "something_else"); final ResultSet resultSet = Mockito.mock(ResultSet.class); when(resultSet.getMetaData()).thenReturn(meta); when(resultSet.getString(1)).thenReturn("some_text_42"); when(resultSet.getString(2)).thenReturn("some_other_42"); - TestHelper.mockArrayColumn(meta, resultSet, 3, "array_field1", "_text", Types.VARCHAR, - "text", new String[] {"some_text_42"}); - TestHelper.mockArrayColumn(meta, resultSet, 4, "array_field2", "_varchar", - Types.VARCHAR, "varchar", (Object) new String[] {"some_varchar_42"}); - TestHelper.mockArrayColumn(meta, resultSet, 5, "array_other", "_other", - Types.OTHER, "other", (Object) new String[] {"some_other_42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 3, + "array_field1", + "_text", + Types.VARCHAR, + "text", + new String[] {"some_text_42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 4, + "array_field2", + "_varchar", + Types.VARCHAR, + "varchar", + (Object) new String[] {"some_varchar_42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 5, + "array_other", + "_other", + Types.OTHER, + "other", + (Object) new String[] {"some_other_42"}); when(resultSet.isFirst()).thenReturn(true); String arrayMode = ArrayHandlingMode.TypedMetaFromFirstRow; - final Schema schema = JdbcAvroSchema.createAvroSchema(resultSet, "ns", "conn_url", - Optional.empty(), "doc", true, arrayMode, false); - final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(resultSet, arrayMode, - false); - - GenericRecord actualRecord = bytesToGenericRecords(schema, - converter.convertResultSetIntoAvroBytes())[0]; - Assert.assertEquals(actualRecord.get("text_field"), new Utf8("some_text_42")); - Assert.assertEquals(actualRecord.get("other_field"), new Utf8("some_other_42")); + final Schema schema = + JdbcAvroSchema.createAvroSchema( + resultSet, "ns", "conn_url", Optional.empty(), "doc", true, arrayMode, false); + final JdbcAvroRecordConverter converter = + JdbcAvroRecordConverter.create(resultSet, arrayMode, false); + + GenericRecord actualRecord = + bytesToGenericRecords(schema, converter.convertResultSetIntoAvroBytes())[0]; + Assertions.assertEquals(actualRecord.get("text_field"), new Utf8("some_text_42")); + Assertions.assertEquals(actualRecord.get("other_field"), new Utf8("some_other_42")); assertGenericRecordArrayField(actualRecord, "array_field1", "some_text_42"); assertGenericRecordArrayField(actualRecord, "array_field2", "some_varchar_42"); assertGenericRecordArrayField(actualRecord, "array_other", "some_other_42"); @@ -168,9 +190,18 @@ public void shouldThrowOnArrayWithNulls() throws SQLException, IOException { when(resultSet.getArray(1)).thenReturn(null); when(resultSet.isFirst()).thenReturn(true); - Assert.assertThrows(RuntimeException.class, () -> JdbcAvroSchema.createAvroSchema(resultSet, - "ns", "conn_url", - Optional.empty(), "doc", true, ArrayHandlingMode.TypedMetaFromFirstRow, false)); + Assertions.assertThrows( + RuntimeException.class, + () -> + JdbcAvroSchema.createAvroSchema( + resultSet, + "ns", + "conn_url", + Optional.empty(), + "doc", + true, + ArrayHandlingMode.TypedMetaFromFirstRow, + false)); } @Test @@ -186,15 +217,16 @@ public void shouldHandleArrayWithNullsUsingArrayAsBytes() throws SQLException, I when(resultSet.isFirst()).thenReturn(true); String arrayMode = ArrayHandlingMode.Bytes; - final Schema schema = JdbcAvroSchema.createAvroSchema(resultSet, "ns", "conn_url", - Optional.empty(), "doc", true, arrayMode, false); - final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(resultSet, arrayMode, - false); + final Schema schema = + JdbcAvroSchema.createAvroSchema( + resultSet, "ns", "conn_url", Optional.empty(), "doc", true, arrayMode, false); + final JdbcAvroRecordConverter converter = + JdbcAvroRecordConverter.create(resultSet, arrayMode, false); - GenericRecord actualRecord = bytesToGenericRecords(schema, - converter.convertResultSetIntoAvroBytes())[0]; - Assert.assertArrayEquals(expectedValue, - ((java.nio.ByteBuffer) actualRecord.get("array_field")).array()); + GenericRecord actualRecord = + bytesToGenericRecords(schema, converter.convertResultSetIntoAvroBytes())[0]; + Assertions.assertArrayEquals( + expectedValue, ((java.nio.ByteBuffer) actualRecord.get("array_field")).array()); } @Test @@ -204,41 +236,92 @@ public void shouldHandleArrayWithNullsWithoutReadingFirstRow() throws SQLExcepti final ResultSet resultSet = Mockito.mock(ResultSet.class); when(resultSet.getMetaData()).thenReturn(meta); - TestHelper.mockArrayColumn(meta, resultSet, 1, "array_field_varchar", "_varchar", - Types.VARCHAR, "varchar", null, (Object) new String[] {"some_varchar_42", "42"}); - TestHelper.mockArrayColumn(meta, resultSet, 2, "array_field_text", "_text", - Types.VARCHAR, "text", null, (Object) new String[] {"some_text_42", "42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 1, + "array_field_varchar", + "_varchar", + Types.VARCHAR, + "varchar", + null, + (Object) new String[] {"some_varchar_42", "42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 2, + "array_field_text", + "_text", + Types.VARCHAR, + "text", + null, + (Object) new String[] {"some_text_42", "42"}); final UUID uuidExpected = UUID.randomUUID(); - TestHelper.mockArrayColumn(meta, resultSet, 3, "array_field_uuid", "_uuid", - Types.VARCHAR, "uuid", null, (Object) new UUID[] {uuidExpected}); - TestHelper.mockArrayColumn(meta, resultSet, 4, "array_field_int", "_int", - Types.VARCHAR, "int", null, (Object) new Integer[] {42}); - TestHelper.mockArrayColumn(meta, resultSet, 5, "array_field_int4", "_int4", - Types.VARCHAR, "int4", null, (Object) new Integer[] {42}); - TestHelper.mockArrayColumn(meta, resultSet, 6, "array_field_int8", "_int8", - Types.VARCHAR, "int8", null, (Object) new Long[] {42L}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 3, + "array_field_uuid", + "_uuid", + Types.VARCHAR, + "uuid", + null, + (Object) new UUID[] {uuidExpected}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 4, + "array_field_int", + "_int", + Types.VARCHAR, + "int", + null, + (Object) new Integer[] {42}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 5, + "array_field_int4", + "_int4", + Types.VARCHAR, + "int4", + null, + (Object) new Integer[] {42}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 6, + "array_field_int8", + "_int8", + Types.VARCHAR, + "int8", + null, + (Object) new Long[] {42L}); when(resultSet.next()).thenReturn(true, true, false); when(resultSet.isFirst()).thenReturn(true); String arrayMode = ArrayHandlingMode.TypedMetaPostgres; - final Schema schema = JdbcAvroSchema.createAvroSchema(resultSet, "ns", "conn_url", - Optional.empty(), "doc", true, arrayMode, false); - final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(resultSet, arrayMode, - false); - GenericRecord[] actualRecords = bytesToGenericRecords(schema, - converter.convertResultSetIntoAvroBytes(), converter.convertResultSetIntoAvroBytes()); - - Assert.assertNull(actualRecords[0].get("array_field_varchar")); - Assert.assertNull(actualRecords[0].get("array_field_text")); - Assert.assertNull(actualRecords[0].get("array_field_uuid")); - Assert.assertNull(actualRecords[0].get("array_field_int")); - Assert.assertNull(actualRecords[0].get("array_field_int4")); - Assert.assertNull(actualRecords[0].get("array_field_int8")); + final Schema schema = + JdbcAvroSchema.createAvroSchema( + resultSet, "ns", "conn_url", Optional.empty(), "doc", true, arrayMode, false); + final JdbcAvroRecordConverter converter = + JdbcAvroRecordConverter.create(resultSet, arrayMode, false); + GenericRecord[] actualRecords = + bytesToGenericRecords( + schema, + converter.convertResultSetIntoAvroBytes(), + converter.convertResultSetIntoAvroBytes()); + + Assertions.assertNull(actualRecords[0].get("array_field_varchar")); + Assertions.assertNull(actualRecords[0].get("array_field_text")); + Assertions.assertNull(actualRecords[0].get("array_field_uuid")); + Assertions.assertNull(actualRecords[0].get("array_field_int")); + Assertions.assertNull(actualRecords[0].get("array_field_int4")); + Assertions.assertNull(actualRecords[0].get("array_field_int8")); assertGenericRecordArrayField(actualRecords[1], "array_field_varchar", "some_varchar_42", "42"); assertGenericRecordArrayField(actualRecords[1], "array_field_text", "some_text_42", "42"); - assertGenericRecordArrayField(actualRecords[1], "array_field_uuid", - uuidExpected.toString()); + assertGenericRecordArrayField(actualRecords[1], "array_field_uuid", uuidExpected.toString()); assertGenericRecordArrayField(actualRecords[1], "array_field_int", 42); assertGenericRecordArrayField(actualRecords[1], "array_field_int4", 42); assertGenericRecordArrayField(actualRecords[1], "array_field_int8", 42L); @@ -251,30 +334,61 @@ public void shouldHandleArrayWithNullsIfEnabled() throws SQLException, IOExcepti final ResultSet resultSet = Mockito.mock(ResultSet.class); when(resultSet.getMetaData()).thenReturn(meta); - TestHelper.mockArrayColumn(meta, resultSet, 1, "array_field_varchar", "_varchar", - Types.VARCHAR, "varchar", new String[] { null, "some_varchar_42", "42"}); - TestHelper.mockArrayColumn(meta, resultSet, 2, "array_field_text", "_text", - Types.VARCHAR, "text", new String[] { "some_text_42", null, "42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 1, + "array_field_varchar", + "_varchar", + Types.VARCHAR, + "varchar", + new String[] {null, "some_varchar_42", "42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 2, + "array_field_text", + "_text", + Types.VARCHAR, + "text", + new String[] {"some_text_42", null, "42"}); final UUID uuidExpected = UUID.randomUUID(); - TestHelper.mockArrayColumn(meta, resultSet, 3, "array_field_uuid", "_uuid", - Types.VARCHAR, "uuid", new UUID[] { uuidExpected, null}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 3, + "array_field_uuid", + "_uuid", + Types.VARCHAR, + "uuid", + new UUID[] {uuidExpected, null}); when(resultSet.next()).thenReturn(true, false); when(resultSet.isFirst()).thenReturn(true, false); String arrayMode = ArrayHandlingMode.TypedMetaPostgres; boolean nullableArrayItems = true; - final Schema schema = JdbcAvroSchema.createAvroSchema(resultSet, "ns", "conn_url", - Optional.empty(), "doc", true, arrayMode, nullableArrayItems); - final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(resultSet, arrayMode, - nullableArrayItems); - GenericRecord actualRecord = bytesToGenericRecords(schema, - converter.convertResultSetIntoAvroBytes(), converter.convertResultSetIntoAvroBytes())[0]; - - assertGenericRecordArrayField(actualRecord, "array_field_varchar", null, "some_varchar_42", - "42"); + final Schema schema = + JdbcAvroSchema.createAvroSchema( + resultSet, + "ns", + "conn_url", + Optional.empty(), + "doc", + true, + arrayMode, + nullableArrayItems); + final JdbcAvroRecordConverter converter = + JdbcAvroRecordConverter.create(resultSet, arrayMode, nullableArrayItems); + GenericRecord actualRecord = + bytesToGenericRecords( + schema, + converter.convertResultSetIntoAvroBytes(), + converter.convertResultSetIntoAvroBytes())[0]; + + assertGenericRecordArrayField( + actualRecord, "array_field_varchar", null, "some_varchar_42", "42"); assertGenericRecordArrayField(actualRecord, "array_field_text", "some_text_42", null, "42"); - assertGenericRecordArrayField(actualRecord, "array_field_uuid", - uuidExpected.toString(), null); + assertGenericRecordArrayField(actualRecord, "array_field_uuid", uuidExpected.toString(), null); } @Test @@ -284,24 +398,47 @@ public void shouldThrowOnArrayWithNullsIfDisabled() throws SQLException { final ResultSet resultSet = Mockito.mock(ResultSet.class); when(resultSet.getMetaData()).thenReturn(meta); - TestHelper.mockArrayColumn(meta, resultSet, 1, "array_field_varchar", "_varchar", - Types.VARCHAR, "varchar", new String[] { null, "some_varchar_42", "42"}); - TestHelper.mockArrayColumn(meta, resultSet, 2, "array_field_text", "_text", - Types.VARCHAR, "text", new String[] { "some_text_42", null, "42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 1, + "array_field_varchar", + "_varchar", + Types.VARCHAR, + "varchar", + new String[] {null, "some_varchar_42", "42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 2, + "array_field_text", + "_text", + Types.VARCHAR, + "text", + new String[] {"some_text_42", null, "42"}); final UUID uuidExpected = UUID.randomUUID(); - TestHelper.mockArrayColumn(meta, resultSet, 3, "array_field_uuid", "_uuid", - Types.VARCHAR, "uuid", new UUID[] { uuidExpected, null}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 3, + "array_field_uuid", + "_uuid", + Types.VARCHAR, + "uuid", + new UUID[] {uuidExpected, null}); when(resultSet.next()).thenReturn(true, false); when(resultSet.isFirst()).thenReturn(true, false); String arrayMode = ArrayHandlingMode.TypedMetaPostgres; boolean nullableArrayItems = false; - final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(resultSet, arrayMode, - nullableArrayItems); - RuntimeException thrown = Assert.assertThrows(RuntimeException.class, - () -> converter.convertResultSetIntoAvroBytes()); - Assert.assertEquals("Array item is null in column 'array_field_varchar', use " - + "--nullableArrayItems", thrown.getMessage()); + final JdbcAvroRecordConverter converter = + JdbcAvroRecordConverter.create(resultSet, arrayMode, nullableArrayItems); + RuntimeException thrown = + Assertions.assertThrows( + RuntimeException.class, () -> converter.convertResultSetIntoAvroBytes()); + Assertions.assertEquals( + "Array item is null in column 'array_field_varchar', use " + "--nullableArrayItems", + thrown.getMessage()); } @Test @@ -311,19 +448,28 @@ public void shouldThrowOnUnknownType() throws SQLException { final ResultSet resultSet = Mockito.mock(ResultSet.class); when(resultSet.getMetaData()).thenReturn(meta); - TestHelper.mockArrayColumn(meta, resultSet, 1, "invalid_array", "_uuid", - Types.VARCHAR, "uuid", new File[] { new File("/some/file")}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 1, + "invalid_array", + "_uuid", + Types.VARCHAR, + "uuid", + new File[] {new File("/some/file")}); when(resultSet.next()).thenReturn(true, false); when(resultSet.isFirst()).thenReturn(true, false); String arrayMode = ArrayHandlingMode.TypedMetaPostgres; boolean nullableArrayItems = false; - final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(resultSet, arrayMode, - nullableArrayItems); - RuntimeException thrown = Assert.assertThrows(RuntimeException.class, - () -> converter.convertResultSetIntoAvroBytes()); - Assert.assertEquals("Value of type class java.io.File in column 'invalid_array' is not " - + "supported", thrown.getMessage()); + final JdbcAvroRecordConverter converter = + JdbcAvroRecordConverter.create(resultSet, arrayMode, nullableArrayItems); + RuntimeException thrown = + Assertions.assertThrows( + RuntimeException.class, () -> converter.convertResultSetIntoAvroBytes()); + Assertions.assertEquals( + "Value of type class java.io.File in column 'invalid_array' is not " + "supported", + thrown.getMessage()); } @Test @@ -333,17 +479,35 @@ public void shouldThrowOnInvalidArrayColumnTypeName() throws SQLException { final ResultSet resultSet = Mockito.mock(ResultSet.class); when(resultSet.getMetaData()).thenReturn(meta); - TestHelper.mockArrayColumn(meta, resultSet, 1, "array_field_text", "text", - Types.VARCHAR, "text", new String[] { "some_text_42", null, "42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 1, + "array_field_text", + "text", + Types.VARCHAR, + "text", + new String[] {"some_text_42", null, "42"}); when(resultSet.next()).thenReturn(true, false); when(resultSet.isFirst()).thenReturn(true, false); String arrayMode = ArrayHandlingMode.TypedMetaPostgres; boolean nullableArrayItems = true; - RuntimeException thrown = Assert.assertThrows(RuntimeException.class, - () -> JdbcAvroSchema.createAvroSchema(resultSet, "ns", "conn_url", - Optional.empty(), "doc", true, arrayMode, nullableArrayItems)); - Assert.assertEquals("columnName=array_field_text columnTypeName=text should start with '_'", + RuntimeException thrown = + Assertions.assertThrows( + RuntimeException.class, + () -> + JdbcAvroSchema.createAvroSchema( + resultSet, + "ns", + "conn_url", + Optional.empty(), + "doc", + true, + arrayMode, + nullableArrayItems)); + Assertions.assertEquals( + "columnName=array_field_text columnTypeName=text should start with '_'", thrown.getMessage()); } @@ -354,17 +518,34 @@ public void shouldThrowOnNotSupportedArrayColumnTypeName() throws SQLException { final ResultSet resultSet = Mockito.mock(ResultSet.class); when(resultSet.getMetaData()).thenReturn(meta); - TestHelper.mockArrayColumn(meta, resultSet, 1, "array_field_text", "_not_supported", - Types.VARCHAR, "not_supported", new String[] { "some_text_42", null, "42"}); + TestHelper.mockArrayColumn( + meta, + resultSet, + 1, + "array_field_text", + "_not_supported", + Types.VARCHAR, + "not_supported", + new String[] {"some_text_42", null, "42"}); when(resultSet.next()).thenReturn(true, false); when(resultSet.isFirst()).thenReturn(true, false); String arrayMode = ArrayHandlingMode.TypedMetaPostgres; boolean nullableArrayItems = true; - RuntimeException thrown = Assert.assertThrows(RuntimeException.class, - () -> JdbcAvroSchema.createAvroSchema(resultSet, "ns", "conn_url", - Optional.empty(), "doc", true, arrayMode, nullableArrayItems)); - Assert.assertEquals( + RuntimeException thrown = + Assertions.assertThrows( + RuntimeException.class, + () -> + JdbcAvroSchema.createAvroSchema( + resultSet, + "ns", + "conn_url", + Optional.empty(), + "doc", + true, + arrayMode, + nullableArrayItems)); + Assertions.assertEquals( "columnName=array_field_text Postgres type 'not_supported' is not supported", thrown.getMessage()); } diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/beam/BeamHelperTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/beam/BeamHelperTest.java index 1a348f00..892ebada 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/beam/BeamHelperTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/beam/BeamHelperTest.java @@ -20,13 +20,13 @@ package com.spotify.dbeam.beam; -import static org.junit.Assert.assertTrue; -import static org.junit.Assert.fail; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; import java.io.IOException; import org.apache.beam.sdk.io.FileSystems; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.junit.Test; +import org.junit.jupiter.api.Test; public class BeamHelperTest { @@ -40,9 +40,10 @@ public void shouldWorkWithGcs() throws IOException { } catch (IOException e) { // Beam changed behavior - now throws IOException instead of FileNotFoundException // Verify it's the expected GCS error message - assertTrue("Exception should indicate GCS file matching error", - e.getMessage().contains("Error matching file spec") - && e.getMessage().contains("gs://does-not-exist-1")); + assertTrue( + e.getMessage().contains("Error matching file spec") + && e.getMessage().contains("gs://does-not-exist-1"), + "Exception should indicate GCS file matching error"); } } } diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/BeamHelperTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/BeamHelperTest.java index 9c20ffd0..0b5cb392 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/BeamHelperTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/BeamHelperTest.java @@ -26,8 +26,8 @@ import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.MetricResults; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; public class BeamHelperTest { @@ -62,9 +62,9 @@ public MetricResults metrics() { }; try { BeamHelper.waitUntilDone(mockResult, Duration.ofMinutes(1)); - Assert.fail("A PipelineExecutionException should be thrown"); + Assertions.fail("A PipelineExecutionException should be thrown"); } catch (Pipeline.PipelineExecutionException exception) { - Assert.assertEquals( + Assertions.assertEquals( "java.lang.Exception: Job finished with terminalState FAILED", exception.getMessage()); } } @@ -100,9 +100,9 @@ public MetricResults metrics() { }; try { BeamHelper.waitUntilDone(mockResult, Duration.ofMinutes(1)); - Assert.fail("A PipelineExecutionException should be thrown"); + Assertions.fail("A PipelineExecutionException should be thrown"); } catch (Pipeline.PipelineExecutionException exception) { - Assert.assertEquals( + Assertions.assertEquals( "java.lang.Exception: Job cancelled after exceeding timeout PT1M", exception.getMessage()); } @@ -139,9 +139,9 @@ public MetricResults metrics() { }; try { BeamHelper.waitUntilDone(mockResult, Duration.ofMinutes(1)); - Assert.fail("A PipelineExecutionException should be thrown"); + Assertions.fail("A PipelineExecutionException should be thrown"); } catch (Pipeline.PipelineExecutionException exception) { - Assert.assertEquals( + Assertions.assertEquals( "java.lang.Exception: Job cancelled after exceeding timeout PT1M", exception.getMessage()); } @@ -178,9 +178,9 @@ public MetricResults metrics() { }; try { BeamHelper.waitUntilDone(mockResult, Duration.ofMinutes(1)); - Assert.fail("A PipelineExecutionException should be thrown"); + Assertions.fail("A PipelineExecutionException should be thrown"); } catch (Pipeline.PipelineExecutionException exception) { - Assert.assertEquals( + Assertions.assertEquals( "java.lang.Exception: Job exceeded timeout of PT1M, " + "but was not possible to cancel, finished with terminalState RUNNING", exception.getMessage()); diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/BenchJdbcAvroJobTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/BenchJdbcAvroJobTest.java index 5ff27485..6354b17d 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/BenchJdbcAvroJobTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/BenchJdbcAvroJobTest.java @@ -28,8 +28,8 @@ import java.io.IOException; import java.nio.file.Path; import java.sql.SQLException; -import org.junit.BeforeClass; -import org.junit.Test; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; public class BenchJdbcAvroJobTest { @@ -37,7 +37,7 @@ public class BenchJdbcAvroJobTest { "jdbc:h2:mem:test3;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; private static Path testDir; - @BeforeClass + @BeforeAll public static void beforeAll() throws SQLException, ClassNotFoundException, IOException { testDir = TestHelper.createTmpDirPath("jdbc-export-args-test"); DbTestHelper.createFixtures(CONNECTION_URL); diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/JdbcAvroJobTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/JdbcAvroJobTest.java index 0a8174d8..3dcd5b39 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/JdbcAvroJobTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/JdbcAvroJobTest.java @@ -27,7 +27,7 @@ import com.spotify.dbeam.DbTestHelper; import com.spotify.dbeam.TestHelper; -import com.spotify.dbeam.avro.JdbcAvroMetering; +import com.spotify.dbeam.avro.JdbcMetering; import com.spotify.dbeam.options.DBeamPipelineOptions; import com.spotify.dbeam.options.OutputOptions; import java.io.File; @@ -48,9 +48,9 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.hamcrest.CoreMatchers; -import org.junit.Assert; -import org.junit.BeforeClass; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; public class JdbcAvroJobTest { @@ -69,7 +69,7 @@ private List readAvroRecords(File avroFile, Schema schema) throws return records; } - @BeforeClass + @BeforeAll public static void beforeAll() throws SQLException, ClassNotFoundException, IOException { testDir = TestHelper.createTmpDirPath("jdbc-avro-test-"); passwordPath = testDir.resolve(".password"); @@ -219,23 +219,26 @@ public void shouldRunAvroJobPreCommands() throws SQLException, ClassNotFoundExce assertThat(schemas, CoreMatchers.hasItems(expectedSchemas)); } - @Test(expected = FailedValidationException.class) - public void shouldFailWithNotEnoughRows() throws Exception { + @Test + public void shouldFailWithNotEnoughRows() { final Path outputPath = testDir.resolve("shouldRunJdbcAvroJob"); - JdbcAvroJob.create( - new String[] { - "--targetParallelism=1", // no need for more threads when testing - "--partition=2025-02-28", - "--skipPartitionCheck", - "--connectionUrl=" + CONNECTION_URL, - "--username=", - "--passwordFile=" + passwordPath.toString(), - "--table=COFFEES", - "--output=" + outputPath, - "--minRows=1000" - }) - .runExport(); + Assertions.assertThrows( + FailedValidationException.class, + () -> + JdbcAvroJob.create( + new String[] { + "--targetParallelism=1", // no need for more threads when testing + "--partition=2025-02-28", + "--skipPartitionCheck", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--minRows=1000" + }) + .runExport()); } @Test @@ -255,38 +258,43 @@ public void shouldConfigureDBeamVersionPipelineOptions() throws Exception { }); jdbcAvroJob.prepareExport(); - Assert.assertEquals( + Assertions.assertEquals( this.getClass().getPackage().getImplementationVersion(), jdbcAvroJob.getPipelineOptions().as(DBeamPipelineOptions.class).getDBeamVersion()); } @Test public void shouldHaveDefaultExitCode() { - Assert.assertEquals( + Assertions.assertEquals( Integer.valueOf(49), ExceptionHandling.exitCode(new IllegalStateException())); } @Test public void shouldExit50OnFailedValidationException() { - Assert.assertEquals( + Assertions.assertEquals( Integer.valueOf(50), ExceptionHandling.exitCode(new FailedValidationException(""))); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnMissingInput() throws IOException, ClassNotFoundException { - JdbcAvroJob.create(PipelineOptionsFactory.create()); + @Test + public void shouldFailOnMissingInput() { + Assertions.assertThrows( + IllegalArgumentException.class, () -> JdbcAvroJob.create(PipelineOptionsFactory.create())); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnEmptyInput() throws IOException, ClassNotFoundException { - final PipelineOptions pipelineOptions = PipelineOptionsFactory.create(); - pipelineOptions.as(OutputOptions.class).setOutput(""); - JdbcAvroJob.create(PipelineOptionsFactory.create()); + @Test + public void shouldFailOnEmptyInput() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> { + final PipelineOptions pipelineOptions = PipelineOptionsFactory.create(); + pipelineOptions.as(OutputOptions.class).setOutput(""); + JdbcAvroJob.create(PipelineOptionsFactory.create()); + }); } @Test public void shouldIncrementCounterMetrics() { - final JdbcAvroMetering metering = new JdbcAvroMetering(1, 1); + final JdbcMetering metering = new JdbcMetering(1, 1, "jdbcavroio"); metering.startWriteMeter(); metering.exposeWriteElapsed(); metering.incrementRecordCount(); diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/PsqlAvroJobTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/PsqlAvroJobTest.java index 676c58eb..1ab537f4 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/PsqlAvroJobTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/PsqlAvroJobTest.java @@ -21,7 +21,7 @@ package com.spotify.dbeam.jobs; import java.io.IOException; -import org.junit.Test; +import org.junit.jupiter.api.Test; public class PsqlAvroJobTest { diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/PsqlReplicationCheckTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/PsqlReplicationCheckTest.java index b5491463..8147de98 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/jobs/PsqlReplicationCheckTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/PsqlReplicationCheckTest.java @@ -28,8 +28,8 @@ import java.time.Instant; import java.time.Period; import java.util.Optional; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; public class PsqlReplicationCheckTest { private static String CONNECTION_URL = @@ -48,20 +48,22 @@ private static JdbcExportArgs createArgs(String url, QueryBuilderArgs queryBuild Optional.empty()); } - @Test(expected = IllegalArgumentException.class) + @Test public void shouldFailOnInvalidDriver() throws ClassNotFoundException { final JdbcExportArgs args = createArgs("jdbc:mysql://some_db", QueryBuilderArgs.create("some_table")); - PsqlReplicationCheck.validateOptions(args); + Assertions.assertThrows( + IllegalArgumentException.class, () -> PsqlReplicationCheck.validateOptions(args)); } - @Test(expected = IllegalArgumentException.class) + @Test public void shouldFailOnMissingPartition() throws ClassNotFoundException { final JdbcExportArgs args = createArgs("jdbc:postgresql://some_db", QueryBuilderArgs.create("some_table")); - PsqlReplicationCheck.validateOptions(args); + Assertions.assertThrows( + IllegalArgumentException.class, () -> PsqlReplicationCheck.validateOptions(args)); } @Test @@ -82,7 +84,7 @@ public void shouldBeNotReplicationDelayedWhenReplicatedUntilEndOfPartition() { final Instant lastReplication = Instant.parse("2027-08-01T00:00:00Z"); final Period partitionPeriod = Period.ofDays(1); - Assert.assertFalse( + Assertions.assertFalse( PsqlReplicationCheck.isReplicationDelayed(partition, lastReplication, partitionPeriod)); } @@ -92,7 +94,7 @@ public void shouldBeNotReplicationDelayedWhenReplicatedUntilEndTheNextDay() { final Instant lastReplication = Instant.parse("2027-08-02T00:00:00Z"); final Period partitionPeriod = Period.ofDays(1); - Assert.assertFalse( + Assertions.assertFalse( PsqlReplicationCheck.isReplicationDelayed(partition, lastReplication, partitionPeriod)); } @@ -102,7 +104,7 @@ public void shouldBeReplicationDelayedWhenReplicatedUpToPartitionStart() { final Instant lastReplication = Instant.parse("2027-07-31T00:00:00Z"); final Period partitionPeriod = Period.ofDays(1); - Assert.assertTrue( + Assertions.assertTrue( PsqlReplicationCheck.isReplicationDelayed(partition, lastReplication, partitionPeriod)); } @@ -112,7 +114,7 @@ public void shouldBeReplicationDelayedWhenReplicatedBeforePartitionStart() { final Instant lastReplication = Instant.parse("2027-07-30T22:00:00Z"); final Period partitionPeriod = Period.ofDays(1); - Assert.assertTrue( + Assertions.assertTrue( PsqlReplicationCheck.isReplicationDelayed(partition, lastReplication, partitionPeriod)); } @@ -122,7 +124,7 @@ public void shouldBeReplicationDelayedWhenReplicatedInsidePartition() { final Instant lastReplication = Instant.parse("2027-07-31T23:59:59Z"); final Period partitionPeriod = Period.ofDays(1); - Assert.assertTrue( + Assertions.assertTrue( PsqlReplicationCheck.isReplicationDelayed(partition, lastReplication, partitionPeriod)); } @@ -132,7 +134,7 @@ public void shouldWorkWithMonthlyPartitionPeriod() { final Instant lastReplication = Instant.parse("2027-07-31T23:59:59Z"); final Period partitionPeriod = Period.ofMonths(1); - Assert.assertTrue( + Assertions.assertTrue( PsqlReplicationCheck.isReplicationDelayed(partition, lastReplication, partitionPeriod)); } @@ -142,7 +144,7 @@ public void shouldBeDelayedWithHourlyPartitionPeriod() { final Instant lastReplication = Instant.parse("2027-07-31T00:59:59Z"); final Duration partitionPeriod = Duration.ofHours(1); - Assert.assertTrue( + Assertions.assertTrue( PsqlReplicationCheck.isReplicationDelayed(partition, lastReplication, partitionPeriod)); } @@ -152,11 +154,11 @@ public void shouldBeNotDelayedWithHourlyPartitionPeriod() { final Instant lastReplication = Instant.parse("2027-07-31T00:59:59Z"); final Duration partitionPeriod = Duration.ofHours(1); - Assert.assertTrue( + Assertions.assertTrue( PsqlReplicationCheck.isReplicationDelayed(partition, lastReplication, partitionPeriod)); } - @Test(expected = NotReadyException.class) + @Test public void shouldRunQueryAndReturnReplicationDelayed() throws Exception { final String query = "SELECT parsedatetime('2017-02-01 23.58.57 UTC', 'yyyy-MM-dd HH.mm.ss z', 'en', 'UTC')" @@ -175,9 +177,9 @@ public void shouldRunQueryAndReturnReplicationDelayed() throws Exception { final Instant actual = replicationCheck.queryReplication(); - Assert.assertEquals(expectedLastReplication, actual); - Assert.assertTrue(replicationCheck.isReplicationDelayed()); - replicationCheck.checkReplication(); + Assertions.assertEquals(expectedLastReplication, actual); + Assertions.assertTrue(replicationCheck.isReplicationDelayed()); + Assertions.assertThrows(NotReadyException.class, () -> replicationCheck.checkReplication()); } @Test @@ -199,8 +201,8 @@ public void shouldRunQueryAndReturnReplicationNotDelayed() throws Exception { final Instant actual = replicationCheck.queryReplication(); - Assert.assertEquals(expectedLastReplication, actual); - Assert.assertFalse(replicationCheck.isReplicationDelayed()); + Assertions.assertEquals(expectedLastReplication, actual); + Assertions.assertFalse(replicationCheck.isReplicationDelayed()); replicationCheck.checkReplication(); } } diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/options/ArrayHandlingModeTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/options/ArrayHandlingModeTest.java index f6b37592..47105d3a 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/options/ArrayHandlingModeTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/options/ArrayHandlingModeTest.java @@ -20,26 +20,30 @@ package com.spotify.dbeam.options; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; public class ArrayHandlingModeTest { @Test public void validValuesShouldPass() { - Assert.assertEquals(ArrayHandlingMode.Bytes, - ArrayHandlingMode.validateValue(ArrayHandlingMode.Bytes)); - Assert.assertEquals(ArrayHandlingMode.TypedMetaFromFirstRow, + Assertions.assertEquals( + ArrayHandlingMode.Bytes, ArrayHandlingMode.validateValue(ArrayHandlingMode.Bytes)); + Assertions.assertEquals( + ArrayHandlingMode.TypedMetaFromFirstRow, ArrayHandlingMode.validateValue(ArrayHandlingMode.TypedMetaFromFirstRow)); - Assert.assertEquals(ArrayHandlingMode.TypedMetaPostgres, + Assertions.assertEquals( + ArrayHandlingMode.TypedMetaPostgres, ArrayHandlingMode.validateValue(ArrayHandlingMode.TypedMetaPostgres)); } @Test public void invalidValueShouldThrow() { - RuntimeException thrown = Assert.assertThrows(RuntimeException.class, - () -> ArrayHandlingMode.validateValue("invalid")); - Assert.assertEquals("Invalid value 'invalid' for array handling mode. Allowed values: " - + "[bytes, typed_first_row, typed_postgres]", + RuntimeException thrown = + Assertions.assertThrows( + RuntimeException.class, () -> ArrayHandlingMode.validateValue("invalid")); + Assertions.assertEquals( + "Invalid value 'invalid' for array handling mode. Allowed values: " + + "[bytes, typed_first_row, typed_postgres]", thrown.getMessage()); } } diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/options/InputAvroSchemaTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/options/InputAvroSchemaTest.java index 1b03dcc9..8d93bd6e 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/options/InputAvroSchemaTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/options/InputAvroSchemaTest.java @@ -37,11 +37,11 @@ import org.apache.avro.Schema; import org.apache.avro.SchemaParseException; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.junit.AfterClass; -import org.junit.Assert; -import org.junit.BeforeClass; -import org.junit.Ignore; -import org.junit.Test; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; public class InputAvroSchemaTest { @@ -49,7 +49,7 @@ public class InputAvroSchemaTest { private static Path avroSchemaFilePath; private static String avroSchemaFilePathStr; - @BeforeClass + @BeforeAll public static void beforeAll() throws IOException { final String jsonSchema = "{\n" @@ -104,7 +104,7 @@ private Schema createRecordSchema( return inputSchema; } - @AfterClass + @AfterAll public static void afterAll() throws IOException { Files.delete(avroSchemaFile.toPath()); } @@ -115,17 +115,17 @@ public void checkReadAvroSchema() throws IOException { PipelineOptionsFactory.create().as(JdbcExportPipelineOptions.class); options.setAvroSchemaFilePath(avroSchemaFilePathStr); - Assert.assertEquals(avroSchemaFilePathStr, options.getAvroSchemaFilePath()); + Assertions.assertEquals(avroSchemaFilePathStr, options.getAvroSchemaFilePath()); final Schema inputSchema = BeamJdbcAvroSchema.parseInputAvroSchemaFile(avroSchemaFilePathStr); - Assert.assertEquals("Record description", inputSchema.getDoc()); - Assert.assertEquals("Field1 description", inputSchema.getField("field1").doc()); - Assert.assertEquals("Field2 description", inputSchema.getField("field2").doc()); + Assertions.assertEquals("Record description", inputSchema.getDoc()); + Assertions.assertEquals("Field1 description", inputSchema.getField("field1").doc()); + Assertions.assertEquals("Field2 description", inputSchema.getField("field2").doc()); } @Test - @Ignore + @Disabled public void checkFullPath() { // TODO // Check provide input string to args and verify final schema @@ -146,8 +146,8 @@ public void checkReadAvroSchemaWithEmptyParameter() throws IOException { PipelineOptionsFactory.create().as(JdbcExportPipelineOptions.class); options.setAvroSchemaFilePath(path); - Assert.assertEquals(path, options.getAvroSchemaFilePath()); - Assert.assertEquals( + Assertions.assertEquals(path, options.getAvroSchemaFilePath()); + Assertions.assertEquals( Optional.empty(), BeamJdbcAvroSchema.parseOptionalInputAvroSchemaFile(path)); } @@ -158,12 +158,12 @@ public void checkReadAvroSchemaWithNullParameter() throws IOException { PipelineOptionsFactory.create().as(JdbcExportPipelineOptions.class); options.setAvroSchemaFilePath(path); - Assert.assertEquals(path, options.getAvroSchemaFilePath()); - Assert.assertEquals( + Assertions.assertEquals(path, options.getAvroSchemaFilePath()); + Assertions.assertEquals( Optional.empty(), BeamJdbcAvroSchema.parseOptionalInputAvroSchemaFile(path)); } - @Test(expected = SchemaParseException.class) + @Test public void checkReadAvroSchemaWithInvalidFormat() throws IOException { final String invalidJson = "{"; final File invalidFile = createTestAvroSchemaFile(invalidJson); @@ -172,19 +172,23 @@ public void checkReadAvroSchemaWithInvalidFormat() throws IOException { PipelineOptionsFactory.create().as(JdbcExportPipelineOptions.class); options.setAvroSchemaFilePath(path); - Assert.assertEquals(path, options.getAvroSchemaFilePath()); - BeamJdbcAvroSchema.parseOptionalInputAvroSchemaFile(path); + Assertions.assertEquals(path, options.getAvroSchemaFilePath()); + Assertions.assertThrows( + SchemaParseException.class, + () -> BeamJdbcAvroSchema.parseOptionalInputAvroSchemaFile(path)); } - @Test(expected = FileNotFoundException.class) - public void checkReadAvroSchemaWithNonExistentFile() throws IOException { + @Test + public void checkReadAvroSchemaWithNonExistentFile() { final String path = "non_existent_schema.avsc"; final JdbcExportPipelineOptions options = PipelineOptionsFactory.create().as(JdbcExportPipelineOptions.class); options.setAvroSchemaFilePath(path); - Assert.assertEquals(path, options.getAvroSchemaFilePath()); - BeamJdbcAvroSchema.parseOptionalInputAvroSchemaFile(path); + Assertions.assertEquals(path, options.getAvroSchemaFilePath()); + Assertions.assertThrows( + FileNotFoundException.class, + () -> BeamJdbcAvroSchema.parseOptionalInputAvroSchemaFile(path)); } @Test @@ -194,7 +198,7 @@ public void checkValidCommandLineArgIsParsedAsOptions() throws IOException, SQLE "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--avroSchemaFilePath=/temp/record1.avsc --partition=2027-07-31"); - Assert.assertEquals("/temp/record1.avsc", options.getAvroSchemaFilePath()); + Assertions.assertEquals("/temp/record1.avsc", options.getAvroSchemaFilePath()); } @Test @@ -204,7 +208,7 @@ public void checkEmptyCommandLineArgIsParsedAsOptions() throws IOException, SQLE "--connectionUrl=jdbc:postgresql://some_db --table=some_table " + "--partition=2027-07-31"); - Assert.assertNull(options.getAvroSchemaFilePath()); + Assertions.assertNull(options.getAvroSchemaFilePath()); } private QueryBuilderArgs pareOptions(String cmdLineArgs) throws IOException { diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/options/JobNameConfigurationTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/options/JobNameConfigurationTest.java index 983b337e..6995fcb4 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/options/JobNameConfigurationTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/options/JobNameConfigurationTest.java @@ -26,8 +26,8 @@ import org.apache.beam.sdk.options.ApplicationNameOptions; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; public class JobNameConfigurationTest { @@ -38,7 +38,7 @@ public void shouldConfigureJobName() { JobNameConfiguration.configureJobName(pipelineOptions, "some_db", "some_table"); - Assert.assertEquals( + Assertions.assertEquals( "JdbcAvroJob", pipelineOptions.as(ApplicationNameOptions.class).getAppName()); assertThat(pipelineOptions.getJobName(), startsWith("dbeam-somedb-sometable-")); } @@ -50,7 +50,7 @@ public void shouldConfigureJobNameWhenJobNameIsAuto() { JobNameConfiguration.configureJobName(pipelineOptions, "some_db", "some_table"); - Assert.assertEquals( + Assertions.assertEquals( "JdbcAvroJob", pipelineOptions.as(ApplicationNameOptions.class).getAppName()); assertThat(pipelineOptions.getJobName(), startsWith("dbeam-somedb-sometable-")); } @@ -62,10 +62,10 @@ public void shouldConfigureJobNameWithEmptyTableName() { JobNameConfiguration.configureJobName(pipelineOptions, "some_db", null); - Assert.assertEquals( + Assertions.assertEquals( "JdbcAvroJob", pipelineOptions.as(ApplicationNameOptions.class).getAppName()); assertThat(pipelineOptions.getJobName(), startsWith("dbeam-somedb-")); - Assert.assertEquals(3, pipelineOptions.getJobName().split("-").length); + Assertions.assertEquals(3, pipelineOptions.getJobName().split("-").length); } } diff --git a/dbeam-core/src/test/java/com/spotify/dbeam/options/PasswordReaderTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/options/PasswordReaderTest.java index 25c0475f..17eff000 100644 --- a/dbeam-core/src/test/java/com/spotify/dbeam/options/PasswordReaderTest.java +++ b/dbeam-core/src/test/java/com/spotify/dbeam/options/PasswordReaderTest.java @@ -32,23 +32,23 @@ import java.util.Optional; import org.apache.beam.sdk.extensions.gcp.auth.NoopCredentialFactory; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.junit.AfterClass; -import org.junit.Assert; -import org.junit.BeforeClass; -import org.junit.Test; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; public class PasswordReaderTest { private static File passwordFile; - @BeforeClass + @BeforeAll public static void beforeAll() throws IOException { passwordFile = File.createTempFile("pattern", ".suffix"); passwordFile.deleteOnExit(); Files.write(passwordFile.toPath(), "something_encrypted".getBytes(), StandardOpenOption.CREATE); } - @AfterClass + @AfterAll public static void afterAll() throws IOException { Files.delete(passwordFile.toPath()); } @@ -83,6 +83,6 @@ public void shouldDecryptPasswordOnpasswordFileKmsEncryptedParameter() throws IO final Optional actualPassword = passwordReader.readPassword(options); - Assert.assertEquals(Optional.of("something_decrypted"), actualPassword); + Assertions.assertEquals(Optional.of("something_decrypted"), actualPassword); } } diff --git a/dbeam-parquet/pom.xml b/dbeam-parquet/pom.xml new file mode 100644 index 00000000..adfe8e3f --- /dev/null +++ b/dbeam-parquet/pom.xml @@ -0,0 +1,105 @@ + + + + 4.0.0 + + + com.spotify + dbeam-parent + 0.10.30-SNAPSHOT + + dbeam-parquet + + DBeam Parquet + DBeam JDBC to Parquet export + + + + com.spotify + dbeam-core + ${project.version} + + + com.google.auto.value + auto-value-annotations + provided + + + + + org.apache.parquet + parquet-hadoop + ${parquet.version} + + + org.apache.hadoop + * + + + + + + org.apache.hadoop + hadoop-common + ${hadoop.version} + + + + + org.slf4j + slf4j-api + + + org.slf4j + slf4j-jdk14 + runtime + + + + + org.apache.hadoop + hadoop-mapreduce-client-core + ${hadoop.version} + test + + + + + com.spotify + dbeam-core + ${project.version} + test-jar + test + + + org.junit.jupiter + junit-jupiter + test + + + org.hamcrest + hamcrest-core + test + + + org.hamcrest + hamcrest-library + test + + + org.mockito + mockito-core + test + + + com.h2database + h2 + test + + + + + com.spotify.dbeam.parquet.JdbcParquetJob + + + diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java new file mode 100644 index 00000000..8b70a4f3 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java @@ -0,0 +1,82 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import com.google.common.io.CharStreams; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.channels.Channels; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.Optional; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.io.FileSystems; +import org.apache.beam.sdk.io.fs.MatchResult; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.values.TypeDescriptors; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class BeamJdbcParquetSchema { + + private static final Logger LOGGER = LoggerFactory.getLogger(BeamJdbcParquetSchema.class); + + public static void exposeSchemaMetrics(final Pipeline pipeline, final long elapsedNanos) { + final long elapsedMs = elapsedNanos / 1000000; + LOGGER.info("Elapsed time to schema {} seconds", elapsedMs / 1000.0); + final Counter cnt = + Metrics.counter(BeamJdbcParquetSchema.class.getCanonicalName(), "schemaElapsedTimeMs"); + pipeline + .apply( + "ExposeSchemaCountersSeed", + Create.of(Collections.singletonList(0)).withType(TypeDescriptors.integers())) + .apply( + "ExposeSchemaCounters", + MapElements.into(TypeDescriptors.integers()) + .via( + v -> { + cnt.inc(elapsedMs); + return v; + })); + } + + public static Optional parseOptionalInputParquetSchemaFile(final String filename) + throws IOException { + if (filename == null || filename.isEmpty()) { + return Optional.empty(); + } + return Optional.of(parseInputParquetSchemaFile(filename)); + } + + public static MessageType parseInputParquetSchemaFile(final String filename) throws IOException { + final MatchResult.Metadata m = FileSystems.matchSingleFileSpec(filename); + try (InputStream inputStream = Channels.newInputStream(FileSystems.open(m.resourceId())); + InputStreamReader reader = new InputStreamReader(inputStream, StandardCharsets.UTF_8)) { + return MessageTypeParser.parseMessageType(CharStreams.toString(reader)); + } + } +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BenchJdbcParquetJob.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BenchJdbcParquetJob.java new file mode 100644 index 00000000..62920841 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BenchJdbcParquetJob.java @@ -0,0 +1,46 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import com.spotify.dbeam.jobs.BenchJdbcJob; +import com.spotify.dbeam.jobs.ExceptionHandling; +import org.apache.beam.sdk.options.PipelineOptionsFactory; + +/** Used on e2e test, allows benchmarking with different configuration parameters. */ +public class BenchJdbcParquetJob { + + public static BenchJdbcJob create(final String[] cmdLineArgs) { + PipelineOptionsFactory.register(BenchJdbcJob.BenchJdbcOptions.class); + return BenchJdbcJob.create( + "BenchJdbcParquetJob", + cmdLineArgs, + JdbcParquetJob::buildPipelineOptions, + (opts, output) -> JdbcParquetJob.create(opts, output).runExport()); + } + + public static void main(String[] cmdLineArgs) { + try { + create(cmdLineArgs).run(); + } catch (Exception e) { + ExceptionHandling.handleException(e); + } + } +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java new file mode 100644 index 00000000..5f8bfb16 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java @@ -0,0 +1,124 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import java.io.IOException; +import java.io.OutputStream; +import java.nio.channels.Channels; +import java.nio.channels.WritableByteChannel; +import org.apache.parquet.io.OutputFile; +import org.apache.parquet.io.PositionOutputStream; + +/** + * An {@link OutputFile} implementation backed by a {@link WritableByteChannel}, enabling Parquet + * writing without Hadoop dependencies. + */ +public class ChannelOutputFile implements OutputFile { + + private final WritableByteChannel channel; + private ChannelPositionOutputStream lastStream; + + public ChannelOutputFile(WritableByteChannel channel) { + this.channel = channel; + } + + /** Returns the number of bytes written, or 0 if no stream was created yet. */ + public long getBytesWritten() { + if (lastStream == null) { + return 0; + } + try { + return lastStream.getPos(); + } catch (IOException e) { + return 0; + } + } + + @Override + public PositionOutputStream create(long blockSizeHint) throws IOException { + lastStream = new ChannelPositionOutputStream(Channels.newOutputStream(channel)); + return lastStream; + } + + @Override + public PositionOutputStream createOrOverwrite(long blockSizeHint) throws IOException { + return create(blockSizeHint); + } + + @Override + public boolean supportsBlockSize() { + return false; + } + + @Override + public long defaultBlockSize() { + return 0; + } + + @Override + public String getPath() { + return channel.toString(); + } + + static class ChannelPositionOutputStream extends PositionOutputStream { + + private final OutputStream out; + private long position; + + ChannelPositionOutputStream(OutputStream out) { + this.out = out; + this.position = 0; + } + + @Override + public long getPos() throws IOException { + return position; + } + + @Override + public void write(int b) throws IOException { + out.write(b); + position++; + } + + @Override + public void write(byte[] b) throws IOException { + out.write(b); + position += b.length; + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + out.write(b, off, len); + position += len; + } + + @Override + public void flush() throws IOException { + out.flush(); + } + + @Override + public void close() throws IOException { + out.close(); + } + } +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java new file mode 100644 index 00000000..feca87f6 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java @@ -0,0 +1,129 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import com.google.auto.value.AutoValue; +import com.google.common.base.Preconditions; +import com.spotify.dbeam.args.JdbcAvroArgs; +import com.spotify.dbeam.args.JdbcConnectionArgs; +import java.io.Serializable; +import java.util.Collections; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; + +@AutoValue +public abstract class JdbcParquetArgs implements Serializable { + + private static final long serialVersionUID = 774966613L; + + public abstract JdbcConnectionArgs jdbcConnectionConfiguration(); + + @Nullable + public abstract JdbcAvroArgs.StatementPreparator statementPreparator(); + + public abstract int fetchSize(); + + public abstract String parquetCodec(); + + public abstract long rowGroupSize(); + + public abstract long pageSize(); + + public abstract List preCommand(); + + public abstract String arrayMode(); + + abstract Builder builder(); + + public CompressionCodecName getCompressionCodecName() { + switch (parquetCodec().toLowerCase()) { + case "snappy": + return CompressionCodecName.SNAPPY; + case "gzip": + return CompressionCodecName.GZIP; + case "zstd": + return CompressionCodecName.ZSTD; + case "lz4": + return CompressionCodecName.LZ4_RAW; + case "none": + case "uncompressed": + return CompressionCodecName.UNCOMPRESSED; + default: + throw new IllegalArgumentException("Invalid parquetCodec: " + parquetCodec()); + } + } + + @AutoValue.Builder + abstract static class Builder { + + abstract Builder setJdbcConnectionConfiguration(JdbcConnectionArgs jdbcConnectionArgs); + + abstract Builder setStatementPreparator(JdbcAvroArgs.StatementPreparator statementPreparator); + + abstract Builder setFetchSize(int fetchSize); + + abstract Builder setParquetCodec(String parquetCodec); + + abstract Builder setRowGroupSize(long rowGroupSize); + + abstract Builder setPageSize(long pageSize); + + abstract Builder setPreCommand(List preCommand); + + abstract Builder setArrayMode(String arrayMode); + + abstract JdbcParquetArgs build(); + } + + public static JdbcParquetArgs create( + final JdbcConnectionArgs jdbcConnectionArgs, + final int fetchSize, + final String parquetCodec, + final long rowGroupSize, + final long pageSize, + final List preCommand, + final String arrayMode) { + Preconditions.checkArgument( + parquetCodec.matches("snappy|gzip|zstd|lz4|none|uncompressed"), + "Parquet codec should be one of: snappy, gzip, zstd, lz4, none, uncompressed"); + return new AutoValue_JdbcParquetArgs.Builder() + .setJdbcConnectionConfiguration(jdbcConnectionArgs) + .setFetchSize(fetchSize) + .setParquetCodec(parquetCodec) + .setRowGroupSize(rowGroupSize) + .setPageSize(pageSize) + .setPreCommand(preCommand) + .setArrayMode(arrayMode) + .build(); + } + + public static JdbcParquetArgs create(final JdbcConnectionArgs jdbcConnectionArgs) { + return create( + jdbcConnectionArgs, + 10000, + "snappy", + 128 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); + } +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java new file mode 100644 index 00000000..cae8cc03 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java @@ -0,0 +1,320 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import static com.google.common.base.Preconditions.checkArgument; + +import com.google.common.collect.ImmutableMap; +import com.spotify.dbeam.avro.JdbcMetering; +import java.nio.channels.WritableByteChannel; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.Collections; +import java.util.Map; +import org.apache.beam.sdk.io.DefaultFilenamePolicy; +import org.apache.beam.sdk.io.DynamicFileDestinations; +import org.apache.beam.sdk.io.FileBasedSink; +import org.apache.beam.sdk.io.ShardNameTemplate; +import org.apache.beam.sdk.io.WriteFiles; +import org.apache.beam.sdk.io.WriteFilesResult; +import org.apache.beam.sdk.io.fs.ResourceId; +import org.apache.beam.sdk.options.ValueProvider; +import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.SerializableFunctions; +import org.apache.beam.sdk.util.MimeTypes; +import org.apache.beam.sdk.values.PCollection; +import org.apache.hadoop.conf.Configuration; +import org.apache.parquet.hadoop.ParquetFileWriter; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.api.WriteSupport; +import org.apache.parquet.io.OutputFile; +import org.apache.parquet.io.api.RecordConsumer; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@SuppressWarnings("checkstyle:AbbreviationAsWordInName") +public class JdbcParquetIO { + + private static final String DEFAULT_SHARD_TEMPLATE = ShardNameTemplate.INDEX_OF_MAX; + + public static final String PARQUET_AVRO_SCHEMA_KEY = "parquet.avro.schema"; + + public static PTransform, WriteFilesResult> createWrite( + final String filenamePrefix, + final String filenameSuffix, + final MessageType schema, + final JdbcParquetArgs jdbcParquetArgs, + final String avroSchemaJson) { + final ValueProvider prefixProvider = + StaticValueProvider.of( + FileBasedSink.convertToFileResourceIfPossible( + filenamePrefix.replaceAll("/+$", "") + "/part")); + final FileBasedSink.FilenamePolicy filenamePolicy = + DefaultFilenamePolicy.fromStandardParameters( + prefixProvider, DEFAULT_SHARD_TEMPLATE, filenameSuffix, false); + + final FileBasedSink.DynamicDestinations destinations = + DynamicFileDestinations.constant(filenamePolicy, SerializableFunctions.identity()); + + // Store schema as String for serialization (MessageType is not Serializable) + final String schemaString = schema.toString(); + final FileBasedSink sink = + new JdbcParquetSink( + prefixProvider, destinations, schemaString, jdbcParquetArgs, avroSchemaJson); + return WriteFiles.to(sink); + } + + static class JdbcParquetSink extends FileBasedSink { + + private static final long serialVersionUID = 937707428039L; + private final String schemaString; + private final JdbcParquetArgs jdbcParquetArgs; + private final String avroSchemaJson; + + JdbcParquetSink( + final ValueProvider filenamePrefix, + final DynamicDestinations destinations, + final String schemaString, + final JdbcParquetArgs jdbcParquetArgs, + final String avroSchemaJson) { + super(filenamePrefix, destinations); + this.schemaString = schemaString; + this.jdbcParquetArgs = jdbcParquetArgs; + this.avroSchemaJson = avroSchemaJson; + } + + @Override + public WriteOperation createWriteOperation() { + return new JdbcParquetWriteOperation(this, schemaString, jdbcParquetArgs, avroSchemaJson); + } + } + + private static class JdbcParquetWriteOperation + extends FileBasedSink.WriteOperation { + + private static final long serialVersionUID = 305340251351L; + private final String schemaString; + private final JdbcParquetArgs jdbcParquetArgs; + private final String avroSchemaJson; + + private JdbcParquetWriteOperation( + final FileBasedSink sink, + final String schemaString, + final JdbcParquetArgs jdbcParquetArgs, + final String avroSchemaJson) { + super(sink); + this.schemaString = schemaString; + this.jdbcParquetArgs = jdbcParquetArgs; + this.avroSchemaJson = avroSchemaJson; + } + + @Override + public FileBasedSink.Writer createWriter() { + return new JdbcParquetWriter(this, schemaString, jdbcParquetArgs, avroSchemaJson); + } + } + + private static class JdbcParquetWriter extends FileBasedSink.Writer { + private static final Logger LOGGER = LoggerFactory.getLogger(JdbcParquetWriter.class); + private final String schemaString; + private final JdbcParquetArgs jdbcParquetArgs; + private final String avroSchemaJson; + private ParquetWriter parquetWriter; + private Connection connection; + private JdbcMetering metering; + private ChannelOutputFile channelOutputFile; + + JdbcParquetWriter( + FileBasedSink.WriteOperation writeOperation, + String schemaString, + JdbcParquetArgs jdbcParquetArgs, + String avroSchemaJson) { + super(writeOperation, MimeTypes.BINARY); + this.schemaString = schemaString; + this.jdbcParquetArgs = jdbcParquetArgs; + this.avroSchemaJson = avroSchemaJson; + this.metering = JdbcMetering.create("jdbcparquetio"); + } + + public Void getDestination() { + return null; + } + + @Override + protected void prepareWrite(final WritableByteChannel channel) throws Exception { + LOGGER.info("jdbcparquetio : Preparing write..."); + connection = jdbcParquetArgs.jdbcConnectionConfiguration().createConnection(); + + final MessageType schema = MessageTypeParser.parseMessageType(schemaString); + channelOutputFile = new ChannelOutputFile(channel); + parquetWriter = + new ResultSetParquetWriterBuilder( + channelOutputFile, schema, avroSchemaJson, jdbcParquetArgs.arrayMode()) + .withCompressionCodec(jdbcParquetArgs.getCompressionCodecName()) + .withRowGroupSize(jdbcParquetArgs.rowGroupSize()) + .withPageSize((int) jdbcParquetArgs.pageSize()) + .withWriteMode(ParquetFileWriter.Mode.CREATE) + .build(); + LOGGER.info("jdbcparquetio : Write prepared"); + } + + private ResultSet executeQuery(final String query) throws Exception { + checkArgument(connection != null, "JDBC connection was not properly created"); + final PreparedStatement statement = + connection.prepareStatement( + query, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); + statement.setFetchSize(jdbcParquetArgs.fetchSize()); + if (jdbcParquetArgs.statementPreparator() != null) { + jdbcParquetArgs.statementPreparator().setParameters(statement); + } + + if (jdbcParquetArgs.preCommand() != null && !jdbcParquetArgs.preCommand().isEmpty()) { + final Statement stmt = connection.createStatement(); + for (String command : jdbcParquetArgs.preCommand()) { + stmt.execute(command); + } + } + + final long startTime = System.nanoTime(); + LOGGER.info( + "jdbcparquetio : Executing query with fetchSize={} (this might take a few minutes) ...", + statement.getFetchSize()); + final ResultSet resultSet = statement.executeQuery(); + this.metering.exposeExecuteQueryMs((System.nanoTime() - startTime) / 1000000L); + checkArgument(resultSet != null, "JDBC resultSet was not properly created"); + return resultSet; + } + + @Override + public void write(final String query) throws Exception { + checkArgument(parquetWriter != null, "ParquetWriter was not properly created"); + LOGGER.info("jdbcparquetio : Starting write..."); + try (ResultSet resultSet = executeQuery(query)) { + metering.startWriteMeter(); + while (resultSet.next()) { + parquetWriter.write(resultSet); + this.metering.incrementRecordCount(); + } + this.metering.exposeWriteElapsed(); + if (channelOutputFile != null) { + this.metering.exposeWrittenBytes(channelOutputFile.getBytesWritten()); + } + } + } + + @Override + protected void finishWrite() throws Exception { + LOGGER.info("jdbcparquetio : Closing connection, flushing writer..."); + if (parquetWriter != null) { + parquetWriter.close(); + } + if (connection != null) { + connection.close(); + } + LOGGER.info("jdbcparquetio : Write finished"); + } + } + + static class ResultSetParquetWriterBuilder + extends ParquetWriter.Builder { + + private final MessageType schema; + private final String avroSchemaJson; + private final String arrayMode; + + ResultSetParquetWriterBuilder(OutputFile outputFile, MessageType schema) { + this(outputFile, schema, null, "typed_first_row"); + } + + ResultSetParquetWriterBuilder( + OutputFile outputFile, MessageType schema, String avroSchemaJson) { + this(outputFile, schema, avroSchemaJson, "typed_first_row"); + } + + ResultSetParquetWriterBuilder( + OutputFile outputFile, MessageType schema, String avroSchemaJson, String arrayMode) { + super(outputFile); + this.schema = schema; + this.avroSchemaJson = avroSchemaJson; + this.arrayMode = arrayMode; + } + + @Override + protected ResultSetParquetWriterBuilder self() { + return this; + } + + @Override + protected WriteSupport getWriteSupport(Configuration conf) { + return new ResultSetWriteSupport(schema, avroSchemaJson, arrayMode); + } + } + + static class ResultSetWriteSupport extends WriteSupport { + + private final MessageType schema; + private final String avroSchemaJson; + private final String arrayMode; + private JdbcParquetWriteSupport writeSupport; + private RecordConsumer recordConsumer; + private boolean initialized = false; + + ResultSetWriteSupport(MessageType schema, String avroSchemaJson, String arrayMode) { + this.schema = schema; + this.avroSchemaJson = avroSchemaJson; + this.arrayMode = arrayMode; + } + + @Override + public WriteContext init(Configuration configuration) { + final Map extraMetadata; + if (avroSchemaJson != null && !avroSchemaJson.isEmpty()) { + extraMetadata = ImmutableMap.of(PARQUET_AVRO_SCHEMA_KEY, avroSchemaJson); + } else { + extraMetadata = Collections.emptyMap(); + } + return new WriteContext(schema, extraMetadata); + } + + @Override + public void prepareForWrite(RecordConsumer recordConsumer) { + this.recordConsumer = recordConsumer; + } + + @Override + public void write(ResultSet resultSet) { + try { + if (!initialized) { + this.writeSupport = JdbcParquetWriteSupport.create(resultSet, schema, arrayMode); + initialized = true; + } + writeSupport.writeRecord(recordConsumer, resultSet); + } catch (Exception e) { + throw new RuntimeException("Failed to write ResultSet row to Parquet", e); + } + } + } +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java new file mode 100644 index 00000000..084e8c3f --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -0,0 +1,287 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import com.google.common.base.Preconditions; +import com.spotify.dbeam.args.JdbcExportArgs; +import com.spotify.dbeam.avro.JdbcAvroSchema; +import com.spotify.dbeam.avro.JdbcMetering; +import com.spotify.dbeam.beam.BeamHelper; +import com.spotify.dbeam.beam.MetricsHelper; +import com.spotify.dbeam.jobs.ExceptionHandling; +import com.spotify.dbeam.jobs.FailedValidationException; +import com.spotify.dbeam.options.DBeamPipelineOptions; +import com.spotify.dbeam.options.JdbcExportArgsFactory; +import com.spotify.dbeam.options.JdbcExportPipelineOptions; +import com.spotify.dbeam.options.JobNameConfiguration; +import com.spotify.dbeam.options.OutputOptions; +import java.io.IOException; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.List; +import java.util.Map; +import org.apache.avro.Schema; +import org.apache.beam.runners.direct.DirectOptions; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.Create; +import org.apache.parquet.schema.MessageType; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class JdbcParquetJob { + + private static final Logger LOGGER = LoggerFactory.getLogger(JdbcParquetJob.class); + + private final PipelineOptions pipelineOptions; + private final Pipeline pipeline; + private final JdbcExportArgs jdbcExportArgs; + private final String output; + private final boolean dataOnly; + private final long minRows; + + public JdbcParquetJob( + final PipelineOptions pipelineOptions, + final Pipeline pipeline, + final JdbcExportArgs jdbcExportArgs, + final String output, + final boolean dataOnly, + final long minRows) { + this.pipelineOptions = pipelineOptions; + this.pipeline = pipeline; + this.jdbcExportArgs = jdbcExportArgs; + this.output = output; + this.dataOnly = dataOnly; + this.minRows = minRows; + Preconditions.checkArgument( + this.output != null && this.output.length() > 0, "'output' must be defined"); + } + + public static JdbcParquetJob create(final PipelineOptions pipelineOptions, final String output) + throws IOException, ClassNotFoundException { + pipelineOptions.as(DirectOptions.class).setBlockOnRun(false); + return new JdbcParquetJob( + pipelineOptions, + Pipeline.create(pipelineOptions), + JdbcExportArgsFactory.fromPipelineOptions(pipelineOptions), + output, + pipelineOptions.as(OutputOptions.class).getDataOnly(), + pipelineOptions.as(JdbcExportPipelineOptions.class).getMinRows()); + } + + public static JdbcParquetJob create(final PipelineOptions pipelineOptions) + throws IOException, ClassNotFoundException { + return create(pipelineOptions, pipelineOptions.as(OutputOptions.class).getOutput()); + } + + public static JdbcParquetJob create(final String[] cmdLineArgs) + throws IOException, ClassNotFoundException { + return create(buildPipelineOptions(cmdLineArgs)); + } + + public static PipelineOptions buildPipelineOptions(final String[] cmdLineArgs) { + PipelineOptionsFactory.register(JdbcExportPipelineOptions.class); + PipelineOptionsFactory.register(OutputOptions.class); + PipelineOptionsFactory.register(ParquetPipelineOptions.class); + return PipelineOptionsFactory.fromArgs(cmdLineArgs).withValidation().create(); + } + + private void configureVersion() { + final String dbeamVersion = this.getClass().getPackage().getImplementationVersion(); + LOGGER.info( + "{} {} version {}", + this.getClass().getPackage().getImplementationTitle(), + this.getClass().getSimpleName(), + dbeamVersion); + pipelineOptions.as(DBeamPipelineOptions.class).setDBeamVersion(dbeamVersion); + } + + public void prepareExport() throws Exception { + configureVersion(); + final List queries; + final MessageType generatedSchema; + final Schema avroSchema; + final String arrayMode = jdbcExportArgs.jdbcAvroOptions().arrayMode(); + try (Connection connection = jdbcExportArgs.createConnection()) { + final String schemaFilePath = + pipelineOptions.as(ParquetPipelineOptions.class).getParquetSchemaFilePath(); + final java.util.Optional inputSchema = + BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(schemaFilePath); + if (inputSchema.isPresent()) { + final long startTime = System.nanoTime(); + generatedSchema = inputSchema.get(); + avroSchema = generateAvroSchema(connection); + BeamJdbcParquetSchema.exposeSchemaMetrics(pipeline, System.nanoTime() - startTime); + } else { + // Query DB once for both Parquet and Avro schemas + final long startTime = System.nanoTime(); + try (Statement stmt = connection.createStatement()) { + final ResultSet rs = + stmt.executeQuery(jdbcExportArgs.queryBuilderArgs().sqlQueryWithLimitOne()); + rs.next(); + generatedSchema = + JdbcParquetSchema.createParquetSchema( + rs, + jdbcExportArgs.avroSchemaName(), + jdbcExportArgs.useAvroLogicalTypes(), + arrayMode); + final String dbUrl = connection.getMetaData().getURL(); + final String avroDoc = + jdbcExportArgs + .avroDoc() + .orElseGet( + () -> String.format("Generate schema from JDBC ResultSet from %s", dbUrl)); + avroSchema = + JdbcAvroSchema.createAvroSchema( + rs, + jdbcExportArgs.avroSchemaNamespace(), + dbUrl, + jdbcExportArgs.avroSchemaName(), + avroDoc, + jdbcExportArgs.useAvroLogicalTypes(), + arrayMode, + jdbcExportArgs.jdbcAvroOptions().nullableArrayItems()); + } + BeamJdbcParquetSchema.exposeSchemaMetrics(pipeline, System.nanoTime() - startTime); + } + queries = jdbcExportArgs.queryBuilderArgs().buildQueries(connection); + + final String tableName = pipelineOptions.as(DBeamPipelineOptions.class).getTable(); + JobNameConfiguration.configureJobName( + pipeline.getOptions(), connection.getCatalog(), tableName); + } + if (!this.dataOnly) { + BeamHelper.saveStringOnSubPath(output, "/_PARQUET_SCHEMA.json", generatedSchema.toString()); + BeamHelper.saveStringOnSubPath(output, "/_AVRO_SCHEMA.avsc", avroSchema.toString(true)); + for (int i = 0; i < queries.size(); i++) { + BeamHelper.saveStringOnSubPath( + this.output, String.format("/_queries/query_%d.sql", i), queries.get(i)); + } + } + LOGGER.info("Running queries: {}", queries.toString()); + + final String parquetCodec = resolveParquetCodec(pipelineOptions); + final ParquetPipelineOptions parquetOptions = pipelineOptions.as(ParquetPipelineOptions.class); + final JdbcParquetArgs parquetArgs = + JdbcParquetArgs.create( + jdbcExportArgs.jdbcAvroOptions().jdbcConnectionConfiguration(), + jdbcExportArgs.jdbcAvroOptions().fetchSize(), + parquetCodec, + parquetOptions.getRowGroupSize(), + parquetOptions.getPageSize(), + jdbcExportArgs.jdbcAvroOptions().preCommand(), + arrayMode); + + pipeline + .apply("JdbcQueries", Create.of(queries)) + .apply( + "JdbcParquetSave", + JdbcParquetIO.createWrite( + output, ".parquet", generatedSchema, parquetArgs, avroSchema.toString())); + } + + private void checkMetrics(PipelineResult pipelineResult) throws FailedValidationException { + final Map metrics = MetricsHelper.getMetrics(pipelineResult); + if (!this.dataOnly) { + BeamHelper.saveMetrics(metrics, output); + } + final Long recordCount = metrics.getOrDefault(JdbcMetering.RECORD_COUNT_METRIC_NAME, 0L); + if (recordCount < this.minRows) { + throw new FailedValidationException( + String.format( + "Unexpected number of rows in the output: got %d, expecting at least %d", + recordCount, this.minRows)); + } + } + + public PipelineResult runAndWait() { + return BeamHelper.waitUntilDone(this.pipeline.run(), jdbcExportArgs.exportTimeout()); + } + + public PipelineResult runExport() throws Exception { + prepareExport(); + final PipelineResult pipelineResult = runAndWait(); + checkMetrics(pipelineResult); + return pipelineResult; + } + + private Schema generateAvroSchema(final Connection connection) throws Exception { + final String dbUrl = connection.getMetaData().getURL(); + final String avroDoc = + jdbcExportArgs + .avroDoc() + .orElseGet(() -> String.format("Generate schema from JDBC ResultSet from %s", dbUrl)); + return JdbcAvroSchema.createSchemaByReadingOneRow( + connection, + jdbcExportArgs.queryBuilderArgs(), + jdbcExportArgs.avroSchemaNamespace(), + jdbcExportArgs.avroSchemaName(), + avroDoc, + jdbcExportArgs.useAvroLogicalTypes(), + jdbcExportArgs.jdbcAvroOptions().arrayMode(), + jdbcExportArgs.jdbcAvroOptions().nullableArrayItems()); + } + + public Pipeline getPipeline() { + return pipeline; + } + + public JdbcExportArgs getJdbcExportArgs() { + return jdbcExportArgs; + } + + public String getOutput() { + return output; + } + + static String resolveParquetCodec(final PipelineOptions options) { + final String parquetCodec = options.as(ParquetPipelineOptions.class).getParquetCodec(); + if (parquetCodec != null && !parquetCodec.isEmpty()) { + return parquetCodec; + } + return mapAvroCodecToParquetCodec(options.as(JdbcExportPipelineOptions.class).getAvroCodec()); + } + + static String mapAvroCodecToParquetCodec(final String avroCodec) { + if (avroCodec == null) { + return "snappy"; + } + if (avroCodec.equals("snappy")) { + return "snappy"; + } else if (avroCodec.startsWith("deflate")) { + return "gzip"; + } else if (avroCodec.startsWith("zstandard")) { + return "zstd"; + } + return avroCodec; + } + + public static void main(String[] cmdLineArgs) { + try { + JdbcParquetJob.create(cmdLineArgs).runExport(); + } catch (Exception e) { + ExceptionHandling.handleException(e); + } + } +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetMetering.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetMetering.java new file mode 100644 index 00000000..c7baead0 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetMetering.java @@ -0,0 +1,38 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import com.spotify.dbeam.avro.JdbcMetering; + +/** + * @deprecated Use {@link JdbcMetering} directly instead. + */ +@Deprecated +public class JdbcParquetMetering extends JdbcMetering { + + public JdbcParquetMetering(int countReportEvery, int logEvery) { + super(countReportEvery, logEvery, "jdbcparquetio"); + } + + public static JdbcParquetMetering create() { + return new JdbcParquetMetering(100000, 100000); + } +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java new file mode 100644 index 00000000..92a03184 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -0,0 +1,276 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import static java.sql.Types.ARRAY; +import static java.sql.Types.BIGINT; +import static java.sql.Types.BINARY; +import static java.sql.Types.BIT; +import static java.sql.Types.BLOB; +import static java.sql.Types.BOOLEAN; +import static java.sql.Types.CHAR; +import static java.sql.Types.CLOB; +import static java.sql.Types.DATE; +import static java.sql.Types.DOUBLE; +import static java.sql.Types.FLOAT; +import static java.sql.Types.INTEGER; +import static java.sql.Types.LONGNVARCHAR; +import static java.sql.Types.LONGVARBINARY; +import static java.sql.Types.LONGVARCHAR; +import static java.sql.Types.NCHAR; +import static java.sql.Types.OTHER; +import static java.sql.Types.REAL; +import static java.sql.Types.SMALLINT; +import static java.sql.Types.TIME; +import static java.sql.Types.TIMESTAMP; +import static java.sql.Types.TIME_WITH_TIMEZONE; +import static java.sql.Types.TINYINT; +import static java.sql.Types.VARBINARY; +import static java.sql.Types.VARCHAR; + +import com.spotify.dbeam.args.QueryBuilderArgs; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; +import org.apache.parquet.schema.Types; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@SuppressWarnings("checkstyle:AbbreviationAsWordInName") +public class JdbcParquetSchema { + + private static final Logger LOGGER = LoggerFactory.getLogger(JdbcParquetSchema.class); + + public static MessageType createSchemaByReadingOneRow( + final Connection connection, + final QueryBuilderArgs queryBuilderArgs, + final Optional schemaName, + final boolean useLogicalTypes, + final String arrayMode) + throws SQLException { + LOGGER.debug("Creating Parquet schema based on the first read row from the database"); + try (Statement statement = connection.createStatement()) { + final ResultSet resultSet = statement.executeQuery(queryBuilderArgs.sqlQueryWithLimitOne()); + resultSet.next(); + final MessageType schema = + createParquetSchema(resultSet, schemaName, useLogicalTypes, arrayMode); + LOGGER.info( + "Parquet schema created successfully. useLogicalTypes={}. Generated schema: {}", + useLogicalTypes, + schema); + return schema; + } + } + + public static MessageType createParquetSchema( + final ResultSet resultSet, + final Optional maybeSchemaName, + final boolean useLogicalTypes, + final String arrayMode) + throws SQLException { + final ResultSetMetaData meta = resultSet.getMetaData(); + final String tableName = getDatabaseTableName(meta); + final String schemaName = maybeSchemaName.orElse(tableName); + + final List fields = new ArrayList<>(); + 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); + } + + final int columnType = meta.getColumnType(i); + final int precision = meta.getPrecision(i); + final String columnClassName = meta.getColumnClassName(i); + final String columnTypeName = meta.getColumnTypeName(i); + + fields.add( + buildParquetFieldType( + normalizeFieldName(columnName), + columnType, + precision, + columnClassName, + columnTypeName, + useLogicalTypes, + arrayMode)); + } + + return new MessageType(schemaName, fields); + } + + 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()) { + return normalizeFieldName(metaTableName); + } + } + return defaultTableName; + } + + static Type buildParquetFieldType( + final String columnName, + final int columnType, + final int precision, + final String columnClassName, + final String columnTypeName, + final boolean useLogicalTypes, + final String arrayMode) { + switch (columnType) { + case BIGINT: + return Types.optional(PrimitiveType.PrimitiveTypeName.INT64).named(columnName); + case INTEGER: + case SMALLINT: + case TINYINT: + if (Long.class.getCanonicalName().equals(columnClassName)) { + return Types.optional(PrimitiveType.PrimitiveTypeName.INT64).named(columnName); + } else { + return Types.optional(PrimitiveType.PrimitiveTypeName.INT32).named(columnName); + } + case TIMESTAMP: + case DATE: // stored as TIMESTAMP(MILLIS) for Avro-path compatibility + case TIME: + case TIME_WITH_TIMEZONE: + if (useLogicalTypes) { + return Types.optional(PrimitiveType.PrimitiveTypeName.INT64) + .as(LogicalTypeAnnotation.timestampType(true, LogicalTypeAnnotation.TimeUnit.MILLIS)) + .named(columnName); + } else { + return Types.optional(PrimitiveType.PrimitiveTypeName.INT64).named(columnName); + } + case BOOLEAN: + return Types.optional(PrimitiveType.PrimitiveTypeName.BOOLEAN).named(columnName); + case BIT: + if (precision <= 1) { + return Types.optional(PrimitiveType.PrimitiveTypeName.BOOLEAN).named(columnName); + } else { + return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY).named(columnName); + } + case BINARY: + case VARBINARY: + case LONGVARBINARY: + case BLOB: + return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY).named(columnName); + case DOUBLE: + return Types.optional(PrimitiveType.PrimitiveTypeName.DOUBLE).named(columnName); + case FLOAT: + case REAL: + return Types.optional(PrimitiveType.PrimitiveTypeName.FLOAT).named(columnName); + case ARRAY: + if ("bytes".equals(arrayMode)) { + return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY).named(columnName); + } + // Parquet 3-level LIST convention with typed elements. + // Element type is inferred from columnTypeName (e.g. _int4, _text for PostgreSQL). + return Types.optionalList() + .element(buildArrayElementType(columnTypeName)) + .named(columnName); + case OTHER: + if (useLogicalTypes && "uuid".equals(columnTypeName)) { + return Types.optional(PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY) + .length(16) + .as(LogicalTypeAnnotation.uuidType()) + .named(columnName); + } + // fall through to string + case VARCHAR: + case CHAR: + case CLOB: + case LONGNVARCHAR: + case LONGVARCHAR: + case NCHAR: + default: + return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named(columnName); + } + } + + /** + * Determine the Parquet element type for an ARRAY column based on the column type name. For + * PostgreSQL, array column type names are prefixed with underscore (e.g. _int4, _text). Falls + * back to STRING for unrecognized types. + */ + static Type buildArrayElementType(final String columnTypeName) { + final String elementType = resolveArrayElementTypeName(columnTypeName); + switch (elementType) { + case "int": + case "int4": + case "int2": + return Types.optional(PrimitiveType.PrimitiveTypeName.INT32).named("element"); + case "int8": + return Types.optional(PrimitiveType.PrimitiveTypeName.INT64).named("element"); + case "float4": + return Types.optional(PrimitiveType.PrimitiveTypeName.FLOAT).named("element"); + case "float8": + return Types.optional(PrimitiveType.PrimitiveTypeName.DOUBLE).named("element"); + case "bool": + return Types.optional(PrimitiveType.PrimitiveTypeName.BOOLEAN).named("element"); + default: + return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("element"); + } + } + + /** + * Extract the element type name from an array column type name. PostgreSQL uses underscore prefix + * (e.g. _int4 -> int4, _text -> text). H2 and others use "INTEGER ARRAY" style names. + */ + static String resolveArrayElementTypeName(final String columnTypeName) { + if (columnTypeName == null) { + return "text"; + } + if (columnTypeName.startsWith("_")) { + return columnTypeName.substring(1); + } + // H2 style: "INTEGER ARRAY", "VARCHAR ARRAY", etc. + final String lower = columnTypeName.toLowerCase(); + if (lower.startsWith("integer")) { + return "int4"; + } else if (lower.startsWith("bigint")) { + return "int8"; + } else if (lower.startsWith("real") || lower.startsWith("float")) { + return "float4"; + } else if (lower.startsWith("double")) { + return "float8"; + } else if (lower.startsWith("boolean")) { + return "bool"; + } + return "text"; + } + + static String normalizeFieldName(final String input) { + return input.replaceAll("[^A-Za-z0-9_]", "_"); + } +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java new file mode 100644 index 00000000..bbf0da05 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -0,0 +1,345 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import static java.sql.Types.ARRAY; +import static java.sql.Types.BIGINT; +import static java.sql.Types.BINARY; +import static java.sql.Types.BIT; +import static java.sql.Types.BLOB; +import static java.sql.Types.BOOLEAN; +import static java.sql.Types.CHAR; +import static java.sql.Types.CLOB; +import static java.sql.Types.DATE; +import static java.sql.Types.DOUBLE; +import static java.sql.Types.FLOAT; +import static java.sql.Types.INTEGER; +import static java.sql.Types.LONGNVARCHAR; +import static java.sql.Types.LONGVARBINARY; +import static java.sql.Types.LONGVARCHAR; +import static java.sql.Types.NCHAR; +import static java.sql.Types.OTHER; +import static java.sql.Types.REAL; +import static java.sql.Types.SMALLINT; +import static java.sql.Types.TIME; +import static java.sql.Types.TIMESTAMP; +import static java.sql.Types.TIME_WITH_TIMEZONE; +import static java.sql.Types.TINYINT; +import static java.sql.Types.VARBINARY; +import static java.sql.Types.VARCHAR; + +import java.nio.ByteBuffer; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.util.Calendar; +import java.util.GregorianCalendar; +import java.util.Objects; +import java.util.TimeZone; +import java.util.UUID; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.io.api.RecordConsumer; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.Type; + +/** + * Writes values from a JDBC ResultSet directly to a Parquet RecordConsumer. Each column is written + * using the appropriate typed method, skipping null fields (Parquet's standard null representation + * for optional fields). + */ +public class JdbcParquetWriteSupport { + + private static final ThreadLocal CALENDAR = + ThreadLocal.withInitial(() -> new GregorianCalendar(TimeZone.getTimeZone("UTC"))); + + @FunctionalInterface + interface ColumnWriter { + void write(RecordConsumer consumer, ResultSet rs) throws SQLException; + } + + private final MessageType schema; + private final ColumnWriter[] columnWriters; + private final int columnCount; + + public JdbcParquetWriteSupport( + MessageType schema, ColumnWriter[] columnWriters, int columnCount) { + this.schema = schema; + this.columnWriters = columnWriters; + this.columnCount = columnCount; + } + + public static JdbcParquetWriteSupport create( + ResultSet resultSet, MessageType schema, String arrayMode) throws SQLException { + final ResultSetMetaData meta = resultSet.getMetaData(); + final int columnCount = meta.getColumnCount(); + final ColumnWriter[] writers = new ColumnWriter[columnCount + 1]; + + for (int i = 1; i <= columnCount; i++) { + writers[i] = computeColumnWriter(meta, i, arrayMode, schema.getType(i - 1)); + } + + return new JdbcParquetWriteSupport(schema, writers, columnCount); + } + + public MessageType getSchema() { + return schema; + } + + /** Write the current row of the ResultSet to the RecordConsumer. */ + public void writeRecord(RecordConsumer consumer, ResultSet resultSet) throws SQLException { + consumer.startMessage(); + for (int i = 1; i <= columnCount; i++) { + columnWriters[i].write(consumer, resultSet); + } + consumer.endMessage(); + } + + static ColumnWriter computeColumnWriter( + final ResultSetMetaData meta, final int column, final String arrayMode, final Type fieldType) + throws SQLException { + final int columnType = meta.getColumnType(column); + final int fieldIndex = column - 1; + final String fieldName = + meta.getColumnName(column).isEmpty() + ? meta.getColumnLabel(column) + : meta.getColumnName(column); + final String normalizedName = JdbcParquetSchema.normalizeFieldName(fieldName); + + switch (columnType) { + case VARCHAR: + case CHAR: + case CLOB: + case LONGNVARCHAR: + case LONGVARCHAR: + case NCHAR: + return (consumer, rs) -> { + final String val = rs.getString(column); + if (val != null && !rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addBinary(Binary.fromString(val)); + consumer.endField(normalizedName, fieldIndex); + } + }; + case BIGINT: + return (consumer, rs) -> { + final long val = rs.getLong(column); + if (!rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addLong(val); + consumer.endField(normalizedName, fieldIndex); + } + }; + case INTEGER: + case SMALLINT: + case TINYINT: + if (Long.class.getCanonicalName().equals(meta.getColumnClassName(column))) { + return (consumer, rs) -> { + final long val = rs.getLong(column); + if (!rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addLong(val); + consumer.endField(normalizedName, fieldIndex); + } + }; + } + return (consumer, rs) -> { + final int val = rs.getInt(column); + if (!rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addInteger(val); + consumer.endField(normalizedName, fieldIndex); + } + }; + case TIMESTAMP: + case DATE: // written as epoch millis for Avro-path compatibility + case TIME: + case TIME_WITH_TIMEZONE: + return (consumer, rs) -> { + final Timestamp timestamp = rs.getTimestamp(column, CALENDAR.get()); + if (timestamp != null && !rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addLong(timestamp.getTime()); + consumer.endField(normalizedName, fieldIndex); + } + }; + case BOOLEAN: + return (consumer, rs) -> { + final boolean val = rs.getBoolean(column); + if (!rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addBoolean(val); + consumer.endField(normalizedName, fieldIndex); + } + }; + case BIT: + if (meta.getPrecision(column) <= 1) { + return (consumer, rs) -> { + final boolean val = rs.getBoolean(column); + if (!rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addBoolean(val); + consumer.endField(normalizedName, fieldIndex); + } + }; + } else { + return (consumer, rs) -> { + final byte[] val = rs.getBytes(column); + if (val != null && !rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addBinary(Binary.fromConstantByteArray(val)); + consumer.endField(normalizedName, fieldIndex); + } + }; + } + case ARRAY: + if ("bytes".equals(arrayMode)) { + return (consumer, rs) -> { + final byte[] val = rs.getBytes(column); + if (val != null && !rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addBinary(Binary.fromConstantByteArray(val)); + consumer.endField(normalizedName, fieldIndex); + } + }; + } + final String arrayElementType = + JdbcParquetSchema.resolveArrayElementTypeName(meta.getColumnTypeName(column)); + return (consumer, rs) -> { + final java.sql.Array sqlArray = rs.getArray(column); + if (sqlArray != null && !rs.wasNull()) { + final Object[] items = (Object[]) sqlArray.getArray(); + consumer.startField(normalizedName, fieldIndex); + consumer.startGroup(); // LIST group + consumer.startField("list", 0); // repeated list field + for (Object item : items) { + consumer.startGroup(); // list element group + if (item != null) { + consumer.startField("element", 0); + writeArrayElement(consumer, item, arrayElementType); + consumer.endField("element", 0); + } + consumer.endGroup(); + } + consumer.endField("list", 0); + consumer.endGroup(); + consumer.endField(normalizedName, fieldIndex); + } + }; + case BINARY: + case VARBINARY: + case LONGVARBINARY: + case BLOB: + return (consumer, rs) -> { + final byte[] val = rs.getBytes(column); + if (val != null && !rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addBinary(Binary.fromConstantByteArray(val)); + consumer.endField(normalizedName, fieldIndex); + } + }; + case DOUBLE: + return (consumer, rs) -> { + final double val = rs.getDouble(column); + if (!rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addDouble(val); + consumer.endField(normalizedName, fieldIndex); + } + }; + case FLOAT: + case REAL: + return (consumer, rs) -> { + final float val = rs.getFloat(column); + if (!rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addFloat(val); + consumer.endField(normalizedName, fieldIndex); + } + }; + case OTHER: + if (Objects.equals(meta.getColumnTypeName(column), "uuid")) { + final boolean isUuidLogicalType = + fieldType.getLogicalTypeAnnotation() != null + && fieldType.getLogicalTypeAnnotation().equals(LogicalTypeAnnotation.uuidType()); + if (isUuidLogicalType) { + return (consumer, rs) -> { + final Object val = rs.getObject(column); + if (val != null && !rs.wasNull()) { + final UUID uuid = + val instanceof UUID ? (UUID) val : UUID.fromString(val.toString()); + final ByteBuffer buf = ByteBuffer.allocate(16); + buf.putLong(uuid.getMostSignificantBits()); + buf.putLong(uuid.getLeastSignificantBits()); + consumer.startField(normalizedName, fieldIndex); + consumer.addBinary(Binary.fromConstantByteArray(buf.array())); + consumer.endField(normalizedName, fieldIndex); + } + }; + } + return (consumer, rs) -> { + final Object val = rs.getObject(column); + if (val != null && !rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addBinary(Binary.fromString(val.toString())); + consumer.endField(normalizedName, fieldIndex); + } + }; + } + // fall through to string + default: + return (consumer, rs) -> { + final String val = rs.getString(column); + if (val != null && !rs.wasNull()) { + consumer.startField(normalizedName, fieldIndex); + consumer.addBinary(Binary.fromString(val)); + consumer.endField(normalizedName, fieldIndex); + } + }; + } + } + + static void writeArrayElement(RecordConsumer consumer, Object item, String elementType) { + switch (elementType) { + case "int": + case "int4": + case "int2": + consumer.addInteger(((Number) item).intValue()); + break; + case "int8": + consumer.addLong(((Number) item).longValue()); + break; + case "float4": + consumer.addFloat(((Number) item).floatValue()); + break; + case "float8": + consumer.addDouble(((Number) item).doubleValue()); + break; + case "bool": + consumer.addBoolean((Boolean) item); + break; + default: + consumer.addBinary(Binary.fromString(item.toString())); + break; + } + } +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java new file mode 100644 index 00000000..8a176537 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java @@ -0,0 +1,57 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import org.apache.beam.sdk.options.Default; +import org.apache.beam.sdk.options.Description; +import org.apache.beam.sdk.options.PipelineOptions; + +@Description("Parquet-specific export options") +public interface ParquetPipelineOptions extends PipelineOptions { + + @Description("Path to file with a target Parquet schema (MessageType text format).") + String getParquetSchemaFilePath(); + + void setParquetSchemaFilePath(String value); + + @Description( + "Parquet compression codec (snappy, gzip, zstd, lz4, none). " + + "Overrides --avroCodec when set.") + String getParquetCodec(); + + void setParquetCodec(String value); + + @Description( + "Parquet row group size in bytes. Larger values improve read performance " + + "but use more memory during writes.") + @Default.Long(134217728) + Long getRowGroupSize(); + + void setRowGroupSize(Long value); + + @Description( + "Parquet page size in bytes. Controls the granularity of encoding " + + "and compression within a column chunk.") + @Default.Long(1048576) + Long getPageSize(); + + void setPageSize(Long value); +} diff --git a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/PsqlParquetJob.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/PsqlParquetJob.java new file mode 100644 index 00000000..964854ad --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/PsqlParquetJob.java @@ -0,0 +1,55 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import com.spotify.dbeam.jobs.ExceptionHandling; +import com.spotify.dbeam.jobs.PsqlReplicationCheck; +import java.io.IOException; + +public class PsqlParquetJob { + + private final JdbcParquetJob job; + private final PsqlReplicationCheck psqlReplicationCheck; + + public PsqlParquetJob(final JdbcParquetJob job, final PsqlReplicationCheck psqlReplicationCheck) { + this.job = job; + this.psqlReplicationCheck = psqlReplicationCheck; + } + + public static PsqlParquetJob create(final String[] cmdLineArgs) + throws IOException, ClassNotFoundException { + final JdbcParquetJob job = JdbcParquetJob.create(cmdLineArgs); + PsqlReplicationCheck.validateOptions(job.getJdbcExportArgs()); + final PsqlReplicationCheck psqlReplicationCheck = + PsqlReplicationCheck.create(job.getJdbcExportArgs()); + return new PsqlParquetJob(job, psqlReplicationCheck); + } + + public static void main(String[] cmdLineArgs) { + try { + final PsqlParquetJob psqlParquetJob = create(cmdLineArgs); + psqlParquetJob.psqlReplicationCheck.checkReplication(); + psqlParquetJob.job.runExport(); + } catch (Exception e) { + ExceptionHandling.handleException(e); + } + } +} diff --git a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java new file mode 100644 index 00000000..31fe33ac --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java @@ -0,0 +1,69 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Optional; +import org.apache.parquet.schema.MessageType; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class BeamJdbcParquetSchemaTest { + + @Test + public void shouldParseInputSchemaFile() throws IOException { + final Path schemaFile = Files.createTempFile("test-schema-", ".parquet.txt"); + Files.write( + schemaFile, + ("message test_table {\n" + + " optional int64 id;\n" + + " optional binary name (STRING);\n" + + "}") + .getBytes()); + + final Optional schema = + BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(schemaFile.toString()); + + Assertions.assertTrue(schema.isPresent()); + Assertions.assertEquals("test_table", schema.get().getName()); + Assertions.assertEquals(2, schema.get().getFieldCount()); + + Files.deleteIfExists(schemaFile); + } + + @Test + public void shouldReturnEmptyForNullFilename() throws IOException { + final Optional schema = + BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(null); + + Assertions.assertFalse(schema.isPresent()); + } + + @Test + public void shouldReturnEmptyForEmptyFilename() throws IOException { + final Optional schema = + BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(""); + + Assertions.assertFalse(schema.isPresent()); + } +} diff --git a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java new file mode 100644 index 00000000..f2425260 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java @@ -0,0 +1,103 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.greaterThan; + +import com.spotify.dbeam.DbTestHelper; +import com.spotify.dbeam.TestHelper; +import java.io.ByteArrayOutputStream; +import java.io.PrintStream; +import java.nio.file.Path; +import java.sql.SQLException; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +public class BenchJdbcParquetJobTest { + + private static final String CONNECTION_URL = + "jdbc:h2:mem:testbench;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; + + @BeforeAll + public static void beforeAll() throws SQLException, ClassNotFoundException { + DbTestHelper.createFixtures(CONNECTION_URL); + } + + @Test + public void shouldRunBenchJdbcParquetJob() throws Exception { + final Path benchDir = TestHelper.createTmpDirPath("jdbc-parquet-bench-run"); + BenchJdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--skipPartitionCheck", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--table=COFFEES", + "--output=" + benchDir.toString(), + "--avroCodec=snappy", + "--executions=2" + }) + .run(); + assertThat(TestHelper.listDir(benchDir.toFile()), containsInAnyOrder("run_0", "run_1")); + } + + @Test + public void shouldOutputMetricsTsvSummary() throws Exception { + final Path benchDir = TestHelper.createTmpDirPath("jdbc-parquet-bench-metrics"); + final ByteArrayOutputStream capturedOut = new ByteArrayOutputStream(); + final PrintStream originalOut = System.out; + System.setOut(new PrintStream(capturedOut)); + try { + BenchJdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--skipPartitionCheck", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--table=COFFEES", + "--output=" + benchDir.toString(), + "--avroCodec=snappy", + "--executions=2" + }) + .run(); + } finally { + System.setOut(originalOut); + } + + final String output = capturedOut.toString(); + Assertions.assertTrue(output.contains("Summary for BenchJdbcParquetJob")); + Assertions.assertTrue(output.contains("recordCount")); + Assertions.assertTrue(output.contains("writeElapsedMs")); + Assertions.assertTrue(output.contains("bytesWritten")); + Assertions.assertTrue(output.contains("KbWritePerSec")); + Assertions.assertTrue(output.contains("run_00")); + Assertions.assertTrue(output.contains("run_01")); + Assertions.assertTrue(output.contains("max")); + Assertions.assertTrue(output.contains("mean")); + Assertions.assertTrue(output.contains("min")); + Assertions.assertTrue(output.contains("stddev")); + // Each run should report 2 records + assertThat(output.split("run_").length, greaterThan(2)); + } +} diff --git a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/ChannelOutputFileTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/ChannelOutputFileTest.java new file mode 100644 index 00000000..b82d94a9 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/ChannelOutputFileTest.java @@ -0,0 +1,94 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.channels.Channels; +import java.nio.channels.WritableByteChannel; +import org.apache.parquet.io.PositionOutputStream; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class ChannelOutputFileTest { + + @Test + public void shouldTrackPosition() throws IOException { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final WritableByteChannel channel = Channels.newChannel(baos); + final ChannelOutputFile outputFile = new ChannelOutputFile(channel); + + try (PositionOutputStream out = outputFile.create(0)) { + Assertions.assertEquals(0, out.getPos()); + + out.write(42); + Assertions.assertEquals(1, out.getPos()); + + out.write(new byte[] {1, 2, 3, 4, 5}); + Assertions.assertEquals(6, out.getPos()); + + out.write(new byte[] {0, 1, 2, 3, 4, 5, 6, 7}, 2, 3); + Assertions.assertEquals(9, out.getPos()); + } + + Assertions.assertEquals(9, baos.size()); + } + + @Test + public void shouldWriteCorrectBytes() throws IOException { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final WritableByteChannel channel = Channels.newChannel(baos); + final ChannelOutputFile outputFile = new ChannelOutputFile(channel); + + try (PositionOutputStream out = outputFile.create(0)) { + out.write(65); // 'A' + out.write(new byte[] {66, 67}); // 'B', 'C' + } + + byte[] result = baos.toByteArray(); + Assertions.assertEquals(3, result.length); + Assertions.assertEquals(65, result[0]); + Assertions.assertEquals(66, result[1]); + Assertions.assertEquals(67, result[2]); + } + + @Test + public void shouldNotSupportBlockSize() throws IOException { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final WritableByteChannel channel = Channels.newChannel(baos); + final ChannelOutputFile outputFile = new ChannelOutputFile(channel); + + Assertions.assertFalse(outputFile.supportsBlockSize()); + Assertions.assertEquals(0, outputFile.defaultBlockSize()); + } + + @Test + public void shouldReturnSameStreamFromCreateOrOverwrite() throws IOException { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final WritableByteChannel channel = Channels.newChannel(baos); + final ChannelOutputFile outputFile = new ChannelOutputFile(channel); + + try (PositionOutputStream out = outputFile.createOrOverwrite(0)) { + Assertions.assertNotNull(out); + Assertions.assertEquals(0, out.getPos()); + } + } +} diff --git a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java new file mode 100644 index 00000000..2ee9a33a --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java @@ -0,0 +1,178 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import com.spotify.dbeam.args.JdbcConnectionArgs; +import java.util.Collections; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class JdbcParquetArgsTest { + + private static JdbcConnectionArgs connArgs() { + try { + return JdbcConnectionArgs.create("jdbc:h2:mem:test"); + } catch (ClassNotFoundException e) { + throw new RuntimeException(e); + } + } + + private static final JdbcConnectionArgs CONN_ARGS = connArgs(); + + @Test + public void shouldCreateWithDefaults() { + final JdbcParquetArgs args = JdbcParquetArgs.create(CONN_ARGS); + + Assertions.assertEquals(10000, args.fetchSize()); + Assertions.assertEquals("snappy", args.parquetCodec()); + Assertions.assertEquals(128L * 1024 * 1024, args.rowGroupSize()); + Assertions.assertEquals(1024L * 1024, args.pageSize()); + Assertions.assertEquals(Collections.emptyList(), args.preCommand()); + } + + @Test + public void shouldMapSnappyCodec() { + final JdbcParquetArgs args = JdbcParquetArgs.create(CONN_ARGS); + + Assertions.assertEquals(CompressionCodecName.SNAPPY, args.getCompressionCodecName()); + } + + @Test + public void shouldMapGzipCodec() { + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "gzip", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); + + Assertions.assertEquals(CompressionCodecName.GZIP, args.getCompressionCodecName()); + } + + @Test + public void shouldMapZstdCodec() { + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "zstd", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); + + Assertions.assertEquals(CompressionCodecName.ZSTD, args.getCompressionCodecName()); + } + + @Test + public void shouldMapLz4Codec() { + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "lz4", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); + + Assertions.assertEquals(CompressionCodecName.LZ4_RAW, args.getCompressionCodecName()); + } + + @Test + public void shouldMapUncompressedCodec() { + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "uncompressed", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); + + Assertions.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); + } + + @Test + public void shouldMapNoneCodec() { + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "none", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); + + Assertions.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); + } + + @Test + public void shouldRejectInvalidCodec() { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "invalid", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row")); + } + + @Test + public void shouldAcceptCustomFetchSize() { + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 50000, + "snappy", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); + + Assertions.assertEquals(50000, args.fetchSize()); + } + + @Test + public void shouldAcceptCustomRowGroupSize() { + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "snappy", + 128 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); + + Assertions.assertEquals(128L * 1024 * 1024, args.rowGroupSize()); + } +} diff --git a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java new file mode 100644 index 00000000..8405a26f --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java @@ -0,0 +1,361 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.greaterThan; + +import com.spotify.dbeam.DbTestHelper; +import com.spotify.dbeam.TestHelper; +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.sql.SQLException; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +public class JdbcParquetJobTest { + + private static final String CONNECTION_URL = + "jdbc:h2:mem:testparquet;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; + private static Path testDir; + private static Path passwordPath; + private static Path sqlPath; + + @BeforeAll + public static void beforeAll() throws SQLException, ClassNotFoundException, IOException { + testDir = TestHelper.createTmpDirPath("jdbc-parquet-test-"); + passwordPath = testDir.resolve(".password"); + sqlPath = testDir.resolve("query.sql"); + passwordPath.toFile().createNewFile(); + Files.write(sqlPath, "SELECT COF_NAME, SIZE, TOTAL FROM COFFEES WHERE SIZE >= 300".getBytes()); + DbTestHelper.createFixtures(CONNECTION_URL); + } + + @Test + public void shouldRunJdbcParquetJob() throws Exception { + final Path outputPath = testDir.resolve("shouldRunJdbcParquetJob"); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--avroCodec=snappy" + }) + .runExport(); + + assertThat( + TestHelper.listDir(outputPath.toFile()), + containsInAnyOrder( + "_AVRO_SCHEMA.avsc", + "_PARQUET_SCHEMA.json", + "_METRICS.json", + "_SERVICE_METRICS.json", + "_queries", + "part-00000-of-00001.parquet")); + assertThat( + TestHelper.listDir(outputPath.resolve("_queries").toFile()), + containsInAnyOrder("query_0.sql")); + + // Verify parquet schema file was written + final String schemaJson = + new String(Files.readAllBytes(outputPath.resolve("_PARQUET_SCHEMA.json"))); + Assertions.assertTrue(schemaJson.contains("COFFEES")); + + // Verify the parquet file has data by checking file size + final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); + assertThat(parquetFile.length(), greaterThan(0L)); + } + + @Test + public void shouldRunJdbcParquetJobDataOnly() throws Exception { + final Path outputPath = testDir.resolve("shouldRunJdbcParquetJobDataOnly"); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--dataOnly=true", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath.toString(), + "--avroCodec=snappy" + }) + .runExport(); + + assertThat( + TestHelper.listDir(outputPath.toFile()), containsInAnyOrder("part-00000-of-00001.parquet")); + } + + @Test + public void shouldRunJdbcParquetJobSqlFile() throws Exception { + final Path outputPath = testDir.resolve("shouldRunJdbcParquetJobSqlFile"); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--output=" + outputPath, + "--avroCodec=snappy", + "--sqlFile=" + sqlPath.toString() + }) + .runExport(); + + assertThat( + TestHelper.listDir(outputPath.toFile()), + containsInAnyOrder( + "_AVRO_SCHEMA.avsc", + "_PARQUET_SCHEMA.json", + "_METRICS.json", + "_SERVICE_METRICS.json", + "_queries", + "part-00000-of-00001.parquet")); + + // Verify schema only has the 3 selected columns + final String schemaJson = + new String(Files.readAllBytes(outputPath.resolve("_PARQUET_SCHEMA.json"))); + Assertions.assertTrue(schemaJson.contains("COF_NAME")); + Assertions.assertTrue(schemaJson.contains("SIZE")); + Assertions.assertTrue(schemaJson.contains("TOTAL")); + // Should only have a parquet file with data + final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); + assertThat(parquetFile.length(), greaterThan(0L)); + } + + @Test + public void shouldRunParquetJobWithLogicalTypes() throws Exception { + final Path outputPath = testDir.resolve("shouldRunParquetJobWithLogicalTypes"); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--avroCodec=snappy", + "--useAvroLogicalTypes" + }) + .runExport(); + + assertThat( + TestHelper.listDir(outputPath.toFile()), + containsInAnyOrder( + "_AVRO_SCHEMA.avsc", + "_PARQUET_SCHEMA.json", + "_METRICS.json", + "_SERVICE_METRICS.json", + "_queries", + "part-00000-of-00001.parquet")); + + final String schemaJson = + new String(Files.readAllBytes(outputPath.resolve("_PARQUET_SCHEMA.json"))); + Assertions.assertTrue(schemaJson.contains("TIMESTAMP")); + } + + @Test + public void shouldRunParquetJobWithMinRows() throws Exception { + final Path outputPath = testDir.resolve("shouldRunParquetJobWithMinRows"); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--avroCodec=snappy", + "--minRows=2" + }) + .runExport(); + + final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); + assertThat(parquetFile.length(), greaterThan(0L)); + } + + @Test + public void shouldRunParquetJobWithInputSchemaFile() throws Exception { + final Path outputPath = testDir.resolve("shouldRunParquetJobWithInputSchemaFile"); + final Path schemaFile = testDir.resolve("input_schema.parquet.txt"); + Files.write( + schemaFile, + ("message COFFEES {\n" + + " optional binary COF_NAME (STRING);\n" + + " optional double SIZE;\n" + + " optional int64 TOTAL;\n" + + "}") + .getBytes()); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--output=" + outputPath, + "--avroCodec=snappy", + "--parquetSchemaFilePath=" + schemaFile.toString(), + "--sqlFile=" + sqlPath.toString() + }) + .runExport(); + + final String schemaJson = + new String(Files.readAllBytes(outputPath.resolve("_PARQUET_SCHEMA.json"))); + Assertions.assertTrue(schemaJson.contains("COF_NAME")); + Assertions.assertTrue(schemaJson.contains("SIZE")); + Assertions.assertTrue(schemaJson.contains("TOTAL")); + final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); + assertThat(parquetFile.length(), greaterThan(0L)); + } + + @Test + public void shouldMapAvroCodecToParquetCodec() { + Assertions.assertEquals("snappy", JdbcParquetJob.mapAvroCodecToParquetCodec("snappy")); + Assertions.assertEquals("gzip", JdbcParquetJob.mapAvroCodecToParquetCodec("deflate1")); + Assertions.assertEquals("gzip", JdbcParquetJob.mapAvroCodecToParquetCodec("deflate6")); + Assertions.assertEquals("gzip", JdbcParquetJob.mapAvroCodecToParquetCodec("deflate9")); + Assertions.assertEquals("zstd", JdbcParquetJob.mapAvroCodecToParquetCodec("zstandard1")); + Assertions.assertEquals("zstd", JdbcParquetJob.mapAvroCodecToParquetCodec("zstandard9")); + } + + @Test + public void shouldRunParquetJobWithDeflateCodec() throws Exception { + final Path outputPath = testDir.resolve("shouldRunParquetJobWithDeflateCodec"); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--avroCodec=deflate1" + }) + .runExport(); + + final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); + assertThat(parquetFile.length(), greaterThan(0L)); + } + + @Test + public void shouldRunParquetJobWithZstandardCodec() throws Exception { + final Path outputPath = testDir.resolve("shouldRunParquetJobWithZstandardCodec"); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--avroCodec=zstandard1" + }) + .runExport(); + + final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); + assertThat(parquetFile.length(), greaterThan(0L)); + } + + @Test + public void shouldRunParquetJobWithParquetCodecOption() throws Exception { + final Path outputPath = testDir.resolve("shouldRunParquetJobWithParquetCodecOption"); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--parquetCodec=gzip" + }) + .runExport(); + + final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); + assertThat(parquetFile.length(), greaterThan(0L)); + } + + @Test + public void shouldPreferParquetCodecOverAvroCodec() throws Exception { + final Path outputPath = testDir.resolve("shouldPreferParquetCodecOverAvroCodec"); + + JdbcParquetJob.create( + new String[] { + "--targetParallelism=1", + "--partition=2025-02-28", + "--skipPartitionCheck", + "--exportTimeout=PT1M", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--avroCodec=deflate1", + "--parquetCodec=zstd" + }) + .runExport(); + + final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); + assertThat(parquetFile.length(), greaterThan(0L)); + } +} diff --git a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java new file mode 100644 index 00000000..7ee5c69e --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java @@ -0,0 +1,395 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import com.spotify.dbeam.DbTestHelper; +import com.spotify.dbeam.args.QueryBuilderArgs; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Map; +import java.util.Optional; +import org.apache.hadoop.conf.Configuration; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.ParquetFileWriter; +import org.apache.parquet.hadoop.ParquetReader; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.example.GroupReadSupport; +import org.apache.parquet.hadoop.util.HadoopInputFile; +import org.apache.parquet.io.OutputFile; +import org.apache.parquet.schema.MessageType; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +public class JdbcParquetRecordTest { + + private static final String CONNECTION_URL = + "jdbc:h2:mem:testrecord;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; + + @BeforeAll + public static void beforeAll() throws SQLException, ClassNotFoundException { + DbTestHelper.createFixtures(CONNECTION_URL); + } + + @Test + public void shouldCreateSchemaFromDatabase() throws ClassNotFoundException, SQLException { + final Connection connection = DbTestHelper.createConnection(CONNECTION_URL); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.empty(), + false, + "typed_first_row"); + + Assertions.assertNotNull(schema); + Assertions.assertEquals("COFFEES", schema.getName()); + Assertions.assertEquals(14, schema.getFieldCount()); + Assertions.assertEquals("COF_NAME", schema.getFields().get(0).getName()); + Assertions.assertEquals("SUP_ID", schema.getFields().get(1).getName()); + Assertions.assertEquals("PRICE", schema.getFields().get(2).getName()); + Assertions.assertEquals("TEMPERATURE", schema.getFields().get(3).getName()); + Assertions.assertEquals("SIZE", schema.getFields().get(4).getName()); + Assertions.assertEquals("IS_ARABIC", schema.getFields().get(5).getName()); + Assertions.assertEquals("SALES", schema.getFields().get(6).getName()); + Assertions.assertEquals("TOTAL", schema.getFields().get(7).getName()); + Assertions.assertEquals("CREATED", schema.getFields().get(8).getName()); + Assertions.assertEquals("UPDATED", schema.getFields().get(9).getName()); + Assertions.assertEquals("UID", schema.getFields().get(10).getName()); + Assertions.assertEquals("ROWNUM", schema.getFields().get(11).getName()); + Assertions.assertEquals("INT_ARR", schema.getFields().get(12).getName()); + Assertions.assertEquals("TEXT_ARR", schema.getFields().get(13).getName()); + } + + @Test + public void shouldCreateSchemaWithLogicalTypes() throws ClassNotFoundException, SQLException { + final Connection connection = DbTestHelper.createConnection(CONNECTION_URL); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.empty(), + true, + "typed_first_row"); + + Assertions.assertEquals(14, schema.getFieldCount()); + // CREATED and UPDATED should have timestamp logical type + Assertions.assertNotNull( + schema.getFields().get(8).asPrimitiveType().getLogicalTypeAnnotation()); + Assertions.assertNotNull( + schema.getFields().get(9).asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldCreateSchemaWithCustomName() throws ClassNotFoundException, SQLException { + final Connection connection = DbTestHelper.createConnection(CONNECTION_URL); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.of("CustomSchema"), + false, + "typed_first_row"); + + Assertions.assertEquals("CustomSchema", schema.getName()); + } + + @Test + public void shouldWriteAndReadBackParquetRecords() + throws ClassNotFoundException, SQLException, IOException { + final Connection connection = DbTestHelper.createConnection(CONNECTION_URL); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.empty(), + false, + "typed_first_row"); + + // Write records to a temp file + final Path tempFile = Files.createTempFile("parquet-test-", ".parquet"); + Files.delete(tempFile); // ParquetWriter needs to create the file itself + + final ResultSet rs = connection.createStatement().executeQuery("SELECT * FROM COFFEES"); + + final OutputFile outputFile = + new ChannelOutputFile( + java.nio.channels.FileChannel.open( + tempFile, + java.nio.file.StandardOpenOption.CREATE, + java.nio.file.StandardOpenOption.WRITE)); + + try (ParquetWriter writer = + new JdbcParquetIO.ResultSetParquetWriterBuilder(outputFile, schema) + .withWriteMode(ParquetFileWriter.Mode.CREATE) + .build()) { + while (rs.next()) { + writer.write(rs); + } + } + + // Read back and verify + final Configuration conf = new Configuration(); + final org.apache.hadoop.fs.Path hadoopPath = new org.apache.hadoop.fs.Path(tempFile.toUri()); + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { + Group record1 = reader.read(); + Assertions.assertNotNull(record1); + // Verify first record has expected fields + String cofName = record1.getString("COF_NAME", 0); + Assertions.assertNotNull(cofName); + + Group record2 = reader.read(); + Assertions.assertNotNull(record2); + + // Should only be 2 records + Group record3 = reader.read(); + Assertions.assertNull(record3); + } + + // Cleanup + Files.deleteIfExists(tempFile); + } + + @Test + public void shouldWriteCorrectFieldValues() + throws ClassNotFoundException, SQLException, IOException { + final Connection connection = DbTestHelper.createConnection(CONNECTION_URL); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.empty(), + false, + "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-values-test-", ".parquet"); + Files.delete(tempFile); + + final ResultSet rs = + connection.createStatement().executeQuery("SELECT * FROM COFFEES ORDER BY COF_NAME"); + + final OutputFile outputFile = + new ChannelOutputFile( + java.nio.channels.FileChannel.open( + tempFile, + java.nio.file.StandardOpenOption.CREATE, + java.nio.file.StandardOpenOption.WRITE)); + + try (ParquetWriter writer = + new JdbcParquetIO.ResultSetParquetWriterBuilder(outputFile, schema) + .withWriteMode(ParquetFileWriter.Mode.CREATE) + .build()) { + while (rs.next()) { + writer.write(rs); + } + } + + final Configuration conf = new Configuration(); + final org.apache.hadoop.fs.Path hadoopPath = new org.apache.hadoop.fs.Path(tempFile.toUri()); + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { + Group record = reader.read(); + Assertions.assertNotNull(record); + + // Verify typed values for the first record (colombian caffee, sorted) + String cofName = record.getString("COF_NAME", 0); + Assertions.assertEquals("colombian caffee", cofName); + + boolean isArabic = record.getBoolean("IS_ARABIC", 0); + Assertions.assertTrue(isArabic); + + int sales = record.getInteger("SALES", 0); + Assertions.assertEquals(13, sales); + + long total = record.getLong("TOTAL", 0); + Assertions.assertEquals(201L, total); + + float temperature = record.getFloat("TEMPERATURE", 0); + Assertions.assertEquals(87.5f, temperature, 0.01f); + + double size = record.getDouble("SIZE", 0); + Assertions.assertEquals(230.7, size, 0.01); + + long rownum = record.getLong("ROWNUM", 0); + Assertions.assertEquals(2L, rownum); + } + + Files.deleteIfExists(tempFile); + } + + @Test + public void shouldHandleNullValues() throws ClassNotFoundException, SQLException, IOException { + // SUP_ID and UPDATED are null in the Coffee fixtures + final Connection connection = DbTestHelper.createConnection(CONNECTION_URL); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.empty(), + false, + "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-null-test-", ".parquet"); + Files.delete(tempFile); + + final ResultSet rs = connection.createStatement().executeQuery("SELECT * FROM COFFEES LIMIT 1"); + + final OutputFile outputFile = + new ChannelOutputFile( + java.nio.channels.FileChannel.open( + tempFile, + java.nio.file.StandardOpenOption.CREATE, + java.nio.file.StandardOpenOption.WRITE)); + + try (ParquetWriter writer = + new JdbcParquetIO.ResultSetParquetWriterBuilder(outputFile, schema) + .withWriteMode(ParquetFileWriter.Mode.CREATE) + .build()) { + while (rs.next()) { + writer.write(rs); + } + } + + final Configuration conf = new Configuration(); + final org.apache.hadoop.fs.Path hadoopPath = new org.apache.hadoop.fs.Path(tempFile.toUri()); + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { + Group record = reader.read(); + Assertions.assertNotNull(record); + + // SUP_ID is null - should have 0 repetitions in Parquet + Assertions.assertEquals(0, record.getFieldRepetitionCount("SUP_ID")); + + // UPDATED is null + Assertions.assertEquals(0, record.getFieldRepetitionCount("UPDATED")); + + // Non-null fields should have 1 repetition + Assertions.assertEquals(1, record.getFieldRepetitionCount("COF_NAME")); + Assertions.assertEquals(1, record.getFieldRepetitionCount("IS_ARABIC")); + } + + Files.deleteIfExists(tempFile); + } + + @Test + public void shouldWriteAvroSchemaInParquetFooter() + throws ClassNotFoundException, SQLException, IOException { + final Connection connection = DbTestHelper.createConnection(CONNECTION_URL); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.empty(), + false, + "typed_first_row"); + + final String avroSchemaJson = + "{\"type\":\"record\",\"name\":\"COFFEES\"," + + "\"namespace\":\"dbeam_generated\",\"fields\":[]}"; + + final Path tempFile = Files.createTempFile("parquet-footer-test-", ".parquet"); + Files.delete(tempFile); + + final ResultSet rs = connection.createStatement().executeQuery("SELECT * FROM COFFEES LIMIT 1"); + + final OutputFile outputFile = + new ChannelOutputFile( + java.nio.channels.FileChannel.open( + tempFile, + java.nio.file.StandardOpenOption.CREATE, + java.nio.file.StandardOpenOption.WRITE)); + + try (ParquetWriter writer = + new JdbcParquetIO.ResultSetParquetWriterBuilder(outputFile, schema, avroSchemaJson) + .withWriteMode(ParquetFileWriter.Mode.CREATE) + .build()) { + while (rs.next()) { + writer.write(rs); + } + } + + // Read back footer metadata and verify parquet.avro.schema key + final Configuration conf = new Configuration(); + final org.apache.hadoop.fs.Path hadoopPath = new org.apache.hadoop.fs.Path(tempFile.toUri()); + try (ParquetFileReader fileReader = + ParquetFileReader.open(HadoopInputFile.fromPath(hadoopPath, conf))) { + final Map keyValueMetaData = + fileReader.getFooter().getFileMetaData().getKeyValueMetaData(); + Assertions.assertTrue(keyValueMetaData.containsKey(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY)); + final String actualAvroSchema = keyValueMetaData.get(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY); + Assertions.assertEquals(avroSchemaJson, actualAvroSchema); + } + + Files.deleteIfExists(tempFile); + } + + @Test + public void shouldOmitAvroSchemaWhenNotProvided() + throws ClassNotFoundException, SQLException, IOException { + final Connection connection = DbTestHelper.createConnection(CONNECTION_URL); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.empty(), + false, + "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-no-avro-test-", ".parquet"); + Files.delete(tempFile); + + final ResultSet rs = connection.createStatement().executeQuery("SELECT * FROM COFFEES LIMIT 1"); + + final OutputFile outputFile = + new ChannelOutputFile( + java.nio.channels.FileChannel.open( + tempFile, + java.nio.file.StandardOpenOption.CREATE, + java.nio.file.StandardOpenOption.WRITE)); + + // No avroSchemaJson provided (2-arg constructor) + try (ParquetWriter writer = + new JdbcParquetIO.ResultSetParquetWriterBuilder(outputFile, schema) + .withWriteMode(ParquetFileWriter.Mode.CREATE) + .build()) { + while (rs.next()) { + writer.write(rs); + } + } + + final Configuration conf = new Configuration(); + final org.apache.hadoop.fs.Path hadoopPath = new org.apache.hadoop.fs.Path(tempFile.toUri()); + try (ParquetFileReader fileReader = + ParquetFileReader.open(HadoopInputFile.fromPath(hadoopPath, conf))) { + final Map keyValueMetaData = + fileReader.getFooter().getFileMetaData().getKeyValueMetaData(); + Assertions.assertFalse(keyValueMetaData.containsKey(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY)); + } + + Files.deleteIfExists(tempFile); + } +} diff --git a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java new file mode 100644 index 00000000..9979d934 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java @@ -0,0 +1,437 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import static org.mockito.Mockito.when; + +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Types; +import java.util.Optional; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +public class JdbcParquetSchemaTest { + + public static final int COLUMN_NUM = 1; + + @Test + public void shouldGetDatabaseTableNameFromMetaData() throws SQLException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + + Assertions.assertEquals("test_table", JdbcParquetSchema.getDatabaseTableName(meta)); + } + + @Test + public void shouldDefaultTableNameWhenMetaDataHasEmptyTableName() throws SQLException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn(""); + + Assertions.assertEquals("no_table_name", JdbcParquetSchema.getDatabaseTableName(meta)); + } + + @Test + public void shouldDefaultTableNameWhenMetaDataHasNullTableName() throws SQLException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn(null); + + Assertions.assertEquals("no_table_name", JdbcParquetSchema.getDatabaseTableName(meta)); + } + + @Test + public void shouldGetDatabaseTableNameFromFirstNonNullMetaData() throws SQLException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(2); + when(meta.getTableName(1)).thenReturn(""); + when(meta.getTableName(2)).thenReturn("test_table"); + + Assertions.assertEquals("test_table", JdbcParquetSchema.getDatabaseTableName(meta)); + } + + @Test + public void shouldConvertBigIntSqlTypeToInt64() throws SQLException { + final Type fieldType = buildFieldType(Types.BIGINT, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); + } + + @Test + public void shouldConvertIntegerSqlTypeToInt32() throws SQLException { + final Type fieldType = buildFieldType(Types.INTEGER, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT32); + } + + @Test + public void shouldConvertIntegerWithLongColumnClassNameToInt64() throws SQLException { + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.INTEGER, 0, "java.lang.Long", null, false, "typed_first_row"); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); + } + + @Test + public void shouldConvertSmallIntSqlTypeToInt32() throws SQLException { + final Type fieldType = buildFieldType(Types.SMALLINT, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT32); + } + + @Test + public void shouldConvertTinyIntSqlTypeToInt32() throws SQLException { + final Type fieldType = buildFieldType(Types.TINYINT, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT32); + } + + @Test + public void shouldConvertTimestampSqlTypeToInt64() throws SQLException { + final Type fieldType = buildFieldType(Types.TIMESTAMP, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); + Assertions.assertNull(fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertTimestampSqlTypeWithLogicalType() throws SQLException { + final Type fieldType = buildFieldType(Types.TIMESTAMP, true); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); + Assertions.assertEquals( + LogicalTypeAnnotation.timestampType(true, LogicalTypeAnnotation.TimeUnit.MILLIS), + fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertDateSqlTypeToInt64() throws SQLException { + final Type fieldType = buildFieldType(Types.DATE, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); + } + + @Test + public void shouldConvertDateSqlTypeWithLogicalType() throws SQLException { + final Type fieldType = buildFieldType(Types.DATE, true); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); + Assertions.assertNotNull(fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertTimeSqlTypeToInt64() throws SQLException { + final Type fieldType = buildFieldType(Types.TIME, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); + } + + @Test + public void shouldConvertBooleanSqlTypeToBoolean() throws SQLException { + final Type fieldType = buildFieldType(Types.BOOLEAN, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BOOLEAN); + } + + @Test + public void shouldConvertBitSqlTypeWithNoPrecisionToBoolean() throws SQLException { + final Type fieldType = buildFieldType(Types.BIT, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BOOLEAN); + } + + @Test + public void shouldConvertBitSqlTypeWithPrecision2ToBinary() throws SQLException { + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.BIT, 2, "foobar", null, false, "typed_first_row"); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + } + + @Test + public void shouldConvertBinarySqlTypeToBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.BINARY, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + } + + @Test + public void shouldConvertVarbinarySqlTypeToBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.VARBINARY, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + } + + @Test + public void shouldConvertBlobSqlTypeToBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.BLOB, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + } + + @Test + public void shouldConvertDoubleSqlTypeToDouble() throws SQLException { + final Type fieldType = buildFieldType(Types.DOUBLE, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.DOUBLE); + } + + @Test + public void shouldConvertFloatSqlTypeToFloat() throws SQLException { + final Type fieldType = buildFieldType(Types.FLOAT, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.FLOAT); + } + + @Test + public void shouldConvertRealSqlTypeToFloat() throws SQLException { + final Type fieldType = buildFieldType(Types.REAL, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.FLOAT); + } + + @Test + public void shouldConvertVarcharSqlTypeToStringBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.VARCHAR, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assertions.assertEquals( + LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertCharSqlTypeToStringBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.CHAR, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assertions.assertEquals( + LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertClobSqlTypeToStringBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.CLOB, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assertions.assertEquals( + LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertArraySqlTypeToListWithStringElements() throws SQLException { + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.ARRAY, 0, "java.sql.Array", "_text", false, "typed_first_row"); + + Assertions.assertFalse(fieldType.isPrimitive()); + Assertions.assertEquals(LogicalTypeAnnotation.listType(), fieldType.getLogicalTypeAnnotation()); + Assertions.assertEquals(Type.Repetition.OPTIONAL, fieldType.getRepetition()); + } + + @Test + public void shouldConvertIntegerArrayToListOfInt32() { + final Type elementType = JdbcParquetSchema.buildArrayElementType("_int4"); + + assertPrimitiveType(elementType, PrimitiveType.PrimitiveTypeName.INT32); + } + + @Test + public void shouldConvertBigintArrayToListOfInt64() { + final Type elementType = JdbcParquetSchema.buildArrayElementType("_int8"); + + assertPrimitiveType(elementType, PrimitiveType.PrimitiveTypeName.INT64); + } + + @Test + public void shouldConvertFloat4ArrayToListOfFloat() { + final Type elementType = JdbcParquetSchema.buildArrayElementType("_float4"); + + assertPrimitiveType(elementType, PrimitiveType.PrimitiveTypeName.FLOAT); + } + + @Test + public void shouldConvertFloat8ArrayToListOfDouble() { + final Type elementType = JdbcParquetSchema.buildArrayElementType("_float8"); + + assertPrimitiveType(elementType, PrimitiveType.PrimitiveTypeName.DOUBLE); + } + + @Test + public void shouldConvertBoolArrayToListOfBoolean() { + final Type elementType = JdbcParquetSchema.buildArrayElementType("_bool"); + + assertPrimitiveType(elementType, PrimitiveType.PrimitiveTypeName.BOOLEAN); + } + + @Test + public void shouldConvertTextArrayToListOfString() { + final Type elementType = JdbcParquetSchema.buildArrayElementType("_text"); + + assertPrimitiveType(elementType, PrimitiveType.PrimitiveTypeName.BINARY); + Assertions.assertEquals( + LogicalTypeAnnotation.stringType(), + elementType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldResolveH2IntegerArrayType() { + Assertions.assertEquals("int4", JdbcParquetSchema.resolveArrayElementTypeName("INTEGER ARRAY")); + } + + @Test + public void shouldResolveH2VarcharArrayType() { + Assertions.assertEquals("text", JdbcParquetSchema.resolveArrayElementTypeName("VARCHAR ARRAY")); + } + + @Test + public void shouldResolveNullColumnTypeName() { + Assertions.assertEquals("text", JdbcParquetSchema.resolveArrayElementTypeName(null)); + } + + @Test + public void shouldConvertUuidWithLogicalType() throws SQLException { + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.OTHER, 0, "foobar", "uuid", true, "typed_first_row"); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY); + Assertions.assertEquals(16, fieldType.asPrimitiveType().getTypeLength()); + Assertions.assertEquals( + LogicalTypeAnnotation.uuidType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertUuidWithoutLogicalType() throws SQLException { + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.OTHER, 0, "foobar", "uuid", false, "typed_first_row"); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assertions.assertEquals( + LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldDefaultConversionToStringBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.SQLXML, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assertions.assertEquals( + LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldCreateParquetSchemaFromMockResultSet() throws SQLException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(3); + when(meta.getTableName(1)).thenReturn("test_table"); + when(meta.getTableName(2)).thenReturn("test_table"); + when(meta.getTableName(3)).thenReturn("test_table"); + when(meta.getColumnName(1)).thenReturn("id"); + when(meta.getColumnName(2)).thenReturn("name"); + when(meta.getColumnName(3)).thenReturn("active"); + when(meta.getColumnType(1)).thenReturn(Types.BIGINT); + when(meta.getColumnType(2)).thenReturn(Types.VARCHAR); + when(meta.getColumnType(3)).thenReturn(Types.BOOLEAN); + when(meta.getColumnClassName(1)).thenReturn("java.lang.Long"); + when(meta.getColumnClassName(2)).thenReturn("java.lang.String"); + when(meta.getColumnClassName(3)).thenReturn("java.lang.Boolean"); + + final ResultSet resultSet = Mockito.mock(ResultSet.class); + when(resultSet.getMetaData()).thenReturn(meta); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + Assertions.assertEquals("test_table", schema.getName()); + Assertions.assertEquals(3, schema.getFieldCount()); + Assertions.assertEquals("id", schema.getFields().get(0).getName()); + Assertions.assertEquals("name", schema.getFields().get(1).getName()); + Assertions.assertEquals("active", schema.getFields().get(2).getName()); + assertPrimitiveType(schema.getFields().get(0), PrimitiveType.PrimitiveTypeName.INT64); + assertPrimitiveType(schema.getFields().get(1), PrimitiveType.PrimitiveTypeName.BINARY); + assertPrimitiveType(schema.getFields().get(2), PrimitiveType.PrimitiveTypeName.BOOLEAN); + Assertions.assertTrue(schema.getFields().get(0).isRepetition(Type.Repetition.OPTIONAL)); + Assertions.assertTrue(schema.getFields().get(1).isRepetition(Type.Repetition.OPTIONAL)); + Assertions.assertTrue(schema.getFields().get(2).isRepetition(Type.Repetition.OPTIONAL)); + } + + @Test + public void shouldUseCustomSchemaName() throws SQLException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + when(meta.getColumnName(1)).thenReturn("id"); + when(meta.getColumnType(1)).thenReturn(Types.BIGINT); + when(meta.getColumnClassName(1)).thenReturn("java.lang.Long"); + + final ResultSet resultSet = Mockito.mock(ResultSet.class); + when(resultSet.getMetaData()).thenReturn(meta); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.of("CustomName"), false, "typed_first_row"); + + Assertions.assertEquals("CustomName", schema.getName()); + } + + @Test + public void shouldNormalizeSpecialCharactersInColumnNames() throws SQLException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test-table.name"); + when(meta.getColumnName(1)).thenReturn("column-name.with spaces"); + when(meta.getColumnType(1)).thenReturn(Types.INTEGER); + when(meta.getColumnClassName(1)).thenReturn("java.lang.Integer"); + + final ResultSet resultSet = Mockito.mock(ResultSet.class); + when(resultSet.getMetaData()).thenReturn(meta); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + Assertions.assertEquals("test_table_name", schema.getName()); + Assertions.assertEquals("column_name_with_spaces", schema.getFields().get(0).getName()); + } + + private Type buildFieldType(final int sqlType, final boolean useLogicalTypes) { + return JdbcParquetSchema.buildParquetFieldType( + "column1", sqlType, 0, "foobar", null, useLogicalTypes, "typed_first_row"); + } + + private void assertPrimitiveType( + final Type actual, final PrimitiveType.PrimitiveTypeName expected) { + Assertions.assertTrue(actual.isPrimitive()); + Assertions.assertEquals(expected, actual.asPrimitiveType().getPrimitiveTypeName()); + } +} diff --git a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java new file mode 100644 index 00000000..a2d89a40 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java @@ -0,0 +1,437 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import static org.mockito.Mockito.when; + +import com.spotify.dbeam.TestHelper; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.sql.Types; +import java.util.Optional; +import java.util.UUID; +import org.apache.hadoop.conf.Configuration; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.hadoop.ParquetFileWriter; +import org.apache.parquet.hadoop.ParquetReader; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.example.GroupReadSupport; +import org.apache.parquet.io.OutputFile; +import org.apache.parquet.schema.MessageType; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +public class PostgresJdbcParquetTest { + + private ResultSet buildMockResultSet(ResultSetMetaData meta) throws SQLException { + final ResultSet resultSet = Mockito.mock(ResultSet.class); + when(resultSet.getMetaData()).thenReturn(meta); + return resultSet; + } + + @Test + public void shouldEncodeUuidAsString() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.OTHER, "uuid_field", "java.util.UUID", "uuid"); + + final ResultSet resultSet = buildMockResultSet(meta); + final UUID uuidExpected = UUID.fromString("123e4567-e89b-12d3-a456-426655440000"); + when(resultSet.getObject(1)).thenReturn(uuidExpected); + when(resultSet.wasNull()).thenReturn(false); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-uuid-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + String actualUuid = record.getString("uuid_field", 0); + Assertions.assertEquals(uuidExpected.toString(), actualUuid); + }); + } + + @Test + public void shouldEncodeStringAndOtherTypes() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(2); + when(meta.getTableName(1)).thenReturn("test_table"); + when(meta.getTableName(2)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.VARCHAR, "text_field", "java.lang.String", "text"); + TestHelper.mockResultSetMeta( + meta, 2, Types.OTHER, "other_field", "java.util.UUID", "something_else"); + + final ResultSet resultSet = buildMockResultSet(meta); + when(resultSet.getString(1)).thenReturn("some_text_42"); + when(resultSet.getString(2)).thenReturn("some_other_42"); + when(resultSet.wasNull()).thenReturn(false); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-string-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Assertions.assertEquals("some_text_42", record.getString("text_field", 0)); + Assertions.assertEquals("some_other_42", record.getString("other_field", 0)); + }); + } + + @Test + public void shouldEncodeTimestampAsMillis() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + TestHelper.mockResultSetMeta( + meta, 1, Types.TIMESTAMP, "ts_field", "java.sql.Timestamp", "timestamp"); + + final ResultSet resultSet = buildMockResultSet(meta); + final Timestamp ts = new Timestamp(1488300933000L); + when(resultSet.getTimestamp(Mockito.eq(1), Mockito.any())).thenReturn(ts); + when(resultSet.wasNull()).thenReturn(false); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-ts-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + long actualTs = record.getLong("ts_field", 0); + Assertions.assertEquals(1488300933000L, actualTs); + }); + } + + @Test + public void shouldHandleNullFields() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(2); + when(meta.getTableName(1)).thenReturn("test_table"); + when(meta.getTableName(2)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.VARCHAR, "name", "java.lang.String", "varchar"); + TestHelper.mockResultSetMeta(meta, 2, Types.INTEGER, "age", "java.lang.Integer", "int4"); + + final ResultSet resultSet = buildMockResultSet(meta); + when(resultSet.getString(1)).thenReturn("alice"); + when(resultSet.getInt(2)).thenReturn(0); + // First call for name (not null), second for age (null) + when(resultSet.wasNull()).thenReturn(false, true); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-null-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Assertions.assertEquals("alice", record.getString("name", 0)); + Assertions.assertEquals(0, record.getFieldRepetitionCount("age")); + }); + } + + @Test + public void shouldEncodeMultipleNumericTypes() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(4); + when(meta.getTableName(1)).thenReturn("test_table"); + when(meta.getTableName(2)).thenReturn("test_table"); + when(meta.getTableName(3)).thenReturn("test_table"); + when(meta.getTableName(4)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.INTEGER, "int_col", "java.lang.Integer", "int4"); + TestHelper.mockResultSetMeta(meta, 2, Types.BIGINT, "long_col", "java.lang.Long", "int8"); + TestHelper.mockResultSetMeta(meta, 3, Types.FLOAT, "float_col", "java.lang.Float", "float4"); + TestHelper.mockResultSetMeta(meta, 4, Types.DOUBLE, "double_col", "java.lang.Double", "float8"); + + final ResultSet resultSet = buildMockResultSet(meta); + when(resultSet.getInt(1)).thenReturn(42); + when(resultSet.getLong(2)).thenReturn(9999999999L); + when(resultSet.getFloat(3)).thenReturn(3.14f); + when(resultSet.getDouble(4)).thenReturn(2.71828); + when(resultSet.wasNull()).thenReturn(false); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-numeric-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Assertions.assertEquals(42, record.getInteger("int_col", 0)); + Assertions.assertEquals(9999999999L, record.getLong("long_col", 0)); + Assertions.assertEquals(3.14f, record.getFloat("float_col", 0), 0.001f); + Assertions.assertEquals(2.71828, record.getDouble("double_col", 0), 0.00001); + }); + } + + @Test + public void shouldEncodeBooleanType() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.BOOLEAN, "is_active", "java.lang.Boolean", "bool"); + + final ResultSet resultSet = buildMockResultSet(meta); + when(resultSet.getBoolean(1)).thenReturn(true); + when(resultSet.wasNull()).thenReturn(false); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-bool-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Assertions.assertTrue(record.getBoolean("is_active", 0)); + }); + } + + @Test + public void shouldEncodeArrayAsParquetList() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.ARRAY, "tags", "java.sql.Array", "_text"); + + final ResultSet resultSet = buildMockResultSet(meta); + final java.sql.Array mockArray = Mockito.mock(java.sql.Array.class); + when(mockArray.getArray()).thenReturn(new String[] {"rock", "jazz", "blues"}); + when(resultSet.getArray(1)).thenReturn(mockArray); + when(resultSet.wasNull()).thenReturn(false); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + // Verify schema has LIST type + Assertions.assertFalse(schema.getFields().get(0).isPrimitive()); + Assertions.assertEquals( + org.apache.parquet.schema.LogicalTypeAnnotation.listType(), + schema.getFields().get(0).getLogicalTypeAnnotation()); + + final Path tempFile = Files.createTempFile("parquet-array-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Group tagsList = record.getGroup("tags", 0); + Assertions.assertEquals(3, tagsList.getFieldRepetitionCount("list")); + Assertions.assertEquals("rock", tagsList.getGroup("list", 0).getString("element", 0)); + Assertions.assertEquals("jazz", tagsList.getGroup("list", 1).getString("element", 0)); + Assertions.assertEquals("blues", tagsList.getGroup("list", 2).getString("element", 0)); + }); + } + + @Test + public void shouldEncodeIntegerArrayAsTypedList() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.ARRAY, "scores", "java.sql.Array", "_int4"); + + final ResultSet resultSet = buildMockResultSet(meta); + final java.sql.Array mockArray = Mockito.mock(java.sql.Array.class); + when(mockArray.getArray()).thenReturn(new Integer[] {10, 20, 30}); + when(resultSet.getArray(1)).thenReturn(mockArray); + when(resultSet.wasNull()).thenReturn(false); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-int-array-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Group scoresList = record.getGroup("scores", 0); + Assertions.assertEquals(3, scoresList.getFieldRepetitionCount("list")); + Assertions.assertEquals(10, scoresList.getGroup("list", 0).getInteger("element", 0)); + Assertions.assertEquals(20, scoresList.getGroup("list", 1).getInteger("element", 0)); + Assertions.assertEquals(30, scoresList.getGroup("list", 2).getInteger("element", 0)); + }); + } + + @Test + public void shouldEncodeBigintArrayAsTypedList() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.ARRAY, "ids", "java.sql.Array", "_int8"); + + final ResultSet resultSet = buildMockResultSet(meta); + final java.sql.Array mockArray = Mockito.mock(java.sql.Array.class); + when(mockArray.getArray()).thenReturn(new Long[] {100L, 200L, 300L}); + when(resultSet.getArray(1)).thenReturn(mockArray); + when(resultSet.wasNull()).thenReturn(false); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-long-array-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Group idsList = record.getGroup("ids", 0); + Assertions.assertEquals(3, idsList.getFieldRepetitionCount("list")); + Assertions.assertEquals(100L, idsList.getGroup("list", 0).getLong("element", 0)); + Assertions.assertEquals(200L, idsList.getGroup("list", 1).getLong("element", 0)); + Assertions.assertEquals(300L, idsList.getGroup("list", 2).getLong("element", 0)); + }); + } + + @Test + public void shouldHandleNullArrayElements() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.ARRAY, "tags", "java.sql.Array", "_text"); + + final ResultSet resultSet = buildMockResultSet(meta); + final java.sql.Array mockArray = Mockito.mock(java.sql.Array.class); + when(mockArray.getArray()).thenReturn(new String[] {"first", null, "third"}); + when(resultSet.getArray(1)).thenReturn(mockArray); + when(resultSet.wasNull()).thenReturn(false); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-null-array-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Group tagsList = record.getGroup("tags", 0); + Assertions.assertEquals(3, tagsList.getFieldRepetitionCount("list")); + // First element: present + Assertions.assertEquals("first", tagsList.getGroup("list", 0).getString("element", 0)); + // Second element: null (element field has 0 repetitions) + Assertions.assertEquals( + 0, tagsList.getGroup("list", 1).getFieldRepetitionCount("element")); + // Third element: present + Assertions.assertEquals("third", tagsList.getGroup("list", 2).getString("element", 0)); + }); + } + + @Test + public void shouldHandleNullSqlArray() throws SQLException, IOException { + final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); + when(meta.getColumnCount()).thenReturn(1); + when(meta.getTableName(1)).thenReturn("test_table"); + TestHelper.mockResultSetMeta(meta, 1, Types.ARRAY, "tags", "java.sql.Array", "_text"); + + final ResultSet resultSet = buildMockResultSet(meta); + when(resultSet.getArray(1)).thenReturn(null); + when(resultSet.wasNull()).thenReturn(true); + + final MessageType schema = + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); + + final Path tempFile = Files.createTempFile("parquet-null-sql-array-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + // Entire array field is null (optional, 0 repetitions) + Assertions.assertEquals(0, record.getFieldRepetitionCount("tags")); + }); + } + + @FunctionalInterface + interface RecordAssertion { + void assertRecord(Group record) throws IOException; + } + + private void writeAndVerify( + MessageType schema, ResultSet resultSet, Path tempFile, RecordAssertion assertion) + throws IOException { + final OutputFile outputFile = + new ChannelOutputFile( + java.nio.channels.FileChannel.open( + tempFile, + java.nio.file.StandardOpenOption.CREATE, + java.nio.file.StandardOpenOption.WRITE)); + + try (ParquetWriter writer = + new JdbcParquetIO.ResultSetParquetWriterBuilder(outputFile, schema) + .withWriteMode(ParquetFileWriter.Mode.CREATE) + .build()) { + writer.write(resultSet); + } + + final Configuration conf = new Configuration(); + final org.apache.hadoop.fs.Path hadoopPath = new org.apache.hadoop.fs.Path(tempFile.toUri()); + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { + Group record = reader.read(); + Assertions.assertNotNull(record); + assertion.assertRecord(record); + } + + Files.deleteIfExists(tempFile); + } +} diff --git a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PsqlParquetJobTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PsqlParquetJobTest.java new file mode 100644 index 00000000..0abe565f --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PsqlParquetJobTest.java @@ -0,0 +1,75 @@ +/*- + * -\-\- + * DBeam Parquet + * -- + * Copyright (C) 2016 - 2025 Spotify AB + * -- + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * -/-/- + */ + +package com.spotify.dbeam.parquet; + +import com.spotify.dbeam.TestHelper; +import java.io.IOException; +import java.nio.file.Path; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +public class PsqlParquetJobTest { + + private static Path testDir; + private static Path passwordPath; + + @BeforeAll + public static void beforeAll() throws IOException { + testDir = TestHelper.createTmpDirPath("psql-parquet-test-"); + passwordPath = testDir.resolve(".password"); + passwordPath.toFile().createNewFile(); + } + + @Test + public void shouldFailOnNonPostgresConnection() { + final Path outputPath = testDir.resolve("shouldFailOnNonPostgres"); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + PsqlParquetJob.create( + new String[] { + "--connectionUrl=jdbc:h2:mem:test", + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--partition=2025-02-28" + })); + } + + @Test + public void shouldFailOnMissingPartition() { + final Path outputPath = testDir.resolve("shouldFailOnMissingPartition"); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + PsqlParquetJob.create( + new String[] { + "--connectionUrl=jdbc:postgresql://localhost/test", + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--skipPartitionCheck" + })); + } +} diff --git a/docs/parquet-type-conversion.md b/docs/parquet-type-conversion.md new file mode 100644 index 00000000..c0a6d294 --- /dev/null +++ b/docs/parquet-type-conversion.md @@ -0,0 +1,63 @@ +### Type Conversion Details: Java SQL ==> Parquet + +Java SQL types (java.sql.Types.*) are converted to Parquet types according to the following table. +When applicable and `--useAvroLogicalTypes` parameter is set to `true`, Parquet logical types are used. + +All columns are represented as `optional` (nullable) fields in the Parquet schema. + +| **Java SQL type** | **Parquet physical type** | **Parquet logical type** | **Comments** | +|--------------------------|---------------------------|---------------------------------|---------------------------------------| +| BIGINT | INT64 | | | +| INTEGER | INT32 | | INT64 if column class is java.lang.Long (e.g. MySQL unsigned int) | +| SMALLINT | INT32 | | | +| TINYINT | INT32 | | | +| TIMESTAMP | INT64 | TIMESTAMP(MILLIS, isAdjustedToUTC=true) | Logical type only with `--useAvroLogicalTypes` | +| DATE | INT64 | TIMESTAMP(MILLIS, isAdjustedToUTC=true) | Logical type only with `--useAvroLogicalTypes` | +| TIME | INT64 | TIMESTAMP(MILLIS, isAdjustedToUTC=true) | Logical type only with `--useAvroLogicalTypes` | +| TIME_WITH_TIMEZONE | INT64 | TIMESTAMP(MILLIS, isAdjustedToUTC=true) | Logical type only with `--useAvroLogicalTypes` | +| BOOLEAN | BOOLEAN | | | +| BIT | BOOLEAN / BINARY | | BOOLEAN if precision <= 1, BINARY otherwise | +| BINARY | BINARY | | | +| VARBINARY | BINARY | | | +| LONGVARBINARY | BINARY | | | +| BLOB | BINARY | | | +| ARRAY | LIST (group) | LIST | Typed elements (see table below) | +| DOUBLE | DOUBLE | | | +| FLOAT | FLOAT | | | +| REAL | FLOAT | | | +| VARCHAR | BINARY | STRING (UTF8) | | +| CHAR | BINARY | STRING (UTF8) | | +| CLOB | BINARY | STRING (UTF8) | | +| LONGNVARCHAR | BINARY | STRING (UTF8) | | +| LONGVARCHAR | BINARY | STRING (UTF8) | | +| NCHAR | BINARY | STRING (UTF8) | | +| DECIMAL / NUMERIC | BINARY | STRING (UTF8) | Exported as string representation; no fixed-point Parquet DECIMAL type | +| OTHER (uuid) | FIXED_LEN_BYTE_ARRAY(16) | UUID | Only with `--useAvroLogicalTypes`; otherwise STRING | +| OTHER | BINARY | STRING (UTF8) | Default for unrecognized OTHER types | +| all other Java SQL types | BINARY | STRING (UTF8) | | + +#### Array element type mapping + +Array element types are inferred from the column type name. For PostgreSQL, the type name is prefixed with underscore (e.g. `_int4`). For H2 and other databases, the type name is parsed from patterns like `INTEGER ARRAY`. + +| **Column type name** | **Element Parquet type** | +|--------------------------|--------------------------| +| `_int`, `_int4`, `_int2` | INT32 | +| `_int8` | INT64 | +| `_float4` | FLOAT | +| `_float8` | DOUBLE | +| `_bool` | BOOLEAN | +| `_text`, `_varchar`, `_uuid`, others | BINARY (STRING) | +| `INTEGER ARRAY` | INT32 | +| `BIGINT ARRAY` | INT64 | +| `VARCHAR ARRAY`, others | BINARY (STRING) | + +#### Codec mapping + +When using `JdbcParquetJob`, the `--avroCodec` parameter is mapped to Parquet compression codecs: + +| **--avroCodec value** | **Parquet codec** | +|-----------------------|-------------------| +| snappy | SNAPPY | +| deflate1 .. deflate9 | GZIP | +| zstandard1 .. zstandard9 | ZSTD | diff --git a/e2e/e2e.sh b/e2e/e2e.sh index 2f0c587e..f50c0b93 100755 --- a/e2e/e2e.sh +++ b/e2e/e2e.sh @@ -13,7 +13,7 @@ readonly PROJECT_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." >/dev/null && pw # This file contatins psql views with complex types to validate and troubleshoot dbeam -PSQL_DOCKER_IMAGE=postgres:16 +PSQL_DOCKER_IMAGE=postgres:18 PSQL_USER=postgres PSQL_PASSWORD=tempandnotasecret PSQL_DB=dbeam_test @@ -65,7 +65,7 @@ JAVA_OPTS=( pack() { java -version - # create a fat jar + # create fat jars (cd "$PROJECT_PATH"; mvn package -Ppack -DskipTests -Dmaven.test.skip=true -Dmaven.site.skip=true -Dmaven.javadoc.skip=true) } @@ -80,6 +80,17 @@ run_docker_dbeam() { "${JAVA_OPTS[@]}" -cp /dbeam/dbeam-core-shaded.jar com.spotify.dbeam.jobs.BenchJdbcAvroJob "$@" } +run_docker_dbeam_parquet() { + time docker run --interactive --rm \ + --net="$DOCKER_NETWORK" \ + --mount="type=bind,source=$PROJECT_PATH/dbeam-parquet/target,target=/dbeam" \ + --mount="type=bind,source=$SCRIPT_PATH,target=$SCRIPT_PATH" \ + --memory=1G \ + --entrypoint=/usr/bin/java \ + "$JAVA_DOCKER_IMAGE" \ + "${JAVA_OPTS[@]}" -cp /dbeam/dbeam-parquet-shaded.jar com.spotify.dbeam.parquet.BenchJdbcParquetJob "$@" +} + runDBeamDockerCon() { OUTPUT="$SCRIPT_PATH/results/testn/$(date +%FT%H%M%S)/" set -o xtrace @@ -99,6 +110,38 @@ runDBeamDockerCon() { avro-tools tojson --head=5 $OUTPUT_FILE } +runDBeamParquetDockerCon() { + OUTPUT="$SCRIPT_PATH/results/testn/parquet-$(date +%FT%H%M%S)/" + set -o xtrace + time \ + run_docker_dbeam_parquet \ + --skipPartitionCheck \ + --targetParallelism=1 \ + "--connectionUrl=jdbc:postgresql://dbeam-postgres:5432/$PSQL_DB?binaryTransfer=${BINARY_TRANSFER:-false}" \ + "--username=$PSQL_USER" \ + "--password=$PSQL_PASSWORD" \ + "--table=${table:-demo_table}" \ + "--partition=$(date +%F)" \ + "--output=$OUTPUT" \ + "--minRows=${minRows:-1000000}" \ + "$@" 2>&1 | tee -a /tmp/debeam_e2e.log + OUTPUT_FILE=$(ls ${OUTPUT}run_0/*.parquet | head -n 1) + echo "Parquet output: $OUTPUT_FILE ($(wc -c < "$OUTPUT_FILE" | tr -d ' ') bytes)" + parquet-tools head -n 5 "$OUTPUT_FILE" || echo "parquet-tools not available, skipping content validation" + parquet-tools schema "$OUTPUT_FILE" || echo "parquet-tools not available, skipping schema validation" + + # Verify parquet.avro.schema is present in footer metadata + AVRO_SCHEMA_META=$(parquet-tools meta "$OUTPUT_FILE" 2>/dev/null | grep "parquet.avro.schema" || true) + if [[ -n "$AVRO_SCHEMA_META" ]]; then + echo "OK: parquet.avro.schema found in footer metadata" + # Sanity check: should contain "type" and "record" (valid Avro JSON) + echo "$AVRO_SCHEMA_META" | grep -q '"type"' && echo "OK: Avro schema contains type field" || echo "WARN: Avro schema may be malformed" + else + echo "FAIL: parquet.avro.schema NOT found in footer metadata" + exit 1 + fi +} + runSuite() { table=demo_table BINARY_TRANSFER='false' runDBeamDockerCon --executions=3 --avroCodec=deflate1 @@ -108,12 +151,24 @@ runSuite() { BINARY_TRANSFER='false' runDBeamDockerCon --executions=3 --avroCodec=deflate1 --arrayMode=typed_postgres } +runParquetSuite() { + table=demo_table + BINARY_TRANSFER='false' runDBeamParquetDockerCon --executions=3 + BINARY_TRANSFER='false' runDBeamParquetDockerCon --executions=3 --queryParallelism=5 --splitColumn=row_number +} + light() { pack table=demo_table BINARY_TRANSFER='false' runDBeamDockerCon --executions=3 --avroCodec=deflate1 --arrayMode=typed_postgres } +lightParquet() { + pack + table=demo_table + BINARY_TRANSFER='false' runDBeamParquetDockerCon --executions=3 --avroCodec=snappy +} + main() { if [[ $# -gt 0 ]]; then @@ -124,6 +179,7 @@ main() { time startPostgres runSuite + runParquetSuite dockerClean fi } diff --git a/pom.xml b/pom.xml index b38ac1ae..0432d0a5 100644 --- a/pom.xml +++ b/pom.xml @@ -76,6 +76,7 @@ dbeam-core + dbeam-parquet dbeam-bom @@ -101,45 +102,51 @@ - 8 + 17 UTF-8 UTF-8 - - 2.65.0 + + 2.73.0 - 1.9 - 1.11.4 + 1.11.0 + 1.11.5 3.42.0 1.17.1 1.26.2 - 2.10.0 - 33.1.0-jre + 2.31.0 + 33.5.0-jre 2.1 4.5.13 4.4.14 - 2.15.4 - 2.10.14 - 4.1.121.Final - 1.7.30 + 2.18.2 + 2.14.0 + 4.1.130.Final + 2.0.16 1.6.8 - 1.5.6-3 + 1.5.7-7 + 1.14.4 + 3.5.0 + + 7.0.0 + 4.2.2 - - 26.57.0 + + 26.79.0 0.31.1 - 1.78.1 - 4.13.2 + 1.84 + 5.11.4 8.4.0 - 3.5.3 - 42.7.4 - 1.18.0 + 3.5.8 + 42.7.8 + 1.25.0 @@ -241,6 +248,32 @@ + + org.apache.parquet + parquet-hadoop + ${parquet.version} + + + org.apache.hadoop + * + + + + + org.apache.hadoop + hadoop-common + ${hadoop.version} + + + com.fasterxml.woodstox + woodstox-core + ${woodstox-core.version} + + + org.codehaus.woodstox + stax2-api + ${stax2-api.version} + org.apache.beam beam-sdks-java-bom @@ -265,7 +298,7 @@ com.google.apis google-api-services-cloudkms - v1-rev20240314-2.0.0 + v1-rev20260319-2.0.0 @@ -400,10 +433,11 @@ - junit - junit - ${junit.version} - test + org.junit + junit-bom + ${junit-jupiter.version} + pom + import org.hamcrest @@ -420,7 +454,13 @@ org.mockito mockito-core - 5.17.0 + 5.20.0 + test + + + org.mockito + mockito-junit-jupiter + 5.20.0 test @@ -431,12 +471,6 @@ - - - com.google.auto.value - auto-value - - @@ -458,15 +492,18 @@ org.apache.maven.plugins maven-compiler-plugin - 3.13.0 - - - - compile - - compile - - + 3.14.1 + + + + + com.google.auto.value + auto-value + ${auto-value.version} + + + maven-release-plugin @@ -519,11 +556,11 @@ org.apache.maven.plugins maven-project-info-reports-plugin - 3.8.0 + 3.9.0 maven-surefire-plugin - 3.5.3 + 3.5.5 true @@ -532,7 +569,7 @@ org.apache.maven.plugins maven-enforcer-plugin - 3.5.0 + 3.6.2 enforce @@ -542,7 +579,7 @@ - [11,) + [17,) @@ -555,6 +592,12 @@ com.fasterxml.jackson.core:jackson-core com.fasterxml.jackson.core:jackson-databind com.fasterxml.jackson.datatype:jackson-datatype-jsr310 + com.fasterxml.jackson.datatype:jackson-datatype-jdk8 + + org.apache.commons:commons-lang3 + commons-logging:commons-logging + org.apache.commons:commons-text + com.google.re2j:re2j com.github.luben:zstd-jni com.google.auto.value:auto-value com.google.auto.value:auto-value-annotations @@ -633,7 +676,7 @@ com.puppycrawl.tools checkstyle - 10.23.0 + 10.26.1 @@ -653,7 +696,7 @@ org.apache.maven.plugins maven-gpg-plugin - 3.2.7 + 3.2.8 sign-artifacts @@ -675,7 +718,7 @@ org.apache.maven.plugins maven-javadoc-plugin - 3.11.2 + 3.12.0 8 @@ -732,7 +775,7 @@ org.apache.maven.plugins maven-shade-plugin - 3.6.0 + 3.6.2 bundle-and-repackage @@ -760,7 +803,7 @@ - com.spotify.dbeam.jobs.JdbcAvroJob + ${dbeam.mainClass} ${basedir}/target/${project.artifactId}-shaded.jar