From 5276648b4301edb7588edc0311acd47a44c5923f Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 3 Oct 2025 03:01:40 +0000 Subject: [PATCH 01/80] Bump org.apache.maven.plugins:maven-enforcer-plugin from 3.5.0 to 3.6.2 Bumps [org.apache.maven.plugins:maven-enforcer-plugin](https://github.com/apache/maven-enforcer) from 3.5.0 to 3.6.2. - [Release notes](https://github.com/apache/maven-enforcer/releases) - [Commits](https://github.com/apache/maven-enforcer/compare/enforcer-3.5.0...enforcer-3.6.2) --- updated-dependencies: - dependency-name: org.apache.maven.plugins:maven-enforcer-plugin dependency-version: 3.6.2 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index b38ac1ae..f1574da0 100644 --- a/pom.xml +++ b/pom.xml @@ -532,7 +532,7 @@ org.apache.maven.plugins maven-enforcer-plugin - 3.5.0 + 3.6.2 enforce From a2d1c9a660d8c9fdbdeb65f032e20aedb551a456 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 22 Sep 2025 03:03:22 +0000 Subject: [PATCH 02/80] Bump org.mockito:mockito-core from 5.17.0 to 5.20.0 Bumps [org.mockito:mockito-core](https://github.com/mockito/mockito) from 5.17.0 to 5.20.0. - [Release notes](https://github.com/mockito/mockito/releases) - [Commits](https://github.com/mockito/mockito/compare/v5.17.0...v5.20.0) --- updated-dependencies: - dependency-name: org.mockito:mockito-core dependency-version: 5.20.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index f1574da0..68f9d0b4 100644 --- a/pom.xml +++ b/pom.xml @@ -420,7 +420,7 @@ org.mockito mockito-core - 5.17.0 + 5.20.0 test From 8ef6e503861d4dc6acdc01889e0504d92fd3c7b8 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 22 Sep 2025 03:01:24 +0000 Subject: [PATCH 03/80] Bump org.apache.maven.plugins:maven-compiler-plugin Bumps [org.apache.maven.plugins:maven-compiler-plugin](https://github.com/apache/maven-compiler-plugin) from 3.13.0 to 3.14.1. - [Release notes](https://github.com/apache/maven-compiler-plugin/releases) - [Commits](https://github.com/apache/maven-compiler-plugin/compare/maven-compiler-plugin-3.13.0...maven-compiler-plugin-3.14.1) --- updated-dependencies: - dependency-name: org.apache.maven.plugins:maven-compiler-plugin dependency-version: 3.14.1 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 68f9d0b4..6a9ad3dc 100644 --- a/pom.xml +++ b/pom.xml @@ -458,7 +458,7 @@ org.apache.maven.plugins maven-compiler-plugin - 3.13.0 + 3.14.1 From 592932915e625440263609039b794315201bda5c Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 22 Sep 2025 03:01:13 +0000 Subject: [PATCH 04/80] Bump org.apache.maven.plugins:maven-javadoc-plugin from 3.11.2 to 3.12.0 Bumps [org.apache.maven.plugins:maven-javadoc-plugin](https://github.com/apache/maven-javadoc-plugin) from 3.11.2 to 3.12.0. - [Release notes](https://github.com/apache/maven-javadoc-plugin/releases) - [Commits](https://github.com/apache/maven-javadoc-plugin/compare/maven-javadoc-plugin-3.11.2...maven-javadoc-plugin-3.12.0) --- updated-dependencies: - dependency-name: org.apache.maven.plugins:maven-javadoc-plugin dependency-version: 3.12.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 6a9ad3dc..229a8c5f 100644 --- a/pom.xml +++ b/pom.xml @@ -675,7 +675,7 @@ org.apache.maven.plugins maven-javadoc-plugin - 3.11.2 + 3.12.0 8 From c7f53ca9ded6059794c69d8facdbb9fc20cf1bd9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 3 Jul 2025 03:13:00 +0000 Subject: [PATCH 05/80] Bump org.apache.maven.plugins:maven-gpg-plugin from 3.2.7 to 3.2.8 Bumps [org.apache.maven.plugins:maven-gpg-plugin](https://github.com/apache/maven-gpg-plugin) from 3.2.7 to 3.2.8. - [Release notes](https://github.com/apache/maven-gpg-plugin/releases) - [Commits](https://github.com/apache/maven-gpg-plugin/compare/maven-gpg-plugin-3.2.7...maven-gpg-plugin-3.2.8) --- updated-dependencies: - dependency-name: org.apache.maven.plugins:maven-gpg-plugin dependency-version: 3.2.8 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 229a8c5f..84c89095 100644 --- a/pom.xml +++ b/pom.xml @@ -653,7 +653,7 @@ org.apache.maven.plugins maven-gpg-plugin - 3.2.7 + 3.2.8 sign-artifacts From 489a0a5591b0a8ecaeee9d49f9f2f6dd01b03f98 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 5 Jun 2025 03:08:18 +0000 Subject: [PATCH 06/80] Bump org.codehaus.mojo:build-helper-maven-plugin from 3.6.0 to 3.6.1 Bumps [org.codehaus.mojo:build-helper-maven-plugin](https://github.com/mojohaus/build-helper-maven-plugin) from 3.6.0 to 3.6.1. - [Release notes](https://github.com/mojohaus/build-helper-maven-plugin/releases) - [Commits](https://github.com/mojohaus/build-helper-maven-plugin/compare/3.6.0...3.6.1) --- updated-dependencies: - dependency-name: org.codehaus.mojo:build-helper-maven-plugin dependency-version: 3.6.1 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] --- dbeam-bom/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dbeam-bom/pom.xml b/dbeam-bom/pom.xml index 4e7d364e..5c6d6e56 100644 --- a/dbeam-bom/pom.xml +++ b/dbeam-bom/pom.xml @@ -47,7 +47,7 @@ org.codehaus.mojo build-helper-maven-plugin - 3.6.0 + 3.6.1 attach-effective-pom From ea7e26ac51f9a72cb19c77d1522ec4f24270c674 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 27 Feb 2025 04:54:19 +0000 Subject: [PATCH 07/80] Bump org.apache.maven.plugins:maven-project-info-reports-plugin Bumps [org.apache.maven.plugins:maven-project-info-reports-plugin](https://github.com/apache/maven-project-info-reports-plugin) from 3.8.0 to 3.9.0. - [Release notes](https://github.com/apache/maven-project-info-reports-plugin/releases) - [Commits](https://github.com/apache/maven-project-info-reports-plugin/compare/maven-project-info-reports-plugin-3.8.0...maven-project-info-reports-plugin-3.9.0) --- updated-dependencies: - dependency-name: org.apache.maven.plugins:maven-project-info-reports-plugin dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 84c89095..715a8374 100644 --- a/pom.xml +++ b/pom.xml @@ -519,7 +519,7 @@ org.apache.maven.plugins maven-project-info-reports-plugin - 3.8.0 + 3.9.0 maven-surefire-plugin From 8e4a759ecef18cbca59943f490a3caad32588ee0 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 15:19:59 +0200 Subject: [PATCH 08/80] Bump PostgreSQL JDBC driver from 42.7.4 to 42.7.7 Fixes CVE-2025-49146 (https://nvd.nist.gov/vuln/detail/CVE-2025-49146) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 715a8374..4b690bc6 100644 --- a/pom.xml +++ b/pom.xml @@ -138,7 +138,7 @@ 4.13.2 8.4.0 3.5.3 - 42.7.4 + 42.7.7 1.18.0 From 50dd3dc6655eeb361e5784294207de9eebdf13f2 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 16:16:18 +0200 Subject: [PATCH 09/80] Bump Avro from 1.11.4 to 1.11.5 --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 4b690bc6..26d49486 100644 --- a/pom.xml +++ b/pom.xml @@ -110,7 +110,7 @@ 2.65.0 1.9 - 1.11.4 + 1.11.5 3.42.0 1.17.1 1.26.2 From 94307a9fe954c5501fcc286386ad34efb993680f Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 16:15:20 +0200 Subject: [PATCH 10/80] Bump Apache Beam SDK from 2.65.0 to 2.72.0 Beam 2.65.0 reaches end-of-support in May 2026. Beam 2.72.0 is the latest stable release, supported until March 2027. Updated dependency versions to match Beam 2.72.0: - errorprone: 2.10.0 -> 2.31.0 - joda-time: 2.10.14 -> 2.14.0 - netty: 4.1.121.Final -> 4.1.124.Final - slf4j: 1.7.30 -> 2.0.16 - google-cloud-libraries-bom: 26.57.0 -> 26.76.0 Added dependency management overrides to resolve version convergence between Beam BOM and libraries-bom: - jackson-dataformat-xml 2.18.2 (from google-cloud-storage) - google-cloud-bigtable 2.73.1 and proto stubs (from beam-io-gcp) - j2objc-annotations 3.1 (from libraries-bom) - failureaccess 1.0.3 (from transitives) --- pom.xml | 55 +++++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 47 insertions(+), 8 deletions(-) diff --git a/pom.xml b/pom.xml index 26d49486..1914b853 100644 --- a/pom.xml +++ b/pom.xml @@ -106,30 +106,30 @@ UTF-8 - - 2.65.0 + + 2.72.0 1.9 1.11.5 3.42.0 1.17.1 1.26.2 - 2.10.0 + 2.31.0 33.1.0-jre 2.1 4.5.13 4.4.14 2.15.4 - 2.10.14 - 4.1.121.Final - 1.7.30 + 2.14.0 + 4.1.124.Final + 2.0.16 1.6.8 1.5.6-3 - - 26.57.0 + + 26.76.0 0.31.1 @@ -238,6 +238,45 @@ opencensus-contrib-grpc-metrics ${opencensus.version} + + + com.fasterxml.jackson.dataformat + jackson-dataformat-xml + 2.18.2 + + + + com.google.cloud + google-cloud-bigtable + 2.73.1 + + + com.google.api.grpc + grpc-google-cloud-bigtable-v2 + 2.73.1 + + + com.google.api.grpc + proto-google-cloud-bigtable-v2 + 2.73.1 + + + com.google.api.grpc + proto-google-cloud-bigtable-admin-v2 + 2.73.1 + + + + com.google.j2objc + j2objc-annotations + 3.1 + + + + com.google.guava + failureaccess + 1.0.3 + From 8f94e325e93bd186ea74a3a43c9744abd5011525 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 16:17:56 +0200 Subject: [PATCH 11/80] Bump Jackson from 2.15.4 to 2.18.2 Aligns with the version provided by google-cloud-libraries-bom 26.76.0, removing the separate jackson-dataformat-xml override. --- pom.xml | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/pom.xml b/pom.xml index 1914b853..b01abde2 100644 --- a/pom.xml +++ b/pom.xml @@ -119,7 +119,7 @@ 2.1 4.5.13 4.4.14 - 2.15.4 + 2.18.2 2.14.0 4.1.124.Final 2.0.16 @@ -238,12 +238,6 @@ opencensus-contrib-grpc-metrics ${opencensus.version} - - - com.fasterxml.jackson.dataformat - jackson-dataformat-xml - 2.18.2 - com.google.cloud From 56ed88fe5bd267e3b6b3234fa098062d450d6988 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 16:23:02 +0200 Subject: [PATCH 12/80] Bump Guava from 33.1.0-jre to 33.5.0-jre Aligns with the version provided by google-cloud-libraries-bom 26.76.0. --- pom.xml | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/pom.xml b/pom.xml index b01abde2..2e7aa713 100644 --- a/pom.xml +++ b/pom.xml @@ -115,7 +115,7 @@ 1.17.1 1.26.2 2.31.0 - 33.1.0-jre + 33.5.0-jre 2.1 4.5.13 4.4.14 @@ -259,18 +259,6 @@ proto-google-cloud-bigtable-admin-v2 2.73.1 - - - com.google.j2objc - j2objc-annotations - 3.1 - - - - com.google.guava - failureaccess - 1.0.3 - From 1526fa0b11e51b20a9881c924ab8b81a564f142b Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 16:32:35 +0200 Subject: [PATCH 13/80] Bump PostgreSQL JDBC driver from 42.7.7 to 42.7.8 --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 2e7aa713..b16b0305 100644 --- a/pom.xml +++ b/pom.xml @@ -138,7 +138,7 @@ 4.13.2 8.4.0 3.5.3 - 42.7.7 + 42.7.8 1.18.0 From 46e52481dc0ac5ccede461da8559db78408948d6 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 16:35:07 +0200 Subject: [PATCH 14/80] Bump google-api-services-cloudkms from v1-rev20240314 to v1-rev20260319 --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index b16b0305..0dd025c3 100644 --- a/pom.xml +++ b/pom.xml @@ -286,7 +286,7 @@ com.google.apis google-api-services-cloudkms - v1-rev20240314-2.0.0 + v1-rev20260319-2.0.0 From d1111806b3ae54bc9b8f28849dffb1ff7716ffbd Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 16:36:56 +0200 Subject: [PATCH 15/80] Bump MariaDB JDBC driver from 3.5.3 to 3.5.8 --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 0dd025c3..9ac8fbb4 100644 --- a/pom.xml +++ b/pom.xml @@ -137,7 +137,7 @@ 1.78.1 4.13.2 8.4.0 - 3.5.3 + 3.5.8 42.7.8 1.18.0 From 998d2c9b35a941a195b79f3150df591f9a24220e Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 16:39:55 +0200 Subject: [PATCH 16/80] Bump zstd-jni from 1.5.6-3 to 1.5.7-7 --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 9ac8fbb4..9f888229 100644 --- a/pom.xml +++ b/pom.xml @@ -124,7 +124,7 @@ 4.1.124.Final 2.0.16 1.6.8 - 1.5.6-3 + 1.5.7-7 From ad07d124ceff0867593b70517e0c2cda1b182d1e Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 16:43:26 +0200 Subject: [PATCH 17/80] Bump Cloud SQL socket factory from 1.18.0 to 1.25.0 1.25.0 is the highest version compatible with the current libraries-bom 26.76.0 without introducing dependency conflicts. Versions >= 1.25.1 require a newer google-api-client than what the BOM provides. --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 9f888229..388402f2 100644 --- a/pom.xml +++ b/pom.xml @@ -139,7 +139,7 @@ 8.4.0 3.5.8 42.7.8 - 1.18.0 + 1.25.0 From 900fc1a91e6d702f4781487a66ed5c7c8e935ea5 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 17:01:54 +0200 Subject: [PATCH 18/80] Update avro-tools download URL to 1.11.5 in CI workflow The previous URL for avro-tools 1.11.3 is no longer available on dlcdn.apache.org, causing CI failures. --- .github/workflows/maven.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml index c20bf249..6799ab3a 100644 --- a/.github/workflows/maven.yml +++ b/.github/workflows/maven.yml @@ -52,11 +52,11 @@ jobs: - name: Install avro tools if: matrix.java_version == '17' 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 - name: End to end tests shell: bash if: matrix.java_version == '17' From 067cc76a1280bbc97f17099db10f7e9e3bbda99d Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Mon, 4 May 2026 12:00:51 +0200 Subject: [PATCH 19/80] Bump Apache Beam SDK from 2.72.0 to 2.73.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Also bump google-cloud libraries-bom from 26.76.0 to 26.79.0 (the version that Beam 2.73 expects) and remove the now-stale google-cloud-bigtable 2.73.1 overrides — libraries-bom 26.79.0 plus the Beam BOM converge cleanly without them. Co-Authored-By: Claude --- pom.xml | 29 ++++------------------------- 1 file changed, 4 insertions(+), 25 deletions(-) diff --git a/pom.xml b/pom.xml index 388402f2..3ca18054 100644 --- a/pom.xml +++ b/pom.xml @@ -106,8 +106,8 @@ UTF-8 - - 2.72.0 + + 2.73.0 1.9 1.11.5 @@ -128,8 +128,8 @@ - - 26.76.0 + + 26.79.0 0.31.1 @@ -238,27 +238,6 @@ opencensus-contrib-grpc-metrics ${opencensus.version} - - - com.google.cloud - google-cloud-bigtable - 2.73.1 - - - com.google.api.grpc - grpc-google-cloud-bigtable-v2 - 2.73.1 - - - com.google.api.grpc - proto-google-cloud-bigtable-v2 - 2.73.1 - - - com.google.api.grpc - proto-google-cloud-bigtable-admin-v2 - 2.73.1 - From e02a1e2811da62c2fc1f6365c157d988640959bb Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Mon, 4 May 2026 12:08:23 +0200 Subject: [PATCH 20/80] Bump BouncyCastle from 1.78.1 to 1.84 Routine update for the BouncyCastle crypto provider; pulls in fixes across 1.79, 1.80, 1.81, 1.82, 1.83, and 1.84 releases (April 2026). Co-Authored-By: Claude --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 3ca18054..76b58cd1 100644 --- a/pom.xml +++ b/pom.xml @@ -134,7 +134,7 @@ 0.31.1 - 1.78.1 + 1.84 4.13.2 8.4.0 3.5.8 From e5896776e748a24b2f3bbba110495aff134e9d7c Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Mon, 4 May 2026 12:08:44 +0200 Subject: [PATCH 21/80] Bump Checkstyle from 10.23.0 to 10.26.1 Update to the latest 10.x release (last line that supports JDK 8/11 toolchains; 11.x requires JDK 17 to run the plugin). Co-Authored-By: Claude --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 76b58cd1..ee30c7c6 100644 --- a/pom.xml +++ b/pom.xml @@ -633,7 +633,7 @@ com.puppycrawl.tools checkstyle - 10.23.0 + 10.26.1 From 09af448c5b838da9d6d3e4379b497770e67e7b1d Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Mon, 4 May 2026 12:09:11 +0200 Subject: [PATCH 22/80] Bump maven-surefire-plugin from 3.5.3 to 3.5.5 Co-Authored-By: Claude --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index ee30c7c6..30fbdcfa 100644 --- a/pom.xml +++ b/pom.xml @@ -523,7 +523,7 @@ maven-surefire-plugin - 3.5.3 + 3.5.5 true From ac4625b84c72cfac89b9e4486db9249be39955fe Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Mon, 4 May 2026 12:09:34 +0200 Subject: [PATCH 23/80] Bump maven-shade-plugin from 3.6.0 to 3.6.2 Co-Authored-By: Claude --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 30fbdcfa..109f489c 100644 --- a/pom.xml +++ b/pom.xml @@ -732,7 +732,7 @@ org.apache.maven.plugins maven-shade-plugin - 3.6.0 + 3.6.2 bundle-and-repackage From 9f5ff69bd22798b7d62086490a45007e2e1c72af Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Mon, 4 May 2026 12:15:03 +0200 Subject: [PATCH 24/80] Use Java 21 on deploy --- .github/workflows/maven.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml index 6799ab3a..9df04ab8 100644 --- a/.github/workflows/maven.yml +++ b/.github/workflows/maven.yml @@ -86,12 +86,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: From 774604ef6ac623897e26d45e688c5faacd292e7d Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Mon, 4 May 2026 12:16:03 +0200 Subject: [PATCH 25/80] Test with Java 25 --- .github/workflows/maven.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml index 9df04ab8..73ce5b54 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: ['11', '17', '21', '25'] os: [ubuntu-latest] steps: - uses: actions/checkout@v4 From e1616da88d378a08a7e9d79138587a5d9e2bf6ba Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Mon, 4 May 2026 12:16:35 +0200 Subject: [PATCH 26/80] Java 21 for e2e tests --- .github/workflows/maven.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml index 73ce5b54..55f2d65a 100644 --- a/.github/workflows/maven.yml +++ b/.github/workflows/maven.yml @@ -50,7 +50,7 @@ jobs: fail_ci_if_error: false # optional (default = false) verbose: true # optional (default = false) - name: Install avro tools - if: matrix.java_version == '17' + if: matrix.java_version == '21' run: | wget https://dlcdn.apache.org/avro/avro-1.11.5/java/avro-tools-1.11.5.jar mkdir -p /opt/avro @@ -59,12 +59,12 @@ jobs: echo "alias avro-tools='java -jar /opt/avro/avro-tools-1.11.5.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 From e2858f1c2bc6d8f16ad9a459040435773b21b884 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Mon, 4 May 2026 13:06:35 +0200 Subject: [PATCH 27/80] Support Java 25 build JDK 23 (JEP 472) disabled implicit annotation processing, so AutoValue silently stopped generating AutoValue_* classes on JDK 25. Declare the processor explicitly via maven-compiler-plugin's annotationProcessorPaths so javac runs it again, and bump auto-value 1.9 -> 1.11.0 for JDK 23+ compiler-internals support. Bump maven.compiler.release 8 -> 11 to drop the obsolete-source-8 warnings; matches the enforcer's existing [11,) floor. Co-Authored-By: Claude --- pom.xml | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/pom.xml b/pom.xml index 109f489c..987ed8ce 100644 --- a/pom.xml +++ b/pom.xml @@ -101,7 +101,7 @@ - 8 + 11 UTF-8 UTF-8 @@ -109,7 +109,7 @@ 2.73.0 - 1.9 + 1.11.0 1.11.5 3.42.0 1.17.1 @@ -459,6 +459,17 @@ org.apache.maven.plugins maven-compiler-plugin 3.14.1 + + + + + com.google.auto.value + auto-value + ${auto-value.version} + + + From 0734694300dfad6c48bc30ba41e02e078df7e297 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 5 May 2026 15:16:47 +0200 Subject: [PATCH 28/80] Drop Java 11 support, raise floor to JDK 17 dbeam-parquet pulls hadoop-common 3.5.0, which is compiled to class file v61 (JDK 17). hadoop-common 3.4.x supports JDK 11 but its UGI calls Subject.getSubject which is gone on JDK 25 (JEP 486), so we can't have both Hadoop 3.5 (the JDK 25 fix) and JDK 11. JDK 11 has been past Oracle's free premier support since Sept 2023. - maven.compiler.release 11 -> 17 - enforcer requireJavaVersion [11,) -> [17,) - CI matrix drops '11' - release.yml builds on JDK 17 instead of 11 Co-Authored-By: Claude --- .github/workflows/maven.yml | 2 +- .github/workflows/release.yml | 8 ++++---- pom.xml | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml index 55f2d65a..49621c8d 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', '25'] + java_version: ['17', '21', '25'] os: [ubuntu-latest] steps: - uses: actions/checkout@v4 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/pom.xml b/pom.xml index 987ed8ce..5370cb54 100644 --- a/pom.xml +++ b/pom.xml @@ -101,7 +101,7 @@ - 11 + 17 UTF-8 UTF-8 @@ -553,7 +553,7 @@ - [11,) + [17,) From 53025b801a0994d997bb85b4438bdc8e779ae90a Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 09:40:45 +0200 Subject: [PATCH 29/80] Add dbeam-parquet module for JDBC to Parquet export Add a new dbeam-parquet Maven module that enables exporting SQL databases directly to Parquet files, complementing the existing Avro export in dbeam-core. Key classes: - JdbcParquetJob: main entry point, reuses shared infrastructure from dbeam-core (connection management, query building, Beam helpers) - JdbcParquetIO: Beam FileBasedSink hierarchy for Parquet output - JdbcParquetSchema: JDBC ResultSetMetaData to Parquet MessageType schema conversion with full SQL type mapping - JdbcParquetWriteSupport: writes ResultSet rows directly to Parquet RecordConsumer without intermediate objects - ChannelOutputFile: OutputFile backed by WritableByteChannel, bypassing Hadoop for file I/O - JdbcParquetArgs: configuration for codec, row group size, page size - JdbcParquetMetering: Beam metrics (counters/gauges) Uses parquet-java (parquet-hadoop) with Hadoop dependencies excluded from parquet-hadoop itself. hadoop-common is included for the Configuration class required by parquet-java's API surface. Supported compression codecs: snappy, gzip, zstd, lz4, uncompressed. Also updates: - Parent POM: adds module, parquet/hadoop version properties, dependency management, enforcer excludes - dbeam-core: adds test-jar packaging for test helper sharing - dbeam-bom: includes dbeam-parquet artifact Declare auto-value-annotations explicitly in dbeam-parquet The module uses @AutoValue but only sees the annotation transitively through the parent's bare auto-value dependency. Make the dep explicit (scope=provided, matching how dbeam-core declares it) so the module compiles on its own merits and stops relying on the parent's processor-jar dependency. Co-Authored-By: Claude Pin woodstox-core and stax2-api to satisfy upper-bound deps dbeam-parquet pulls hadoop-common 3.4.1 (woodstox 5.4.0, stax2-api 4.2/4.2.1) alongside google-cloud-storage (woodstox 7.0.0, stax2-api 4.2.2 via beam transitives), failing the enforcer's requireUpperBoundDeps. Pin both to the higher versions in dependencyManagement so the upper bound is actually met. Co-Authored-By: Claude Bump hadoop 3.4.1->3.5.0 (and netty) for JDK 25 support Hadoop 3.4.1's UserGroupInformation calls Subject.getSubject( AccessControlContext), which throws UnsupportedOperationException on JDK 25 (JEP 486 fully removed the Security Manager). The -Djava.security.manager=allow workaround that worked on JDK 24 is rejected at startup on JDK 25. Hadoop 3.5.0 reworks UGI to use the new Subject.current/callAs APIs and is JDK 25 compatible. 3.5.0 brings netty-transport-native-epoll 4.1.130 transitively, so also bump netty.version to satisfy requireUpperBoundDeps. Co-Authored-By: Claude Drop redundant auto-value dependency in parent annotationProcessorPaths in maven-compiler-plugin already pulls the processor jar onto the processor path; the parent's bare auto-value dependency just put the same jar on every module's compile classpath. Modules that need the @AutoValue annotation now declare auto-value-annotations explicitly (dbeam-core and dbeam-parquet already do). Co-Authored-By: Claude Drop redundant compile execution from parent compiler plugin The parent declared an with id "default" bound to the compile phase, alongside Maven's built-in "default-compile" execution. Both ran on every child module, so the AutoValue processor generated each AutoValue_* file twice and the second run logged "A file for type ... already exists on the sourcepath or classpath". Removing the custom execution leaves only the built-in one. Co-Authored-By: Claude Remove Eclipse IDE files from dbeam-parquet --- dbeam-bom/pom.xml | 5 + dbeam-core/pom.xml | 16 + dbeam-parquet/pom.xml | 88 ++++++ .../dbeam/parquet/ChannelOutputFile.java | 110 +++++++ .../dbeam/parquet/JdbcParquetArgs.java | 122 ++++++++ .../spotify/dbeam/parquet/JdbcParquetIO.java | 276 ++++++++++++++++++ .../spotify/dbeam/parquet/JdbcParquetJob.java | 201 +++++++++++++ .../dbeam/parquet/JdbcParquetMetering.java | 105 +++++++ .../dbeam/parquet/JdbcParquetSchema.java | 223 ++++++++++++++ .../parquet/JdbcParquetWriteSupport.java | 271 +++++++++++++++++ .../dbeam/parquet/JdbcParquetJobTest.java | 118 ++++++++ pom.xml | 55 +++- 12 files changed, 1575 insertions(+), 15 deletions(-) create mode 100644 dbeam-parquet/pom.xml create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetMetering.java create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java create mode 100644 dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java diff --git a/dbeam-bom/pom.xml b/dbeam-bom/pom.xml index 5c6d6e56..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} + diff --git a/dbeam-core/pom.xml b/dbeam-core/pom.xml index 6e2cd66e..711bec99 100644 --- a/dbeam-core/pom.xml +++ b/dbeam-core/pom.xml @@ -124,4 +124,20 @@ + + + + org.apache.maven.plugins + maven-jar-plugin + + + + test-jar + + + + + + + diff --git a/dbeam-parquet/pom.xml b/dbeam-parquet/pom.xml new file mode 100644 index 00000000..e241cede --- /dev/null +++ b/dbeam-parquet/pom.xml @@ -0,0 +1,88 @@ + + + + 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 + + + + + com.spotify + dbeam-core + ${project.version} + test-jar + test + + + junit + junit + test + + + org.hamcrest + hamcrest-core + test + + + org.hamcrest + hamcrest-library + test + + + com.h2database + h2 + test + + + + 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..316fee96 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java @@ -0,0 +1,110 @@ +/*- + * -\-\- + * 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; + + public ChannelOutputFile(WritableByteChannel channel) { + this.channel = channel; + } + + @Override + public PositionOutputStream create(long blockSizeHint) throws IOException { + return new ChannelPositionOutputStream(Channels.newOutputStream(channel)); + } + + @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..874658fe --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java @@ -0,0 +1,122 @@ +/*- + * -\-\- + * 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 int rowGroupSize(); + + public abstract int pageSize(); + + public abstract List preCommand(); + + 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(int rowGroupSize); + + abstract Builder setPageSize(int pageSize); + + abstract Builder setPreCommand(List preCommand); + + abstract JdbcParquetArgs build(); + } + + public static JdbcParquetArgs create( + final JdbcConnectionArgs jdbcConnectionArgs, + final int fetchSize, + final String parquetCodec, + final int rowGroupSize, + final int pageSize, + final List preCommand) { + 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) + .build(); + } + + public static JdbcParquetArgs create(final JdbcConnectionArgs jdbcConnectionArgs) { + return create( + jdbcConnectionArgs, + 10000, + "snappy", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList()); + } +} 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..0528d7f6 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.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 com.google.common.base.Preconditions.checkArgument; + +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 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 PTransform, WriteFilesResult> createWrite( + final String filenamePrefix, + final String filenameSuffix, + final MessageType schema, + final JdbcParquetArgs jdbcParquetArgs) { + 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); + return WriteFiles.to(sink); + } + + static class JdbcParquetSink extends FileBasedSink { + + private static final long serialVersionUID = 937707428039L; + private final String schemaString; + private final JdbcParquetArgs jdbcParquetArgs; + + JdbcParquetSink( + final ValueProvider filenamePrefix, + final DynamicDestinations destinations, + final String schemaString, + final JdbcParquetArgs jdbcParquetArgs) { + super(filenamePrefix, destinations); + this.schemaString = schemaString; + this.jdbcParquetArgs = jdbcParquetArgs; + } + + @Override + public WriteOperation createWriteOperation() { + return new JdbcParquetWriteOperation(this, schemaString, jdbcParquetArgs); + } + } + + private static class JdbcParquetWriteOperation + extends FileBasedSink.WriteOperation { + + private static final long serialVersionUID = 305340251351L; + private final String schemaString; + private final JdbcParquetArgs jdbcParquetArgs; + + private JdbcParquetWriteOperation( + final FileBasedSink sink, + final String schemaString, + final JdbcParquetArgs jdbcParquetArgs) { + super(sink); + this.schemaString = schemaString; + this.jdbcParquetArgs = jdbcParquetArgs; + } + + @Override + public FileBasedSink.Writer createWriter() { + return new JdbcParquetWriter(this, schemaString, jdbcParquetArgs); + } + } + + 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 ParquetWriter parquetWriter; + private Connection connection; + private JdbcParquetMetering metering; + + JdbcParquetWriter( + FileBasedSink.WriteOperation writeOperation, + String schemaString, + JdbcParquetArgs jdbcParquetArgs) { + super(writeOperation, MimeTypes.BINARY); + this.schemaString = schemaString; + this.jdbcParquetArgs = jdbcParquetArgs; + this.metering = JdbcParquetMetering.create(); + } + + 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); + final OutputFile outputFile = new ChannelOutputFile(channel); + parquetWriter = new ResultSetParquetWriterBuilder(outputFile, schema) + .withCompressionCodec(jdbcParquetArgs.getCompressionCodecName()) + .withRowGroupSize(jdbcParquetArgs.rowGroupSize()) + .withPageSize(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(); + this.metering.exposeWrittenBytes(0); + } + } + + @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; + + ResultSetParquetWriterBuilder(OutputFile outputFile, MessageType schema) { + super(outputFile); + this.schema = schema; + } + + @Override + protected ResultSetParquetWriterBuilder self() { + return this; + } + + @Override + protected WriteSupport getWriteSupport(Configuration conf) { + return new ResultSetWriteSupport(schema); + } + } + + static class ResultSetWriteSupport extends WriteSupport { + + private final MessageType schema; + private JdbcParquetWriteSupport writeSupport; + private RecordConsumer recordConsumer; + private boolean initialized = false; + + ResultSetWriteSupport(MessageType schema) { + this.schema = schema; + } + + @Override + public WriteContext init(Configuration configuration) { + return new WriteContext(schema, Collections.emptyMap()); + } + + @Override + public void prepareForWrite(RecordConsumer recordConsumer) { + this.recordConsumer = recordConsumer; + } + + @Override + public void write(ResultSet resultSet) { + try { + if (!initialized) { + this.writeSupport = JdbcParquetWriteSupport.create(resultSet, schema); + 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..3f0a657b --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -0,0 +1,201 @@ +/*- + * -\-\- + * 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.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.util.List; +import java.util.Map; +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 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); + 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; + try (Connection connection = jdbcExportArgs.createConnection()) { + generatedSchema = JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + jdbcExportArgs.queryBuilderArgs(), + jdbcExportArgs.avroSchemaName(), + jdbcExportArgs.useAvroLogicalTypes()); + + 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()); + 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 JdbcParquetArgs parquetArgs = JdbcParquetArgs.create( + jdbcExportArgs.jdbcAvroOptions().jdbcConnectionConfiguration(), + jdbcExportArgs.jdbcAvroOptions().fetchSize(), + pipelineOptions.as(JdbcExportPipelineOptions.class).getAvroCodec() + .replace("deflate6", "snappy"), + 64 * 1024 * 1024, + 1024 * 1024, + jdbcExportArgs.jdbcAvroOptions().preCommand()); + + pipeline + .apply("JdbcQueries", Create.of(queries)) + .apply( + "JdbcParquetSave", + JdbcParquetIO.createWrite(output, ".parquet", generatedSchema, parquetArgs)); + } + + 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(JdbcParquetMetering.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; + } + + public Pipeline getPipeline() { + return pipeline; + } + + public JdbcExportArgs getJdbcExportArgs() { + return jdbcExportArgs; + } + + public String getOutput() { + return output; + } + + 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..79157945 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetMetering.java @@ -0,0 +1,105 @@ +/*- + * -\-\- + * 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.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 JdbcParquetMetering { + + public static final String RECORD_COUNT_METRIC_NAME = "recordCount"; + private final int countReportEvery; + private final int logEvery; + private final Logger logger = LoggerFactory.getLogger(JdbcParquetMetering.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 JdbcParquetMetering(int countReportEvery, int logEvery) { + this.countReportEvery = countReportEvery; + this.logEvery = logEvery; + } + + public static JdbcParquetMetering create() { + return new JdbcParquetMetering(100000, 100000); + } + + 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( + "jdbcparquetio : 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("jdbcparquetio : 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("jdbcparquetio : 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-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..8394f22b --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -0,0 +1,223 @@ +/*- + * -\-\- + * 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) + 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); + 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) + 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( + normalizeForAvro(columnName), + columnType, + precision, + columnClassName, + columnTypeName, + useLogicalTypes)); + } + + 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 normalizeForAvro(metaTableName); + } + } + return defaultTableName; + } + + static Type buildParquetFieldType( + final String columnName, + final int columnType, + final int precision, + final String columnClassName, + final String columnTypeName, + final boolean useLogicalTypes) { + 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: + 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: + // For now, serialize arrays as binary (bytes mode) + return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) + .named(columnName); + case OTHER: + if (useLogicalTypes && "uuid".equals(columnTypeName)) { + return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) + .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); + } + } + + private static String normalizeForAvro(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..5ce907df --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -0,0 +1,271 @@ +/*- + * -\-\- + * 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.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 org.apache.parquet.io.api.Binary; +import org.apache.parquet.io.api.RecordConsumer; +import org.apache.parquet.schema.MessageType; + +/** + * 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 Calendar CALENDAR = 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) + 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); + } + + 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) + 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 = fieldName.replaceAll("[^A-Za-z0-9_]", "_"); + + 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: + case TIME: + case TIME_WITH_TIMEZONE: + return (consumer, rs) -> { + final Timestamp timestamp = rs.getTimestamp(column, CALENDAR); + 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: + 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 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")) { + 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); + } + }; + } + } +} 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..19d578aa --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java @@ -0,0 +1,118 @@ +/*- + * -\-\- + * 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.Assert; +import org.junit.BeforeClass; +import org.junit.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; + + @BeforeClass + public static void beforeAll() throws SQLException, ClassNotFoundException, IOException { + testDir = TestHelper.createTmpDirPath("jdbc-parquet-test-"); + passwordPath = testDir.resolve(".password"); + passwordPath.toFile().createNewFile(); + 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( + "_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"))); + Assert.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")); + } +} diff --git a/pom.xml b/pom.xml index 5370cb54..4984fbd3 100644 --- a/pom.xml +++ b/pom.xml @@ -76,6 +76,7 @@ dbeam-core + dbeam-parquet dbeam-bom @@ -121,10 +122,16 @@ 4.4.14 2.18.2 2.14.0 - 4.1.124.Final + 4.1.130.Final 2.0.16 1.6.8 1.5.7-7 + 1.14.4 + 3.5.0 + + 7.0.0 + 4.2.2 @@ -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 @@ -431,12 +464,6 @@ - - - com.google.auto.value - auto-value - - @@ -470,14 +497,6 @@ - - - - compile - - compile - - maven-release-plugin @@ -566,6 +585,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 From 3aaad51971684301b764537039305d38013b9905 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 09:47:51 +0200 Subject: [PATCH 30/80] Add comprehensive tests for dbeam-parquet module Test classes added (58 tests total): - ChannelOutputFileTest: position tracking, byte correctness, block size behavior, createOrOverwrite equivalence - JdbcParquetSchemaTest: SQL type to Parquet type conversion for all supported types (BIGINT, INTEGER, SMALLINT, TINYINT, TIMESTAMP, DATE, TIME, BOOLEAN, BIT, BINARY, BLOB, DOUBLE, FLOAT, REAL, VARCHAR, CHAR, CLOB, ARRAY, UUID, defaults), table name extraction, logical types toggle, custom schema names, character normalization, full schema creation from mock ResultSet - JdbcParquetRecordTest: schema creation from H2 database, write and read-back verification using ParquetReader, field value assertions, null value handling (optional fields), logical types from DB - JdbcParquetArgsTest: default values, all codec mappings (snappy, gzip, zstd, lz4, none, uncompressed), invalid codec rejection, custom fetch/row group sizes - JdbcParquetJobTest: expanded with SQL file input, logical types flag, minRows validation Also fixes UUID logical type to use FIXED_LEN_BYTE_ARRAY(16) per Parquet spec, and adds mockito + hadoop-mapreduce-client-core test dependencies. --- dbeam-parquet/pom.xml | 13 + .../dbeam/parquet/JdbcParquetSchema.java | 3 +- .../dbeam/parquet/ChannelOutputFileTest.java | 94 +++++ .../dbeam/parquet/JdbcParquetArgsTest.java | 120 ++++++ .../dbeam/parquet/JdbcParquetJobTest.java | 100 +++++ .../dbeam/parquet/JdbcParquetRecordTest.java | 273 +++++++++++++ .../dbeam/parquet/JdbcParquetSchemaTest.java | 372 ++++++++++++++++++ 7 files changed, 974 insertions(+), 1 deletion(-) create mode 100644 dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/ChannelOutputFileTest.java create mode 100644 dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java create mode 100644 dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java create mode 100644 dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java diff --git a/dbeam-parquet/pom.xml b/dbeam-parquet/pom.xml index e241cede..b52ab932 100644 --- a/dbeam-parquet/pom.xml +++ b/dbeam-parquet/pom.xml @@ -55,6 +55,14 @@ runtime + + + org.apache.hadoop + hadoop-mapreduce-client-core + ${hadoop.version} + test + + com.spotify @@ -78,6 +86,11 @@ hamcrest-library test + + org.mockito + mockito-core + test + com.h2database h2 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 index 8394f22b..777d572e 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -199,7 +199,8 @@ static Type buildParquetFieldType( .named(columnName); case OTHER: if (useLogicalTypes && "uuid".equals(columnTypeName)) { - return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) + return Types.optional(PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY) + .length(16) .as(LogicalTypeAnnotation.uuidType()) .named(columnName); } 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..e684d1f7 --- /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.Assert; +import org.junit.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)) { + Assert.assertEquals(0, out.getPos()); + + out.write(42); + Assert.assertEquals(1, out.getPos()); + + out.write(new byte[] {1, 2, 3, 4, 5}); + Assert.assertEquals(6, out.getPos()); + + out.write(new byte[] {0, 1, 2, 3, 4, 5, 6, 7}, 2, 3); + Assert.assertEquals(9, out.getPos()); + } + + Assert.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(); + Assert.assertEquals(3, result.length); + Assert.assertEquals(65, result[0]); + Assert.assertEquals(66, result[1]); + Assert.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); + + Assert.assertFalse(outputFile.supportsBlockSize()); + Assert.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)) { + Assert.assertNotNull(out); + Assert.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..4a4c149e --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java @@ -0,0 +1,120 @@ +/*- + * -\-\- + * 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.Assert; +import org.junit.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); + + Assert.assertEquals(10000, args.fetchSize()); + Assert.assertEquals("snappy", args.parquetCodec()); + Assert.assertEquals(64 * 1024 * 1024, args.rowGroupSize()); + Assert.assertEquals(1024 * 1024, args.pageSize()); + Assert.assertEquals(Collections.emptyList(), args.preCommand()); + } + + @Test + public void shouldMapSnappyCodec() { + final JdbcParquetArgs args = JdbcParquetArgs.create(CONN_ARGS); + + Assert.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()); + + Assert.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()); + + Assert.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()); + + Assert.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()); + + Assert.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()); + + Assert.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); + } + + @Test(expected = IllegalArgumentException.class) + public void shouldRejectInvalidCodec() { + JdbcParquetArgs.create( + CONN_ARGS, 10000, "invalid", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + } + + @Test + public void shouldAcceptCustomFetchSize() { + final JdbcParquetArgs args = JdbcParquetArgs.create( + CONN_ARGS, 50000, "snappy", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + + Assert.assertEquals(50000, args.fetchSize()); + } + + @Test + public void shouldAcceptCustomRowGroupSize() { + final JdbcParquetArgs args = JdbcParquetArgs.create( + CONN_ARGS, 10000, "snappy", 128 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + + Assert.assertEquals(128 * 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 index 19d578aa..a3c80391 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java @@ -41,12 +41,15 @@ public class JdbcParquetJobTest { "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; @BeforeClass 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); } @@ -115,4 +118,101 @@ public void shouldRunJdbcParquetJobDataOnly() throws Exception { 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( + "_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"))); + Assert.assertTrue(schemaJson.contains("COF_NAME")); + Assert.assertTrue(schemaJson.contains("SIZE")); + Assert.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( + "_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"))); + Assert.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)); + } } 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..c8355b25 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java @@ -0,0 +1,273 @@ +/*- + * -\-\- + * 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.Optional; +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.Assert; +import org.junit.BeforeClass; +import org.junit.Test; + +public class JdbcParquetRecordTest { + + private static final String CONNECTION_URL = + "jdbc:h2:mem:testrecord;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; + + @BeforeClass + 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); + + Assert.assertNotNull(schema); + Assert.assertEquals("COFFEES", schema.getName()); + Assert.assertEquals(14, schema.getFieldCount()); + Assert.assertEquals("COF_NAME", schema.getFields().get(0).getName()); + Assert.assertEquals("SUP_ID", schema.getFields().get(1).getName()); + Assert.assertEquals("PRICE", schema.getFields().get(2).getName()); + Assert.assertEquals("TEMPERATURE", schema.getFields().get(3).getName()); + Assert.assertEquals("SIZE", schema.getFields().get(4).getName()); + Assert.assertEquals("IS_ARABIC", schema.getFields().get(5).getName()); + Assert.assertEquals("SALES", schema.getFields().get(6).getName()); + Assert.assertEquals("TOTAL", schema.getFields().get(7).getName()); + Assert.assertEquals("CREATED", schema.getFields().get(8).getName()); + Assert.assertEquals("UPDATED", schema.getFields().get(9).getName()); + Assert.assertEquals("UID", schema.getFields().get(10).getName()); + Assert.assertEquals("ROWNUM", schema.getFields().get(11).getName()); + Assert.assertEquals("INT_ARR", schema.getFields().get(12).getName()); + Assert.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); + + Assert.assertEquals(14, schema.getFieldCount()); + // CREATED and UPDATED should have timestamp logical type + Assert.assertNotNull(schema.getFields().get(8).asPrimitiveType().getLogicalTypeAnnotation()); + Assert.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); + + Assert.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); + + // 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(); + Assert.assertNotNull(record1); + // Verify first record has expected fields + String cofName = record1.getString("COF_NAME", 0); + Assert.assertNotNull(cofName); + + Group record2 = reader.read(); + Assert.assertNotNull(record2); + + // Should only be 2 records + Group record3 = reader.read(); + Assert.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); + + 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(); + Assert.assertNotNull(record); + + // Verify typed values for the first record (colombian caffee, sorted) + String cofName = record.getString("COF_NAME", 0); + Assert.assertEquals("colombian caffee", cofName); + + boolean isArabic = record.getBoolean("IS_ARABIC", 0); + Assert.assertTrue(isArabic); + + int sales = record.getInteger("SALES", 0); + Assert.assertEquals(13, sales); + + long total = record.getLong("TOTAL", 0); + Assert.assertEquals(201L, total); + + float temperature = record.getFloat("TEMPERATURE", 0); + Assert.assertEquals(87.5f, temperature, 0.01f); + + double size = record.getDouble("SIZE", 0); + Assert.assertEquals(230.7, size, 0.01); + + long rownum = record.getLong("ROWNUM", 0); + Assert.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); + + 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(); + Assert.assertNotNull(record); + + // SUP_ID is null - should have 0 repetitions in Parquet + Assert.assertEquals(0, record.getFieldRepetitionCount("SUP_ID")); + + // UPDATED is null + Assert.assertEquals(0, record.getFieldRepetitionCount("UPDATED")); + + // Non-null fields should have 1 repetition + Assert.assertEquals(1, record.getFieldRepetitionCount("COF_NAME")); + Assert.assertEquals(1, record.getFieldRepetitionCount("IS_ARABIC")); + } + + 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..0f87b4b6 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java @@ -0,0 +1,372 @@ +/*- + * -\-\- + * 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.Assert; +import org.junit.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"); + + Assert.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(""); + + Assert.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); + + Assert.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"); + + Assert.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); + + 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); + Assert.assertNull(fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertTimestampSqlTypeWithLogicalType() throws SQLException { + final Type fieldType = buildFieldType(Types.TIMESTAMP, true); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); + Assert.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); + Assert.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); + + 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); + Assert.assertEquals( + LogicalTypeAnnotation.stringType(), + fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertCharSqlTypeToStringBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.CHAR, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assert.assertEquals( + LogicalTypeAnnotation.stringType(), + fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertClobSqlTypeToStringBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.CLOB, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assert.assertEquals( + LogicalTypeAnnotation.stringType(), + fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertArraySqlTypeToBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.ARRAY, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + } + + @Test + public void shouldConvertUuidWithLogicalType() throws SQLException { + final Type fieldType = JdbcParquetSchema.buildParquetFieldType( + "column1", Types.OTHER, 0, "foobar", "uuid", true); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY); + Assert.assertEquals(16, fieldType.asPrimitiveType().getTypeLength()); + Assert.assertEquals( + LogicalTypeAnnotation.uuidType(), + fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldConvertUuidWithoutLogicalType() throws SQLException { + final Type fieldType = JdbcParquetSchema.buildParquetFieldType( + "column1", Types.OTHER, 0, "foobar", "uuid", false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assert.assertEquals( + LogicalTypeAnnotation.stringType(), + fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldDefaultConversionToStringBinary() throws SQLException { + final Type fieldType = buildFieldType(Types.SQLXML, false); + + assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assert.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); + + Assert.assertEquals("test_table", schema.getName()); + Assert.assertEquals(3, schema.getFieldCount()); + Assert.assertEquals("id", schema.getFields().get(0).getName()); + Assert.assertEquals("name", schema.getFields().get(1).getName()); + Assert.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); + Assert.assertTrue(schema.getFields().get(0).isRepetition(Type.Repetition.OPTIONAL)); + Assert.assertTrue(schema.getFields().get(1).isRepetition(Type.Repetition.OPTIONAL)); + Assert.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); + + Assert.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); + + Assert.assertEquals("test_table_name", schema.getName()); + Assert.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); + } + + private void assertPrimitiveType( + final Type actual, final PrimitiveType.PrimitiveTypeName expected) { + Assert.assertTrue(actual.isPrimitive()); + Assert.assertEquals(expected, actual.asPrimitiveType().getPrimitiveTypeName()); + } +} From 2fb1cdd4539c6f0401125cd41f13b00d84ddcf42 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 10:19:36 +0200 Subject: [PATCH 31/80] Add BenchJdbcParquetJob and Parquet support in e2e tests Add BenchJdbcParquetJob mirroring BenchJdbcAvroJob: runs the Parquet export N times (configurable via --executions), collects metrics from each run, and outputs a TSV summary table with per-run values and statistical aggregates (max, mean, min, stddev). Update e2e/e2e.sh with Parquet support: - run_docker_dbeam_parquet: runs BenchJdbcParquetJob in Docker using the dbeam-parquet-shaded.jar - runDBeamParquetDockerCon: orchestrates a Parquet export run with connection params, output validation - runParquetSuite: runs snappy codec and parallel query configurations - lightParquet: quick single-run Parquet e2e test Add shade plugin profile (pack) to dbeam-parquet POM for building the shaded JAR with JdbcParquetJob as main class. Add BenchJdbcParquetJobTest verifying multi-execution output structure. --- dbeam-parquet/pom.xml | 49 ++++++ .../dbeam/parquet/BenchJdbcParquetJob.java | 140 ++++++++++++++++++ .../parquet/BenchJdbcParquetJobTest.java | 61 ++++++++ e2e/e2e.sh | 44 +++++- 4 files changed, 293 insertions(+), 1 deletion(-) create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BenchJdbcParquetJob.java create mode 100644 dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java diff --git a/dbeam-parquet/pom.xml b/dbeam-parquet/pom.xml index b52ab932..d5262a01 100644 --- a/dbeam-parquet/pom.xml +++ b/dbeam-parquet/pom.xml @@ -98,4 +98,53 @@ + + + pack + + + + org.apache.maven.plugins + maven-shade-plugin + 3.6.0 + + + bundle-and-repackage + + shade + + package + + ${basedir}/target/dependency-reduced-pom.xml + + + *:* + + + + + *:* + + META-INF/*.SF + META-INF/*.DSA + META-INF/*.RSA + + + + + + + com.spotify.dbeam.parquet.JdbcParquetJob + + + ${basedir}/target/${project.artifactId}-shaded.jar + + + + + + + + + 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..934289aa --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BenchJdbcParquetJob.java @@ -0,0 +1,140 @@ +/*- + * -\-\- + * 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.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.jobs.ExceptionHandling; +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 BenchJdbcParquetJob { + + public interface BenchJdbcParquetOptions extends PipelineOptions { + @Description("Number of benchmark executions.") + @Default.Integer(3) + int getExecutions(); + + void setExecutions(int value); + } + + private final PipelineOptions pipelineOptions; + private List> metrics = newArrayList(); + + public BenchJdbcParquetJob(final PipelineOptions pipelineOptions) { + this.pipelineOptions = pipelineOptions; + } + + public static BenchJdbcParquetJob create(final String[] cmdLineArgs) { + PipelineOptionsFactory.register(BenchJdbcParquetOptions.class); + PipelineOptions options = JdbcParquetJob.buildPipelineOptions(cmdLineArgs); + return new BenchJdbcParquetJob(options); + } + + public void run() throws Exception { + int executions = pipelineOptions.as(BenchJdbcParquetOptions.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 = + JdbcParquetJob.create(pipelineOptions, output).runExport(); + this.metrics.add(MetricsHelper.getMetrics(pipelineResult)); + } + System.out.println("Summary for BenchJdbcParquetJob"); + 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 void main(String[] cmdLineArgs) { + try { + create(cmdLineArgs).run(); + } catch (Exception e) { + ExceptionHandling.handleException(e); + } + } +} 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..e9c52d34 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java @@ -0,0 +1,61 @@ +/*- + * -\-\- + * 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 com.spotify.dbeam.DbTestHelper; +import com.spotify.dbeam.TestHelper; +import java.io.IOException; +import java.nio.file.Path; +import java.sql.SQLException; +import org.junit.BeforeClass; +import org.junit.Test; + +public class BenchJdbcParquetJobTest { + + private static final String CONNECTION_URL = + "jdbc:h2:mem:testbench;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; + private static Path testDir; + + @BeforeClass + public static void beforeAll() throws SQLException, ClassNotFoundException, IOException { + testDir = TestHelper.createTmpDirPath("jdbc-parquet-bench-test"); + DbTestHelper.createFixtures(CONNECTION_URL); + } + + @Test + public void shouldRunBenchJdbcParquetJob() { + BenchJdbcParquetJob.main( + new String[] { + "--targetParallelism=1", + "--skipPartitionCheck", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--table=COFFEES", + "--output=" + testDir.toString(), + "--avroCodec=snappy", + "--executions=2" + }); + assertThat(TestHelper.listDir(testDir.toFile()), containsInAnyOrder("run_0", "run_1")); + } +} diff --git a/e2e/e2e.sh b/e2e/e2e.sh index 2f0c587e..7ab21894 100755 --- a/e2e/e2e.sh +++ b/e2e/e2e.sh @@ -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,25 @@ 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 ($(stat -f%z "$OUTPUT_FILE" 2>/dev/null || stat -c%s "$OUTPUT_FILE") bytes)" +} + runSuite() { table=demo_table BINARY_TRANSFER='false' runDBeamDockerCon --executions=3 --avroCodec=deflate1 @@ -108,12 +138,24 @@ runSuite() { BINARY_TRANSFER='false' runDBeamDockerCon --executions=3 --avroCodec=deflate1 --arrayMode=typed_postgres } +runParquetSuite() { + table=demo_table + BINARY_TRANSFER='false' runDBeamParquetDockerCon --executions=3 --avroCodec=snappy + BINARY_TRANSFER='false' runDBeamParquetDockerCon --executions=3 --avroCodec=snappy --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 From 2af6b6b5f3459d5b8c80b8112075a5110e951497 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 11:32:50 +0200 Subject: [PATCH 32/80] Use runParquetSuite --- e2e/e2e.sh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/e2e/e2e.sh b/e2e/e2e.sh index 7ab21894..c491bf43 100755 --- a/e2e/e2e.sh +++ b/e2e/e2e.sh @@ -140,8 +140,8 @@ runSuite() { runParquetSuite() { table=demo_table - BINARY_TRANSFER='false' runDBeamParquetDockerCon --executions=3 --avroCodec=snappy - BINARY_TRANSFER='false' runDBeamParquetDockerCon --executions=3 --avroCodec=snappy --queryParallelism=5 --splitColumn=row_number + BINARY_TRANSFER='false' runDBeamParquetDockerCon --executions=3 + BINARY_TRANSFER='false' runDBeamParquetDockerCon --executions=3 --queryParallelism=5 --splitColumn=row_number } light() { @@ -166,6 +166,7 @@ main() { time startPostgres runSuite + runParquetSuite dockerClean fi } From 02b988dabffc7a6e44c82db068516083df357909 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 11:43:36 +0200 Subject: [PATCH 33/80] Fix avroCodec to Parquet codec mapping in JdbcParquetJob MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the fragile .replace("deflate6", "snappy") hack with a proper mapAvroCodecToParquetCodec() method that handles all Avro codec names: - deflate1-9 → gzip - zstandard1-9 → zstd - snappy → snappy Previously, only deflate6 was handled — all other deflate/zstandard values would fail JdbcParquetArgs validation. Now all valid --avroCodec values work correctly with Parquet exports. --- .../spotify/dbeam/parquet/JdbcParquetJob.java | 15 ++++- .../dbeam/parquet/JdbcParquetJobTest.java | 56 +++++++++++++++++++ 2 files changed, 69 insertions(+), 2 deletions(-) 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 index 3f0a657b..9df6bb19 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -140,8 +140,8 @@ public void prepareExport() throws Exception { final JdbcParquetArgs parquetArgs = JdbcParquetArgs.create( jdbcExportArgs.jdbcAvroOptions().jdbcConnectionConfiguration(), jdbcExportArgs.jdbcAvroOptions().fetchSize(), - pipelineOptions.as(JdbcExportPipelineOptions.class).getAvroCodec() - .replace("deflate6", "snappy"), + mapAvroCodecToParquetCodec( + pipelineOptions.as(JdbcExportPipelineOptions.class).getAvroCodec()), 64 * 1024 * 1024, 1024 * 1024, jdbcExportArgs.jdbcAvroOptions().preCommand()); @@ -191,6 +191,17 @@ public String getOutput() { return output; } + static String mapAvroCodecToParquetCodec(final String avroCodec) { + 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(); 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 index a3c80391..bd9e6fd8 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java @@ -215,4 +215,60 @@ public void shouldRunParquetJobWithMinRows() throws Exception { final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); assertThat(parquetFile.length(), greaterThan(0L)); } + + @Test + public void shouldMapAvroCodecToParquetCodec() { + Assert.assertEquals("snappy", JdbcParquetJob.mapAvroCodecToParquetCodec("snappy")); + Assert.assertEquals("gzip", JdbcParquetJob.mapAvroCodecToParquetCodec("deflate1")); + Assert.assertEquals("gzip", JdbcParquetJob.mapAvroCodecToParquetCodec("deflate6")); + Assert.assertEquals("gzip", JdbcParquetJob.mapAvroCodecToParquetCodec("deflate9")); + Assert.assertEquals("zstd", JdbcParquetJob.mapAvroCodecToParquetCodec("zstandard1")); + Assert.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)); + } } From 3636d6afbc43a591a129d68da0ff5e658dcc9cfe Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 11:44:40 +0200 Subject: [PATCH 34/80] Add BeamJdbcParquetSchema for schema metrics and input file support Add BeamJdbcParquetSchema mirroring BeamJdbcAvroSchema from dbeam-core: - createSchema(): wraps JdbcParquetSchema.createSchemaByReadingOneRow() and exposes schemaElapsedTimeMs as a Beam counter metric - parseOptionalInputParquetSchemaFile(): reads a Parquet MessageType schema from a file, enabling --avroSchemaFilePath for Parquet exports - parseInputParquetSchemaFile(): parses schema text via MessageTypeParser Update JdbcParquetJob to use BeamJdbcParquetSchema: - Schema creation now goes through createSchema() which checks for input schema file before falling back to database inference - Schema generation timing is now visible in Beam metrics --- .../dbeam/parquet/BeamJdbcParquetSchema.java | 107 ++++++++++++++++++ .../spotify/dbeam/parquet/JdbcParquetJob.java | 19 +++- 2 files changed, 120 insertions(+), 6 deletions(-) create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java 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..f4c2d613 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java @@ -0,0 +1,107 @@ +/*- + * -\-\- + * 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.JdbcExportArgs; +import java.io.IOException; +import java.io.InputStream; +import java.nio.channels.Channels; +import java.sql.Connection; +import java.sql.SQLException; +import java.util.Collections; +import java.util.Optional; +import java.util.Scanner; +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 Logger LOGGER = LoggerFactory.getLogger(BeamJdbcParquetSchema.class); + + /** + * Generate Parquet schema by reading one row. Expose Beam metrics via a Beam PTransform. + * + * @param pipeline Beam SDK pipeline, to expose metrics + * @param args JdbcExportArgs with DBeam configuration + * @param connection JDBC connection to query input schema + * @return Parquet MessageType schema based on configuration + * @throws Exception in case of failure to query database + */ + public static MessageType createSchema( + final Pipeline pipeline, final JdbcExportArgs args, final Connection connection) + throws Exception { + final long startTime = System.nanoTime(); + final MessageType generatedSchema = generateParquetSchema(args, connection); + final long elapsedTimeSchema = (System.nanoTime() - startTime) / 1000000; + LOGGER.info("Elapsed time to schema {} seconds", elapsedTimeSchema / 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(elapsedTimeSchema); + return v; + })); + return generatedSchema; + } + + private static MessageType generateParquetSchema( + final JdbcExportArgs args, final Connection connection) throws SQLException { + return JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + args.queryBuilderArgs(), + args.avroSchemaName(), + args.useAvroLogicalTypes()); + } + + 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); + final InputStream inputStream = Channels.newInputStream(FileSystems.open(m.resourceId())); + final String schemaText = + new Scanner(inputStream, "UTF-8").useDelimiter("\\A").next(); + return MessageTypeParser.parseMessageType(schemaText); + } +} 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 index 9df6bb19..76bcb59c 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -116,12 +116,7 @@ public void prepareExport() throws Exception { final List queries; final MessageType generatedSchema; try (Connection connection = jdbcExportArgs.createConnection()) { - generatedSchema = JdbcParquetSchema.createSchemaByReadingOneRow( - connection, - jdbcExportArgs.queryBuilderArgs(), - jdbcExportArgs.avroSchemaName(), - jdbcExportArgs.useAvroLogicalTypes()); - + generatedSchema = createSchema(connection); queries = jdbcExportArgs.queryBuilderArgs().buildQueries(connection); final String tableName = pipelineOptions.as(DBeamPipelineOptions.class).getTable(); @@ -179,6 +174,18 @@ public PipelineResult runExport() throws Exception { return pipelineResult; } + private MessageType createSchema(final Connection connection) throws Exception { + final String schemaFilePath = + pipelineOptions.as(JdbcExportPipelineOptions.class).getAvroSchemaFilePath(); + final java.util.Optional inputSchema = + BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(schemaFilePath); + if (inputSchema.isPresent()) { + return inputSchema.get(); + } else { + return BeamJdbcParquetSchema.createSchema(this.pipeline, jdbcExportArgs, connection); + } + } + public Pipeline getPipeline() { return pipeline; } From 61b8ce19ce8d53a2c73e5ef4d9b522517a14a830 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 11:47:33 +0200 Subject: [PATCH 35/80] Add --parquetSchemaFilePath option for input schema override Add ParquetPipelineOptions interface with --parquetSchemaFilePath option, allowing users to provide a pre-defined Parquet MessageType schema file instead of inferring from the database. This is the Parquet equivalent of --avroSchemaFilePath. A separate option name is used because the Avro option goes through the Avro schema parser in JdbcExportArgsFactory, which would reject Parquet schema text format. The schema file should contain a Parquet MessageType definition, e.g.: message my_table { optional binary name (STRING); optional int64 id; } --- .../spotify/dbeam/parquet/JdbcParquetJob.java | 3 +- .../dbeam/parquet/ParquetPipelineOptions.java | 33 +++++++++++++++++ .../dbeam/parquet/JdbcParquetJobTest.java | 36 +++++++++++++++++++ 3 files changed, 71 insertions(+), 1 deletion(-) create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java 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 index 76bcb59c..21b29d50 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -98,6 +98,7 @@ public static JdbcParquetJob create(final String[] 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(); } @@ -176,7 +177,7 @@ public PipelineResult runExport() throws Exception { private MessageType createSchema(final Connection connection) throws Exception { final String schemaFilePath = - pipelineOptions.as(JdbcExportPipelineOptions.class).getAvroSchemaFilePath(); + pipelineOptions.as(ParquetPipelineOptions.class).getParquetSchemaFilePath(); final java.util.Optional inputSchema = BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(schemaFilePath); if (inputSchema.isPresent()) { 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..9808fd9e --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java @@ -0,0 +1,33 @@ +/*- + * -\-\- + * 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.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); +} 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 index bd9e6fd8..387727e6 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java @@ -216,6 +216,42 @@ public void shouldRunParquetJobWithMinRows() throws Exception { 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"))); + Assert.assertTrue(schemaJson.contains("COF_NAME")); + Assert.assertTrue(schemaJson.contains("SIZE")); + Assert.assertTrue(schemaJson.contains("TOTAL")); + final File parquetFile = outputPath.resolve("part-00000-of-00001.parquet").toFile(); + assertThat(parquetFile.length(), greaterThan(0L)); + } + @Test public void shouldMapAvroCodecToParquetCodec() { Assert.assertEquals("snappy", JdbcParquetJob.mapAvroCodecToParquetCodec("snappy")); From 220251ae3d081208b77ba2c94a21135fe3fc85d0 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 11:48:30 +0200 Subject: [PATCH 36/80] Add PsqlParquetJob for PostgreSQL exports with replication check Add PsqlParquetJob mirroring PsqlAvroJob from dbeam-core: checks PostgreSQL replication lag before starting the Parquet export. Reuses PsqlReplicationCheck from dbeam-core for the lag check logic. Also make PsqlReplicationCheck.validateOptions() public so it can be called from the dbeam-parquet module (was package-private, only accessible from PsqlAvroJob in the same package). --- .../dbeam/jobs/PsqlReplicationCheck.java | 2 +- .../spotify/dbeam/parquet/PsqlParquetJob.java | 56 +++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) create mode 100644 dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/PsqlParquetJob.java 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-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..6cf531b0 --- /dev/null +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/PsqlParquetJob.java @@ -0,0 +1,56 @@ +/*- + * -\-\- + * 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); + } + } +} From 51a07a69cee9c7ae049fa380cfd4bcf029744978 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 12:45:44 +0200 Subject: [PATCH 37/80] Include dbeam-parquet coverage in Codecov upload Add dbeam-parquet jacoco.xml to the Codecov upload so test coverage for the Parquet module is tracked alongside dbeam-core. --- .github/workflows/maven.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml index 49621c8d..f9a1db5e 100644 --- a/.github/workflows/maven.yml +++ b/.github/workflows/maven.yml @@ -44,7 +44,7 @@ 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) From 534c9647dc9734c91273864093520d3c778702ac Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 12:46:11 +0200 Subject: [PATCH 38/80] Add Parquet type conversion documentation Add docs/parquet-type-conversion.md documenting the SQL to Parquet type mapping, parallel to the existing docs/type-conversion.md for Avro. Covers all supported SQL types, logical type annotations, and the avroCodec-to-Parquet codec mapping table. --- docs/parquet-type-conversion.md | 46 +++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 docs/parquet-type-conversion.md diff --git a/docs/parquet-type-conversion.md b/docs/parquet-type-conversion.md new file mode 100644 index 00000000..8670045c --- /dev/null +++ b/docs/parquet-type-conversion.md @@ -0,0 +1,46 @@ +### 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 | BINARY | | Serialized as string representation | +| 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) | | +| 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) | | + +#### 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 | From 80d400d5e8d26a1895bf4ebfdeb277348c61f572 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 12:47:54 +0200 Subject: [PATCH 39/80] Add tests for PsqlParquetJob and BeamJdbcParquetSchema PsqlParquetJobTest: - Verifies non-PostgreSQL connections are rejected - Verifies missing partition parameter is rejected BeamJdbcParquetSchemaTest: - Verifies schema creation from H2 database with Beam metrics - Verifies input schema file parsing (MessageType text format) - Verifies null/empty filename returns Optional.empty() --- .../parquet/BeamJdbcParquetSchemaTest.java | 114 ++++++++++++++++++ .../dbeam/parquet/PsqlParquetJobTest.java | 68 +++++++++++ 2 files changed, 182 insertions(+) create mode 100644 dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java create mode 100644 dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PsqlParquetJobTest.java 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..30199a02 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java @@ -0,0 +1,114 @@ +/*- + * -\-\- + * 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.JdbcAvroArgs; +import com.spotify.dbeam.args.JdbcConnectionArgs; +import com.spotify.dbeam.args.JdbcExportArgs; +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.SQLException; +import java.time.Duration; +import java.util.Optional; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.parquet.schema.MessageType; +import org.junit.Assert; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; + +public class BeamJdbcParquetSchemaTest { + + private static final String CONNECTION_URL = + "jdbc:h2:mem:testbeamschema;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + @BeforeClass + public static void beforeAll() throws SQLException, ClassNotFoundException { + DbTestHelper.createFixtures(CONNECTION_URL); + } + + private static JdbcExportArgs createArgs() throws ClassNotFoundException { + return JdbcExportArgs.create( + JdbcAvroArgs.create(JdbcConnectionArgs.create(CONNECTION_URL)), + QueryBuilderArgs.create("COFFEES"), + "dbeam_generated", + Optional.empty(), + Optional.empty(), + false, + Duration.ofMinutes(1), + Optional.empty()); + } + + @Test + public void shouldCreateSchemaAndExposeMetrics() throws Exception { + final JdbcExportArgs args = createArgs(); + try (Connection connection = args.createConnection()) { + final MessageType schema = + BeamJdbcParquetSchema.createSchema(pipeline, args, connection); + + Assert.assertNotNull(schema); + Assert.assertEquals("COFFEES", schema.getName()); + Assert.assertEquals(14, schema.getFieldCount()); + } + pipeline.run().waitUntilFinish(); + } + + @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()); + + Assert.assertTrue(schema.isPresent()); + Assert.assertEquals("test_table", schema.get().getName()); + Assert.assertEquals(2, schema.get().getFieldCount()); + + Files.deleteIfExists(schemaFile); + } + + @Test + public void shouldReturnEmptyForNullFilename() throws IOException { + final Optional schema = + BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(null); + + Assert.assertFalse(schema.isPresent()); + } + + @Test + public void shouldReturnEmptyForEmptyFilename() throws IOException { + final Optional schema = + BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(""); + + Assert.assertFalse(schema.isPresent()); + } +} 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..6bee80e7 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PsqlParquetJobTest.java @@ -0,0 +1,68 @@ +/*- + * -\-\- + * 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.BeforeClass; +import org.junit.Test; + +public class PsqlParquetJobTest { + + private static Path testDir; + private static Path passwordPath; + + @BeforeClass + public static void beforeAll() throws IOException { + testDir = TestHelper.createTmpDirPath("psql-parquet-test-"); + passwordPath = testDir.resolve(".password"); + passwordPath.toFile().createNewFile(); + } + + @Test(expected = IllegalArgumentException.class) + public void shouldFailOnNonPostgresConnection() throws Exception { + final Path outputPath = testDir.resolve("shouldFailOnNonPostgres"); + PsqlParquetJob.create( + new String[] { + "--connectionUrl=jdbc:h2:mem:test", + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--partition=2025-02-28" + }); + } + + @Test(expected = IllegalArgumentException.class) + public void shouldFailOnMissingPartition() throws Exception { + final Path outputPath = testDir.resolve("shouldFailOnMissingPartition"); + PsqlParquetJob.create( + new String[] { + "--connectionUrl=jdbc:postgresql://localhost/test", + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--skipPartitionCheck" + }); + } +} From afd4d1f6957ff11849b27ab5f4c53a5b30f6d1c7 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:04:16 +0200 Subject: [PATCH 40/80] Consolidate shade plugin config using dbeam.mainClass property Replace duplicated shade plugin configuration in dbeam-parquet with a parameterized ${dbeam.mainClass} property in the parent POM's pack profile. Each module sets its own main class via a property: - dbeam-core: com.spotify.dbeam.jobs.JdbcAvroJob - dbeam-parquet: com.spotify.dbeam.parquet.JdbcParquetJob This removes ~40 lines of duplicated XML from dbeam-parquet/pom.xml while keeping the shade behavior identical. --- dbeam-core/pom.xml | 4 ++++ dbeam-parquet/pom.xml | 51 +++---------------------------------------- pom.xml | 2 +- 3 files changed, 8 insertions(+), 49 deletions(-) diff --git a/dbeam-core/pom.xml b/dbeam-core/pom.xml index 711bec99..e7a6b22b 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 + + diff --git a/dbeam-parquet/pom.xml b/dbeam-parquet/pom.xml index d5262a01..9c72de4f 100644 --- a/dbeam-parquet/pom.xml +++ b/dbeam-parquet/pom.xml @@ -98,53 +98,8 @@ - - - pack - - - - org.apache.maven.plugins - maven-shade-plugin - 3.6.0 - - - bundle-and-repackage - - shade - - package - - ${basedir}/target/dependency-reduced-pom.xml - - - *:* - - - - - *:* - - META-INF/*.SF - META-INF/*.DSA - META-INF/*.RSA - - - - - - - com.spotify.dbeam.parquet.JdbcParquetJob - - - ${basedir}/target/${project.artifactId}-shaded.jar - - - - - - - - + + com.spotify.dbeam.parquet.JdbcParquetJob + diff --git a/pom.xml b/pom.xml index 4984fbd3..8605aa92 100644 --- a/pom.xml +++ b/pom.xml @@ -796,7 +796,7 @@ - com.spotify.dbeam.jobs.JdbcAvroJob + ${dbeam.mainClass} ${basedir}/target/${project.artifactId}-shaded.jar From 249da71d5b1e0db8d0e16763f2596da7602d1126 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:05:49 +0200 Subject: [PATCH 41/80] Upgrade e2e PostgreSQL image from 16 to 18 --- e2e/e2e.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/e2e/e2e.sh b/e2e/e2e.sh index c491bf43..25c6e478 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 From a38e88d34d0aae99ce8342bbf32d77fb33fb686c Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:16:51 +0200 Subject: [PATCH 42/80] Add test for BenchJdbcParquetJob TSV metrics output Verify that BenchJdbcParquetJob prints the expected summary format: header line with column names (recordCount, writeElapsedMs, etc.), per-run lines (run_00, run_01), and statistical summary lines (max, mean, min, stddev). --- .../parquet/BenchJdbcParquetJobTest.java | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) 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 index e9c52d34..86bb1ac2 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java @@ -22,12 +22,16 @@ 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.IOException; +import java.io.PrintStream; import java.nio.file.Path; import java.sql.SQLException; +import org.junit.Assert; import org.junit.BeforeClass; import org.junit.Test; @@ -58,4 +62,42 @@ public void shouldRunBenchJdbcParquetJob() { }); assertThat(TestHelper.listDir(testDir.toFile()), containsInAnyOrder("run_0", "run_1")); } + + @Test + public void shouldOutputMetricsTsvSummary() throws IOException { + 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.main( + new String[] { + "--targetParallelism=1", + "--skipPartitionCheck", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--table=COFFEES", + "--output=" + benchDir.toString(), + "--avroCodec=snappy", + "--executions=2" + }); + } finally { + System.setOut(originalOut); + } + + final String output = capturedOut.toString(); + Assert.assertTrue(output.contains("Summary for BenchJdbcParquetJob")); + Assert.assertTrue(output.contains("recordCount")); + Assert.assertTrue(output.contains("writeElapsedMs")); + Assert.assertTrue(output.contains("bytesWritten")); + Assert.assertTrue(output.contains("KbWritePerSec")); + Assert.assertTrue(output.contains("run_00")); + Assert.assertTrue(output.contains("run_01")); + Assert.assertTrue(output.contains("max")); + Assert.assertTrue(output.contains("mean")); + Assert.assertTrue(output.contains("min")); + Assert.assertTrue(output.contains("stddev")); + // Each run should report 2 records + assertThat(output.split("run_").length, greaterThan(2)); + } } From d40083b6da706974ccf26e9aee5b7d5e1508fea1 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:17:56 +0200 Subject: [PATCH 43/80] Add PostgreSQL-specific mock tests for Parquet encoding Test Parquet write and read-back roundtrips using mocked ResultSet: - UUID field encoding as string - VARCHAR and OTHER type encoding - Timestamp encoding as millis - Null field handling (optional fields with 0 repetition count) - All numeric types (int, long, float, double) - Boolean type encoding Uses a reusable writeAndVerify helper that writes via ResultSetParquetWriterBuilder and reads back with GroupReadSupport. --- .../parquet/PostgresJdbcParquetTest.java | 242 ++++++++++++++++++ 1 file changed, 242 insertions(+) create mode 100644 dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java 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..b4e77473 --- /dev/null +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java @@ -0,0 +1,242 @@ +/*- + * -\-\- + * 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.Assert; +import org.junit.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); + + final Path tempFile = Files.createTempFile("parquet-uuid-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify(schema, resultSet, tempFile, record -> { + String actualUuid = record.getString("uuid_field", 0); + Assert.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); + + final Path tempFile = Files.createTempFile("parquet-string-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify(schema, resultSet, tempFile, record -> { + Assert.assertEquals("some_text_42", record.getString("text_field", 0)); + Assert.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); + + final Path tempFile = Files.createTempFile("parquet-ts-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify(schema, resultSet, tempFile, record -> { + long actualTs = record.getLong("ts_field", 0); + Assert.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); + + final Path tempFile = Files.createTempFile("parquet-null-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify(schema, resultSet, tempFile, record -> { + Assert.assertEquals("alice", record.getString("name", 0)); + Assert.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); + + final Path tempFile = Files.createTempFile("parquet-numeric-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify(schema, resultSet, tempFile, record -> { + Assert.assertEquals(42, record.getInteger("int_col", 0)); + Assert.assertEquals(9999999999L, record.getLong("long_col", 0)); + Assert.assertEquals(3.14f, record.getFloat("float_col", 0), 0.001f); + Assert.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); + + final Path tempFile = Files.createTempFile("parquet-bool-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify(schema, resultSet, tempFile, record -> { + Assert.assertTrue(record.getBoolean("is_active", 0)); + }); + } + + @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(); + Assert.assertNotNull(record); + assertion.assertRecord(record); + } + + Files.deleteIfExists(tempFile); + } +} From 7e0b6a677337d6d0d140babab9eda4ea994ac60d Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:19:55 +0200 Subject: [PATCH 44/80] Add --parquetCodec option for direct Parquet codec selection Add --parquetCodec to ParquetPipelineOptions, allowing users to specify Parquet compression directly (snappy, gzip, zstd, lz4, none) instead of relying on the --avroCodec mapping. When --parquetCodec is set, it takes precedence over --avroCodec. When not set, --avroCodec is mapped automatically (deflate* -> gzip, zstandard* -> zstd, snappy -> snappy). --- .../spotify/dbeam/parquet/JdbcParquetJob.java | 13 ++++- .../dbeam/parquet/ParquetPipelineOptions.java | 7 +++ .../dbeam/parquet/JdbcParquetJobTest.java | 47 +++++++++++++++++++ 3 files changed, 65 insertions(+), 2 deletions(-) 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 index 21b29d50..361cdd5e 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -133,11 +133,11 @@ public void prepareExport() throws Exception { } LOGGER.info("Running queries: {}", queries.toString()); + final String parquetCodec = resolveParquetCodec(pipelineOptions); final JdbcParquetArgs parquetArgs = JdbcParquetArgs.create( jdbcExportArgs.jdbcAvroOptions().jdbcConnectionConfiguration(), jdbcExportArgs.jdbcAvroOptions().fetchSize(), - mapAvroCodecToParquetCodec( - pipelineOptions.as(JdbcExportPipelineOptions.class).getAvroCodec()), + parquetCodec, 64 * 1024 * 1024, 1024 * 1024, jdbcExportArgs.jdbcAvroOptions().preCommand()); @@ -199,6 +199,15 @@ 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.equals("snappy")) { return "snappy"; 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 index 9808fd9e..219cb282 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java @@ -30,4 +30,11 @@ public interface ParquetPipelineOptions extends PipelineOptions { 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); } 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 index 387727e6..dcdcc98b 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java @@ -307,4 +307,51 @@ public void shouldRunParquetJobWithZstandardCodec() throws Exception { 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)); + } } From 4a1f7242ee152fd540ef8325e031a7499dbe38bb Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:24:32 +0200 Subject: [PATCH 45/80] Track actual bytes written in Parquet export metrics Add getBytesWritten() to ChannelOutputFile, returning the position from the underlying ChannelPositionOutputStream. The JdbcParquetWriter now exposes the actual byte count via Beam metrics after writing rows (instead of the previous hardcoded 0). Also fix BenchJdbcParquetJobTest to use create().run() instead of main() to avoid System.exit() killing the test JVM, and use unique temp directories per test to prevent directory conflicts. --- .../dbeam/parquet/ChannelOutputFile.java | 16 +++++++++++++++- .../spotify/dbeam/parquet/JdbcParquetIO.java | 9 ++++++--- .../dbeam/parquet/BenchJdbcParquetJobTest.java | 17 +++++++++-------- 3 files changed, 30 insertions(+), 12 deletions(-) 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 index 316fee96..ad101953 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java @@ -34,14 +34,28 @@ 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 { - return new ChannelPositionOutputStream(Channels.newOutputStream(channel)); + lastStream = new ChannelPositionOutputStream(Channels.newOutputStream(channel)); + return lastStream; } @Override 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 index 0528d7f6..18ece9b0 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java @@ -131,6 +131,7 @@ private static class JdbcParquetWriter extends FileBasedSink.Writer parquetWriter; private Connection connection; private JdbcParquetMetering metering; + private ChannelOutputFile channelOutputFile; JdbcParquetWriter( FileBasedSink.WriteOperation writeOperation, @@ -152,8 +153,8 @@ protected void prepareWrite(final WritableByteChannel channel) throws Exception connection = jdbcParquetArgs.jdbcConnectionConfiguration().createConnection(); final MessageType schema = MessageTypeParser.parseMessageType(schemaString); - final OutputFile outputFile = new ChannelOutputFile(channel); - parquetWriter = new ResultSetParquetWriterBuilder(outputFile, schema) + channelOutputFile = new ChannelOutputFile(channel); + parquetWriter = new ResultSetParquetWriterBuilder(channelOutputFile, schema) .withCompressionCodec(jdbcParquetArgs.getCompressionCodecName()) .withRowGroupSize(jdbcParquetArgs.rowGroupSize()) .withPageSize(jdbcParquetArgs.pageSize()) @@ -201,7 +202,9 @@ public void write(final String query) throws Exception { this.metering.incrementRecordCount(); } this.metering.exposeWriteElapsed(); - this.metering.exposeWrittenBytes(0); + if (channelOutputFile != null) { + this.metering.exposeWrittenBytes(channelOutputFile.getBytesWritten()); + } } } 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 index 86bb1ac2..3f93488e 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java @@ -48,29 +48,30 @@ public static void beforeAll() throws SQLException, ClassNotFoundException, IOEx } @Test - public void shouldRunBenchJdbcParquetJob() { - BenchJdbcParquetJob.main( + 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=" + testDir.toString(), + "--output=" + benchDir.toString(), "--avroCodec=snappy", "--executions=2" - }); - assertThat(TestHelper.listDir(testDir.toFile()), containsInAnyOrder("run_0", "run_1")); + }).run(); + assertThat(TestHelper.listDir(benchDir.toFile()), containsInAnyOrder("run_0", "run_1")); } @Test - public void shouldOutputMetricsTsvSummary() throws IOException { + 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.main( + BenchJdbcParquetJob.create( new String[] { "--targetParallelism=1", "--skipPartitionCheck", @@ -80,7 +81,7 @@ public void shouldOutputMetricsTsvSummary() throws IOException { "--output=" + benchDir.toString(), "--avroCodec=snappy", "--executions=2" - }); + }).run(); } finally { System.setOut(originalOut); } From 766dafa2080e539740ab010bc4b1e81005d68a5a Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:25:33 +0200 Subject: [PATCH 46/80] Update README to document dbeam-parquet module Add dbeam-parquet to the README: - Overview now mentions Parquet alongside Avro - Link to Parquet type conversion docs - New "dbeam-parquet package features" section - Parquet export usage example with --parquetCodec - Updated build instructions showing both shaded JARs - Updated library dependency examples for both modules --- README.md | 52 ++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 44 insertions(+), 8 deletions(-) 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 ) ``` From da70e824e27eb779e98321789ef7950095a5394e Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:32:09 +0200 Subject: [PATCH 47/80] Expose --rowGroupSize and --pageSize as CLI options Add --rowGroupSize (default 128MB) and --pageSize (default 1MB) to ParquetPipelineOptions. Previously these were hardcoded in JdbcParquetJob. Defaults follow Parquet best practices: 128MB row groups balance read performance with memory usage, 1MB pages provide good compression granularity. --- .../com/spotify/dbeam/parquet/JdbcParquetJob.java | 6 ++++-- .../dbeam/parquet/ParquetPipelineOptions.java | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) 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 index 361cdd5e..69c4a28f 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -134,12 +134,14 @@ public void prepareExport() throws Exception { 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, - 64 * 1024 * 1024, - 1024 * 1024, + parquetOptions.getRowGroupSize(), + parquetOptions.getPageSize(), jdbcExportArgs.jdbcAvroOptions().preCommand()); pipeline 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 index 219cb282..4fcb9054 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java @@ -20,6 +20,7 @@ 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; @@ -37,4 +38,18 @@ public interface ParquetPipelineOptions extends PipelineOptions { 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.Integer(134217728) + Integer getRowGroupSize(); + + void setRowGroupSize(Integer value); + + @Description("Parquet page size in bytes. Controls the granularity of encoding " + + "and compression within a column chunk.") + @Default.Integer(1048576) + Integer getPageSize(); + + void setPageSize(Integer value); } From e9a11bf29393a0668ae62e50a6f283468f79c1ad Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:33:13 +0200 Subject: [PATCH 48/80] Remove unused testDir field from BenchJdbcParquetJobTest --- .../com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 index 3f93488e..e360fc57 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java @@ -39,11 +39,9 @@ public class BenchJdbcParquetJobTest { private static final String CONNECTION_URL = "jdbc:h2:mem:testbench;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; - private static Path testDir; @BeforeClass - public static void beforeAll() throws SQLException, ClassNotFoundException, IOException { - testDir = TestHelper.createTmpDirPath("jdbc-parquet-bench-test"); + public static void beforeAll() throws SQLException, ClassNotFoundException { DbTestHelper.createFixtures(CONNECTION_URL); } From f2a20b3623443f159baa3c385a7c641fcf8423a7 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:33:55 +0200 Subject: [PATCH 49/80] Add parquet-tools validation to e2e tests Install parquet-cli (1.14.4) in CI alongside avro-tools. Update runDBeamParquetDockerCon in e2e.sh to validate Parquet output using parquet-tools: display first 5 records and print the schema. Falls back gracefully if parquet-tools is not available. --- .github/workflows/maven.yml | 6 +++++- e2e/e2e.sh | 2 ++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml index f9a1db5e..3eea9fc0 100644 --- a/.github/workflows/maven.yml +++ b/.github/workflows/maven.yml @@ -49,7 +49,7 @@ jobs: name: codecov-umbrella # optional fail_ci_if_error: false # optional (default = false) verbose: true # optional (default = false) - - name: Install avro tools + - name: Install avro and parquet tools if: matrix.java_version == '21' run: | wget https://dlcdn.apache.org/avro/avro-1.11.5/java/avro-tools-1.11.5.jar @@ -57,6 +57,10 @@ jobs: 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 == '21' diff --git a/e2e/e2e.sh b/e2e/e2e.sh index 25c6e478..cfea0e47 100755 --- a/e2e/e2e.sh +++ b/e2e/e2e.sh @@ -127,6 +127,8 @@ runDBeamParquetDockerCon() { "$@" 2>&1 | tee -a /tmp/debeam_e2e.log OUTPUT_FILE=$(ls ${OUTPUT}run_0/*.parquet | head -n 1) echo "Parquet output: $OUTPUT_FILE ($(stat -f%z "$OUTPUT_FILE" 2>/dev/null || stat -c%s "$OUTPUT_FILE") 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" } runSuite() { From c97fe339506758571319d1afa6e874e962320d09 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 2 Apr 2026 13:36:13 +0200 Subject: [PATCH 50/80] Use Parquet LIST type for SQL ARRAY columns Replace the binary string serialization of SQL arrays with proper Parquet 3-level LIST convention: optional group column_name (LIST) { repeated group list { optional binary element (STRING); } } Array elements are written individually via java.sql.Array.getArray() and serialized as STRING elements using toString(). This produces Parquet files that Spark, Hive, and other analytical tools can read as native array/list types instead of opaque binary blobs. Null array elements are supported (optional element field is simply omitted in the list group). --- .../dbeam/parquet/JdbcParquetSchema.java | 10 ++++-- .../parquet/JdbcParquetWriteSupport.java | 19 ++++++++-- .../dbeam/parquet/JdbcParquetSchemaTest.java | 8 +++-- .../parquet/PostgresJdbcParquetTest.java | 36 +++++++++++++++++++ docs/parquet-type-conversion.md | 2 +- 5 files changed, 67 insertions(+), 8 deletions(-) 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 index 777d572e..1bf16c4a 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -194,8 +194,14 @@ static Type buildParquetFieldType( return Types.optional(PrimitiveType.PrimitiveTypeName.FLOAT) .named(columnName); case ARRAY: - // For now, serialize arrays as binary (bytes mode) - return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) + // Parquet 3-level LIST convention: list > repeated list > element + // Element type defaults to STRING since JDBC array element type + // is not available from ResultSetMetaData alone. + return Types.optionalList() + .element( + Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("element")) .named(columnName); case OTHER: if (useLogicalTypes && "uuid".equals(columnTypeName)) { 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 index 5ce907df..0f92c999 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -207,10 +207,23 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int } case ARRAY: return (consumer, rs) -> { - final String val = rs.getString(column); - if (val != null && !rs.wasNull()) { + final java.sql.Array sqlArray = rs.getArray(column); + if (sqlArray != null && !rs.wasNull()) { + final Object[] items = (Object[]) sqlArray.getArray(); consumer.startField(normalizedName, fieldIndex); - consumer.addBinary(Binary.fromString(val)); + 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); + consumer.addBinary(Binary.fromString(item.toString())); + consumer.endField("element", 0); + } + consumer.endGroup(); + } + consumer.endField("list", 0); + consumer.endGroup(); consumer.endField(normalizedName, fieldIndex); } }; 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 index 0f87b4b6..d47a3a4c 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java @@ -247,10 +247,14 @@ public void shouldConvertClobSqlTypeToStringBinary() throws SQLException { } @Test - public void shouldConvertArraySqlTypeToBinary() throws SQLException { + public void shouldConvertArraySqlTypeToList() throws SQLException { final Type fieldType = buildFieldType(Types.ARRAY, false); - assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); + Assert.assertFalse(fieldType.isPrimitive()); + Assert.assertEquals( + LogicalTypeAnnotation.listType(), + fieldType.getLogicalTypeAnnotation()); + Assert.assertEquals(Type.Repetition.OPTIONAL, fieldType.getRepetition()); } @Test 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 index b4e77473..45ef8588 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java @@ -208,6 +208,42 @@ public void shouldEncodeBooleanType() throws SQLException, IOException { }); } + @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); + + // Verify schema has LIST type + Assert.assertFalse(schema.getFields().get(0).isPrimitive()); + Assert.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); + Assert.assertEquals(3, tagsList.getFieldRepetitionCount("list")); + Assert.assertEquals("rock", + tagsList.getGroup("list", 0).getString("element", 0)); + Assert.assertEquals("jazz", + tagsList.getGroup("list", 1).getString("element", 0)); + Assert.assertEquals("blues", + tagsList.getGroup("list", 2).getString("element", 0)); + }); + } + @FunctionalInterface interface RecordAssertion { void assertRecord(Group record) throws IOException; diff --git a/docs/parquet-type-conversion.md b/docs/parquet-type-conversion.md index 8670045c..b3eab451 100644 --- a/docs/parquet-type-conversion.md +++ b/docs/parquet-type-conversion.md @@ -21,7 +21,7 @@ All columns are represented as `optional` (nullable) fields in the Parquet schem | VARBINARY | BINARY | | | | LONGVARBINARY | BINARY | | | | BLOB | BINARY | | | -| ARRAY | BINARY | | Serialized as string representation | +| ARRAY | LIST (group) | LIST | 3-level list with STRING elements | | DOUBLE | DOUBLE | | | | FLOAT | FLOAT | | | | REAL | FLOAT | | | From a851dd1269f81d9286663ce4ac84912631437d54 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 08:37:28 +0200 Subject: [PATCH 51/80] Add tests for null array elements in Parquet LIST encoding Two new tests: - shouldHandleNullArrayElements: verifies that null elements within an array produce LIST groups with 0 repetitions for the "element" field (Parquet's standard null representation for optional fields) - shouldHandleNullSqlArray: verifies that a completely null SQL array produces 0 repetitions for the entire LIST field --- .../parquet/PostgresJdbcParquetTest.java | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) 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 index 45ef8588..bc8fc31a 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java @@ -244,6 +244,61 @@ public void shouldEncodeArrayAsParquetList() throws SQLException, IOException { }); } + @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); + + final Path tempFile = Files.createTempFile("parquet-null-array-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify(schema, resultSet, tempFile, record -> { + Group tagsList = record.getGroup("tags", 0); + Assert.assertEquals(3, tagsList.getFieldRepetitionCount("list")); + // First element: present + Assert.assertEquals("first", + tagsList.getGroup("list", 0).getString("element", 0)); + // Second element: null (element field has 0 repetitions) + Assert.assertEquals(0, + tagsList.getGroup("list", 1).getFieldRepetitionCount("element")); + // Third element: present + Assert.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); + + 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) + Assert.assertEquals(0, record.getFieldRepetitionCount("tags")); + }); + } + @FunctionalInterface interface RecordAssertion { void assertRecord(Group record) throws IOException; From b19efb10290b3b22174b0b5063fc082ded492878 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 08:40:03 +0200 Subject: [PATCH 52/80] Support typed array elements in Parquet LIST columns Array element types are now inferred from the JDBC column type name instead of always using STRING. For PostgreSQL, the type name prefix is stripped (e.g. _int4 -> INT32, _int8 -> INT64, _float4 -> FLOAT, _float8 -> DOUBLE, _bool -> BOOLEAN). For H2 and other databases, patterns like "INTEGER ARRAY" are parsed. Schema generation (JdbcParquetSchema): - buildArrayElementType() maps column type name to Parquet element type - resolveArrayElementTypeName() handles both PostgreSQL underscore prefix and H2-style "TYPE ARRAY" naming conventions Write support (JdbcParquetWriteSupport): - writeArrayElement() dispatches to the correct consumer.add*() method based on element type (addInteger, addLong, addFloat, addDouble, addBoolean, or addBinary for strings) This means integer[] columns produce LIST of INT32, bigint[] produces LIST of INT64, etc., which downstream tools (Spark, BigQuery, Hive) can read as native typed arrays. --- .../dbeam/parquet/JdbcParquetSchema.java | 70 +++++++++++++++++-- .../parquet/JdbcParquetWriteSupport.java | 30 +++++++- .../dbeam/parquet/JdbcParquetSchemaTest.java | 68 ++++++++++++++++-- .../parquet/PostgresJdbcParquetTest.java | 60 ++++++++++++++++ docs/parquet-type-conversion.md | 18 ++++- 5 files changed, 232 insertions(+), 14 deletions(-) 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 index 1bf16c4a..8f309862 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -194,14 +194,10 @@ static Type buildParquetFieldType( return Types.optional(PrimitiveType.PrimitiveTypeName.FLOAT) .named(columnName); case ARRAY: - // Parquet 3-level LIST convention: list > repeated list > element - // Element type defaults to STRING since JDBC array element type - // is not available from ResultSetMetaData alone. + // Parquet 3-level LIST convention with typed elements. + // Element type is inferred from columnTypeName (e.g. _int4, _text for PostgreSQL). return Types.optionalList() - .element( - Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) - .as(LogicalTypeAnnotation.stringType()) - .named("element")) + .element(buildArrayElementType(columnTypeName)) .named(columnName); case OTHER: if (useLogicalTypes && "uuid".equals(columnTypeName)) { @@ -224,6 +220,66 @@ static Type buildParquetFieldType( } } + /** + * 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"; + } + private static String normalizeForAvro(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 index 0f92c999..90251c72 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -206,6 +206,8 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int }; } case ARRAY: + final String arrayElementType = + JdbcParquetSchema.resolveArrayElementTypeName(meta.getColumnTypeName(column)); return (consumer, rs) -> { final java.sql.Array sqlArray = rs.getArray(column); if (sqlArray != null && !rs.wasNull()) { @@ -217,7 +219,7 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int consumer.startGroup(); // list element group if (item != null) { consumer.startField("element", 0); - consumer.addBinary(Binary.fromString(item.toString())); + writeArrayElement(consumer, item, arrayElementType); consumer.endField("element", 0); } consumer.endGroup(); @@ -281,4 +283,30 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int }; } } + + 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/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java index d47a3a4c..760cb13b 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java @@ -247,16 +247,74 @@ public void shouldConvertClobSqlTypeToStringBinary() throws SQLException { } @Test - public void shouldConvertArraySqlTypeToList() throws SQLException { - final Type fieldType = buildFieldType(Types.ARRAY, false); + public void shouldConvertArraySqlTypeToListWithStringElements() throws SQLException { + final Type fieldType = JdbcParquetSchema.buildParquetFieldType( + "column1", Types.ARRAY, 0, "java.sql.Array", "_text", false); Assert.assertFalse(fieldType.isPrimitive()); - Assert.assertEquals( - LogicalTypeAnnotation.listType(), - fieldType.getLogicalTypeAnnotation()); + Assert.assertEquals(LogicalTypeAnnotation.listType(), fieldType.getLogicalTypeAnnotation()); Assert.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); + Assert.assertEquals(LogicalTypeAnnotation.stringType(), + elementType.asPrimitiveType().getLogicalTypeAnnotation()); + } + + @Test + public void shouldResolveH2IntegerArrayType() { + Assert.assertEquals("int4", JdbcParquetSchema.resolveArrayElementTypeName("INTEGER ARRAY")); + } + + @Test + public void shouldResolveH2VarcharArrayType() { + Assert.assertEquals("text", JdbcParquetSchema.resolveArrayElementTypeName("VARCHAR ARRAY")); + } + + @Test + public void shouldResolveNullColumnTypeName() { + Assert.assertEquals("text", JdbcParquetSchema.resolveArrayElementTypeName(null)); + } + @Test public void shouldConvertUuidWithLogicalType() throws SQLException { final Type fieldType = JdbcParquetSchema.buildParquetFieldType( 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 index bc8fc31a..87c7cce2 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java @@ -244,6 +244,66 @@ public void shouldEncodeArrayAsParquetList() throws SQLException, IOException { }); } + @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); + + final Path tempFile = Files.createTempFile("parquet-int-array-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify(schema, resultSet, tempFile, record -> { + Group scoresList = record.getGroup("scores", 0); + Assert.assertEquals(3, scoresList.getFieldRepetitionCount("list")); + Assert.assertEquals(10, + scoresList.getGroup("list", 0).getInteger("element", 0)); + Assert.assertEquals(20, + scoresList.getGroup("list", 1).getInteger("element", 0)); + Assert.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); + + final Path tempFile = Files.createTempFile("parquet-long-array-test-", ".parquet"); + Files.delete(tempFile); + writeAndVerify(schema, resultSet, tempFile, record -> { + Group idsList = record.getGroup("ids", 0); + Assert.assertEquals(3, idsList.getFieldRepetitionCount("list")); + Assert.assertEquals(100L, + idsList.getGroup("list", 0).getLong("element", 0)); + Assert.assertEquals(200L, + idsList.getGroup("list", 1).getLong("element", 0)); + Assert.assertEquals(300L, + idsList.getGroup("list", 2).getLong("element", 0)); + }); + } + @Test public void shouldHandleNullArrayElements() throws SQLException, IOException { final ResultSetMetaData meta = Mockito.mock(ResultSetMetaData.class); diff --git a/docs/parquet-type-conversion.md b/docs/parquet-type-conversion.md index b3eab451..00cb3096 100644 --- a/docs/parquet-type-conversion.md +++ b/docs/parquet-type-conversion.md @@ -21,7 +21,7 @@ All columns are represented as `optional` (nullable) fields in the Parquet schem | VARBINARY | BINARY | | | | LONGVARBINARY | BINARY | | | | BLOB | BINARY | | | -| ARRAY | LIST (group) | LIST | 3-level list with STRING elements | +| ARRAY | LIST (group) | LIST | Typed elements (see table below) | | DOUBLE | DOUBLE | | | | FLOAT | FLOAT | | | | REAL | FLOAT | | | @@ -35,6 +35,22 @@ All columns are represented as `optional` (nullable) fields in the Parquet schem | 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: From 8032ab5f7db9b140afe9b773e37d313592dc7aa3 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 12:12:16 +0200 Subject: [PATCH 53/80] Write Avro schema in Parquet file footer metadata Parquet files now include the Avro schema (avsc JSON) in the file footer under key `parquet.avro.schema` in keyValueMetaData. This is the standard convention used by parquet-avro and enables tools like Spark, Hive, and BigQuery to read the file with full Avro type information. The Avro schema is generated via JdbcAvroSchema from dbeam-core, ensuring the avsc in the footer matches what JdbcAvroJob would produce for the same table. It is also saved as _AVRO_SCHEMA.avsc alongside _PARQUET_SCHEMA.json in the output directory. The avroSchemaJson string is threaded through the serializable chain: JdbcParquetJob -> JdbcParquetIO.createWrite() -> Sink -> WriteOperation -> Writer -> ResultSetParquetWriterBuilder -> ResultSetWriteSupport.init() -> WriteContext extraMetadata. Tests verify: - Footer metadata contains parquet.avro.schema key with correct value - Footer metadata omits key when Avro schema is not provided - Job output directory includes _AVRO_SCHEMA.avsc file --- .../spotify/dbeam/parquet/JdbcParquetIO.java | 52 +++++++--- .../spotify/dbeam/parquet/JdbcParquetJob.java | 24 ++++- .../dbeam/parquet/JdbcParquetJobTest.java | 3 + .../dbeam/parquet/JdbcParquetRecordTest.java | 96 +++++++++++++++++++ 4 files changed, 163 insertions(+), 12 deletions(-) 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 index 18ece9b0..eefa5be8 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java @@ -22,12 +22,14 @@ import static com.google.common.base.Preconditions.checkArgument; +import com.google.common.collect.ImmutableMap; 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; @@ -57,11 +59,14 @@ 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 JdbcParquetArgs jdbcParquetArgs, + final String avroSchemaJson) { final ValueProvider prefixProvider = StaticValueProvider.of( FileBasedSink.convertToFileResourceIfPossible( @@ -76,7 +81,8 @@ public static PTransform, WriteFilesResult> createWrit // Store schema as String for serialization (MessageType is not Serializable) final String schemaString = schema.toString(); final FileBasedSink sink = - new JdbcParquetSink(prefixProvider, destinations, schemaString, jdbcParquetArgs); + new JdbcParquetSink( + prefixProvider, destinations, schemaString, jdbcParquetArgs, avroSchemaJson); return WriteFiles.to(sink); } @@ -85,20 +91,23 @@ 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 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); + return new JdbcParquetWriteOperation(this, schemaString, jdbcParquetArgs, avroSchemaJson); } } @@ -108,19 +117,22 @@ private static class JdbcParquetWriteOperation 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 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); + return new JdbcParquetWriter(this, schemaString, jdbcParquetArgs, avroSchemaJson); } } @@ -128,6 +140,7 @@ private static class JdbcParquetWriter extends FileBasedSink.Writer parquetWriter; private Connection connection; private JdbcParquetMetering metering; @@ -136,10 +149,12 @@ private static class JdbcParquetWriter extends FileBasedSink.Writer writeOperation, String schemaString, - JdbcParquetArgs jdbcParquetArgs) { + JdbcParquetArgs jdbcParquetArgs, + String avroSchemaJson) { super(writeOperation, MimeTypes.BINARY); this.schemaString = schemaString; this.jdbcParquetArgs = jdbcParquetArgs; + this.avroSchemaJson = avroSchemaJson; this.metering = JdbcParquetMetering.create(); } @@ -154,7 +169,7 @@ protected void prepareWrite(final WritableByteChannel channel) throws Exception final MessageType schema = MessageTypeParser.parseMessageType(schemaString); channelOutputFile = new ChannelOutputFile(channel); - parquetWriter = new ResultSetParquetWriterBuilder(channelOutputFile, schema) + parquetWriter = new ResultSetParquetWriterBuilder(channelOutputFile, schema, avroSchemaJson) .withCompressionCodec(jdbcParquetArgs.getCompressionCodecName()) .withRowGroupSize(jdbcParquetArgs.rowGroupSize()) .withPageSize(jdbcParquetArgs.pageSize()) @@ -225,10 +240,17 @@ static class ResultSetParquetWriterBuilder extends ParquetWriter.Builder { private final MessageType schema; + private final String avroSchemaJson; ResultSetParquetWriterBuilder(OutputFile outputFile, MessageType schema) { + this(outputFile, schema, null); + } + + ResultSetParquetWriterBuilder( + OutputFile outputFile, MessageType schema, String avroSchemaJson) { super(outputFile); this.schema = schema; + this.avroSchemaJson = avroSchemaJson; } @Override @@ -238,24 +260,32 @@ protected ResultSetParquetWriterBuilder self() { @Override protected WriteSupport getWriteSupport(Configuration conf) { - return new ResultSetWriteSupport(schema); + return new ResultSetWriteSupport(schema, avroSchemaJson); } } static class ResultSetWriteSupport extends WriteSupport { private final MessageType schema; + private final String avroSchemaJson; private JdbcParquetWriteSupport writeSupport; private RecordConsumer recordConsumer; private boolean initialized = false; - ResultSetWriteSupport(MessageType schema) { + ResultSetWriteSupport(MessageType schema, String avroSchemaJson) { this.schema = schema; + this.avroSchemaJson = avroSchemaJson; } @Override public WriteContext init(Configuration configuration) { - return new WriteContext(schema, Collections.emptyMap()); + 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 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 index 69c4a28f..53bc2911 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -22,6 +22,7 @@ import com.google.common.base.Preconditions; import com.spotify.dbeam.args.JdbcExportArgs; +import com.spotify.dbeam.avro.JdbcAvroSchema; import com.spotify.dbeam.beam.BeamHelper; import com.spotify.dbeam.beam.MetricsHelper; import com.spotify.dbeam.jobs.ExceptionHandling; @@ -35,6 +36,7 @@ import java.sql.Connection; 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; @@ -116,8 +118,10 @@ public void prepareExport() throws Exception { configureVersion(); final List queries; final MessageType generatedSchema; + final Schema avroSchema; try (Connection connection = jdbcExportArgs.createConnection()) { generatedSchema = createSchema(connection); + avroSchema = generateAvroSchema(connection); queries = jdbcExportArgs.queryBuilderArgs().buildQueries(connection); final String tableName = pipelineOptions.as(DBeamPipelineOptions.class).getTable(); @@ -126,6 +130,7 @@ public void prepareExport() throws Exception { } 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)); @@ -148,7 +153,8 @@ public void prepareExport() throws Exception { .apply("JdbcQueries", Create.of(queries)) .apply( "JdbcParquetSave", - JdbcParquetIO.createWrite(output, ".parquet", generatedSchema, parquetArgs)); + JdbcParquetIO.createWrite( + output, ".parquet", generatedSchema, parquetArgs, avroSchema.toString())); } private void checkMetrics(PipelineResult pipelineResult) throws FailedValidationException { @@ -177,6 +183,22 @@ public PipelineResult runExport() throws Exception { 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()); + } + private MessageType createSchema(final Connection connection) throws Exception { final String schemaFilePath = pipelineOptions.as(ParquetPipelineOptions.class).getParquetSchemaFilePath(); 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 index dcdcc98b..fcfdf8d9 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java @@ -75,6 +75,7 @@ public void shouldRunJdbcParquetJob() throws Exception { assertThat( TestHelper.listDir(outputPath.toFile()), containsInAnyOrder( + "_AVRO_SCHEMA.avsc", "_PARQUET_SCHEMA.json", "_METRICS.json", "_SERVICE_METRICS.json", @@ -141,6 +142,7 @@ public void shouldRunJdbcParquetJobSqlFile() throws Exception { assertThat( TestHelper.listDir(outputPath.toFile()), containsInAnyOrder( + "_AVRO_SCHEMA.avsc", "_PARQUET_SCHEMA.json", "_METRICS.json", "_SERVICE_METRICS.json", @@ -181,6 +183,7 @@ public void shouldRunParquetJobWithLogicalTypes() throws Exception { assertThat( TestHelper.listDir(outputPath.toFile()), containsInAnyOrder( + "_AVRO_SCHEMA.avsc", "_PARQUET_SCHEMA.json", "_METRICS.json", "_SERVICE_METRICS.json", 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 index c8355b25..bbf32db1 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java @@ -28,13 +28,17 @@ 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.metadata.ParquetMetadata; +import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.parquet.io.OutputFile; import org.apache.parquet.schema.MessageType; import org.junit.Assert; @@ -270,4 +274,96 @@ public void shouldHandleNullValues() 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); + + 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(); + Assert.assertTrue(keyValueMetaData.containsKey(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY)); + final String actualAvroSchema = + keyValueMetaData.get(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY); + Assert.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); + + 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(); + Assert.assertFalse(keyValueMetaData.containsKey(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY)); + } + + Files.deleteIfExists(tempFile); + } } From 4f8a2db856d935b84677af5343ffc81049bfd85d Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 13:42:35 +0200 Subject: [PATCH 54/80] Remove unused ParquetMetadata import from JdbcParquetRecordTest --- .../java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java | 1 - 1 file changed, 1 deletion(-) 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 index bbf32db1..3904410a 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java @@ -37,7 +37,6 @@ import org.apache.parquet.hadoop.ParquetReader; import org.apache.parquet.hadoop.ParquetWriter; import org.apache.parquet.hadoop.example.GroupReadSupport; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.parquet.io.OutputFile; import org.apache.parquet.schema.MessageType; From 4752a83d42cbcaff11fb1cef55a0bece264e390c Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 13:43:03 +0200 Subject: [PATCH 55/80] Assert parquet.avro.schema presence in e2e tests After each Parquet e2e run, verify the file footer contains the parquet.avro.schema key using parquet-tools meta. Sanity check that the value contains a "type" field (valid Avro JSON). Fail the e2e if the key is missing. --- e2e/e2e.sh | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/e2e/e2e.sh b/e2e/e2e.sh index cfea0e47..1d57d57c 100755 --- a/e2e/e2e.sh +++ b/e2e/e2e.sh @@ -129,6 +129,17 @@ runDBeamParquetDockerCon() { echo "Parquet output: $OUTPUT_FILE ($(stat -f%z "$OUTPUT_FILE" 2>/dev/null || stat -c%s "$OUTPUT_FILE") 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() { From 12d5859a2d507d72a9a47cab6940eb481409d1b4 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 13:43:41 +0200 Subject: [PATCH 56/80] Add CHANGELOG documenting dbeam-parquet module --- CHANGELOG.md | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 CHANGELOG.md 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. From b413e95e5841ff5efa30730e516c1a9f4c9f0c98 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Tue, 7 Apr 2026 13:52:58 +0200 Subject: [PATCH 57/80] Wire --arrayMode option into Parquet export The --arrayMode option (bytes, typed_first_row, typed_postgres) from JdbcExportPipelineOptions is now respected by the Parquet export: - bytes: serializes arrays as plain BINARY (raw bytes), matching the Avro bytes mode behavior - typed_first_row / typed_postgres: uses Parquet 3-level LIST type with typed elements (current default behavior) The arrayMode is threaded through: JdbcParquetJob -> JdbcParquetArgs -> JdbcParquetIO -> WriteSupport JdbcParquetJob -> BeamJdbcParquetSchema -> JdbcParquetSchema Schema generation and write support both check arrayMode to decide between BINARY and LIST representations for ARRAY columns. --- .../dbeam/parquet/BeamJdbcParquetSchema.java | 11 ++++--- .../dbeam/parquet/JdbcParquetArgs.java | 11 +++++-- .../spotify/dbeam/parquet/JdbcParquetIO.java | 20 ++++++++--- .../spotify/dbeam/parquet/JdbcParquetJob.java | 12 ++++--- .../dbeam/parquet/JdbcParquetSchema.java | 18 +++++++--- .../parquet/JdbcParquetWriteSupport.java | 18 ++++++++-- .../parquet/BeamJdbcParquetSchemaTest.java | 2 +- .../dbeam/parquet/JdbcParquetArgsTest.java | 24 +++++++++----- .../dbeam/parquet/JdbcParquetRecordTest.java | 24 +++++++++----- .../dbeam/parquet/JdbcParquetSchemaTest.java | 21 +++++++----- .../parquet/PostgresJdbcParquetTest.java | 33 ++++++++++++------- 11 files changed, 134 insertions(+), 60 deletions(-) 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 index f4c2d613..1dab6153 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java @@ -56,10 +56,11 @@ public class BeamJdbcParquetSchema { * @throws Exception in case of failure to query database */ public static MessageType createSchema( - final Pipeline pipeline, final JdbcExportArgs args, final Connection connection) + final Pipeline pipeline, final JdbcExportArgs args, final Connection connection, + final String arrayMode) throws Exception { final long startTime = System.nanoTime(); - final MessageType generatedSchema = generateParquetSchema(args, connection); + final MessageType generatedSchema = generateParquetSchema(args, connection, arrayMode); final long elapsedTimeSchema = (System.nanoTime() - startTime) / 1000000; LOGGER.info("Elapsed time to schema {} seconds", elapsedTimeSchema / 1000.0); @@ -81,12 +82,14 @@ public static MessageType createSchema( } private static MessageType generateParquetSchema( - final JdbcExportArgs args, final Connection connection) throws SQLException { + final JdbcExportArgs args, final Connection connection, final String arrayMode) + throws SQLException { return JdbcParquetSchema.createSchemaByReadingOneRow( connection, args.queryBuilderArgs(), args.avroSchemaName(), - args.useAvroLogicalTypes()); + args.useAvroLogicalTypes(), + arrayMode); } public static Optional parseOptionalInputParquetSchemaFile(final String filename) 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 index 874658fe..bc51a347 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java @@ -50,6 +50,8 @@ public abstract class JdbcParquetArgs implements Serializable { public abstract List preCommand(); + public abstract String arrayMode(); + abstract Builder builder(); public CompressionCodecName getCompressionCodecName() { @@ -87,6 +89,8 @@ abstract static class Builder { abstract Builder setPreCommand(List preCommand); + abstract Builder setArrayMode(String arrayMode); + abstract JdbcParquetArgs build(); } @@ -96,7 +100,8 @@ public static JdbcParquetArgs create( final String parquetCodec, final int rowGroupSize, final int pageSize, - final List preCommand) { + 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"); @@ -107,6 +112,7 @@ public static JdbcParquetArgs create( .setRowGroupSize(rowGroupSize) .setPageSize(pageSize) .setPreCommand(preCommand) + .setArrayMode(arrayMode) .build(); } @@ -117,6 +123,7 @@ public static JdbcParquetArgs create(final JdbcConnectionArgs jdbcConnectionArgs "snappy", 64 * 1024 * 1024, 1024 * 1024, - Collections.emptyList()); + 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 index eefa5be8..839fcbcd 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java @@ -169,7 +169,8 @@ protected void prepareWrite(final WritableByteChannel channel) throws Exception final MessageType schema = MessageTypeParser.parseMessageType(schemaString); channelOutputFile = new ChannelOutputFile(channel); - parquetWriter = new ResultSetParquetWriterBuilder(channelOutputFile, schema, avroSchemaJson) + parquetWriter = new ResultSetParquetWriterBuilder( + channelOutputFile, schema, avroSchemaJson, jdbcParquetArgs.arrayMode()) .withCompressionCodec(jdbcParquetArgs.getCompressionCodecName()) .withRowGroupSize(jdbcParquetArgs.rowGroupSize()) .withPageSize(jdbcParquetArgs.pageSize()) @@ -241,16 +242,23 @@ static class ResultSetParquetWriterBuilder private final MessageType schema; private final String avroSchemaJson; + private final String arrayMode; ResultSetParquetWriterBuilder(OutputFile outputFile, MessageType schema) { - this(outputFile, schema, null); + 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 @@ -260,7 +268,7 @@ protected ResultSetParquetWriterBuilder self() { @Override protected WriteSupport getWriteSupport(Configuration conf) { - return new ResultSetWriteSupport(schema, avroSchemaJson); + return new ResultSetWriteSupport(schema, avroSchemaJson, arrayMode); } } @@ -268,13 +276,15 @@ 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) { + ResultSetWriteSupport(MessageType schema, String avroSchemaJson, String arrayMode) { this.schema = schema; this.avroSchemaJson = avroSchemaJson; + this.arrayMode = arrayMode; } @Override @@ -297,7 +307,7 @@ public void prepareForWrite(RecordConsumer recordConsumer) { public void write(ResultSet resultSet) { try { if (!initialized) { - this.writeSupport = JdbcParquetWriteSupport.create(resultSet, schema); + this.writeSupport = JdbcParquetWriteSupport.create(resultSet, schema, arrayMode); initialized = true; } writeSupport.writeRecord(recordConsumer, resultSet); 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 index 53bc2911..94aaf502 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -119,8 +119,9 @@ public void prepareExport() throws Exception { final List queries; final MessageType generatedSchema; final Schema avroSchema; + final String arrayMode = jdbcExportArgs.jdbcAvroOptions().arrayMode(); try (Connection connection = jdbcExportArgs.createConnection()) { - generatedSchema = createSchema(connection); + generatedSchema = createSchema(connection, arrayMode); avroSchema = generateAvroSchema(connection); queries = jdbcExportArgs.queryBuilderArgs().buildQueries(connection); @@ -147,7 +148,8 @@ public void prepareExport() throws Exception { parquetCodec, parquetOptions.getRowGroupSize(), parquetOptions.getPageSize(), - jdbcExportArgs.jdbcAvroOptions().preCommand()); + jdbcExportArgs.jdbcAvroOptions().preCommand(), + arrayMode); pipeline .apply("JdbcQueries", Create.of(queries)) @@ -199,7 +201,8 @@ private Schema generateAvroSchema(final Connection connection) throws Exception jdbcExportArgs.jdbcAvroOptions().nullableArrayItems()); } - private MessageType createSchema(final Connection connection) throws Exception { + private MessageType createSchema(final Connection connection, final String arrayMode) + throws Exception { final String schemaFilePath = pipelineOptions.as(ParquetPipelineOptions.class).getParquetSchemaFilePath(); final java.util.Optional inputSchema = @@ -207,7 +210,8 @@ private MessageType createSchema(final Connection connection) throws Exception { if (inputSchema.isPresent()) { return inputSchema.get(); } else { - return BeamJdbcParquetSchema.createSchema(this.pipeline, jdbcExportArgs, connection); + return BeamJdbcParquetSchema.createSchema( + this.pipeline, jdbcExportArgs, connection, arrayMode); } } 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 index 8f309862..5b9dd9d1 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -72,14 +72,15 @@ public static MessageType createSchemaByReadingOneRow( final Connection connection, final QueryBuilderArgs queryBuilderArgs, final Optional schemaName, - final boolean useLogicalTypes) + 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); + createParquetSchema(resultSet, schemaName, useLogicalTypes, arrayMode); LOGGER.info( "Parquet schema created successfully. useLogicalTypes={}. Generated schema: {}", useLogicalTypes, @@ -91,7 +92,8 @@ public static MessageType createSchemaByReadingOneRow( public static MessageType createParquetSchema( final ResultSet resultSet, final Optional maybeSchemaName, - final boolean useLogicalTypes) + final boolean useLogicalTypes, + final String arrayMode) throws SQLException { final ResultSetMetaData meta = resultSet.getMetaData(); final String tableName = getDatabaseTableName(meta); @@ -118,7 +120,8 @@ public static MessageType createParquetSchema( precision, columnClassName, columnTypeName, - useLogicalTypes)); + useLogicalTypes, + arrayMode)); } return new MessageType(schemaName, fields); @@ -141,7 +144,8 @@ static Type buildParquetFieldType( final int precision, final String columnClassName, final String columnTypeName, - final boolean useLogicalTypes) { + final boolean useLogicalTypes, + final String arrayMode) { switch (columnType) { case BIGINT: return Types.optional(PrimitiveType.PrimitiveTypeName.INT64) @@ -194,6 +198,10 @@ static Type buildParquetFieldType( 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() 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 index 90251c72..96c52aa2 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -83,14 +83,15 @@ public JdbcParquetWriteSupport(MessageType schema, ColumnWriter[] columnWriters, this.columnCount = columnCount; } - public static JdbcParquetWriteSupport create(ResultSet resultSet, MessageType schema) + 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); + writers[i] = computeColumnWriter(meta, i, arrayMode); } return new JdbcParquetWriteSupport(schema, writers, columnCount); @@ -111,7 +112,8 @@ public void writeRecord(RecordConsumer consumer, ResultSet resultSet) throws SQL consumer.endMessage(); } - static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int column) + static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int column, + final String arrayMode) throws SQLException { final int columnType = meta.getColumnType(column); final int fieldIndex = column - 1; @@ -206,6 +208,16 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int }; } 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) -> { 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 index 30199a02..4c404e23 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java @@ -68,7 +68,7 @@ public void shouldCreateSchemaAndExposeMetrics() throws Exception { final JdbcExportArgs args = createArgs(); try (Connection connection = args.createConnection()) { final MessageType schema = - BeamJdbcParquetSchema.createSchema(pipeline, args, connection); + BeamJdbcParquetSchema.createSchema(pipeline, args, connection, "typed_first_row"); Assert.assertNotNull(schema); Assert.assertEquals("COFFEES", schema.getName()); 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 index 4a4c149e..3ea087ef 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java @@ -59,7 +59,8 @@ public void shouldMapSnappyCodec() { @Test public void shouldMapGzipCodec() { final JdbcParquetArgs args = JdbcParquetArgs.create( - CONN_ARGS, 10000, "gzip", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + CONN_ARGS, 10000, "gzip", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), + "typed_first_row"); Assert.assertEquals(CompressionCodecName.GZIP, args.getCompressionCodecName()); } @@ -67,7 +68,8 @@ public void shouldMapGzipCodec() { @Test public void shouldMapZstdCodec() { final JdbcParquetArgs args = JdbcParquetArgs.create( - CONN_ARGS, 10000, "zstd", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + CONN_ARGS, 10000, "zstd", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), + "typed_first_row"); Assert.assertEquals(CompressionCodecName.ZSTD, args.getCompressionCodecName()); } @@ -75,7 +77,8 @@ public void shouldMapZstdCodec() { @Test public void shouldMapLz4Codec() { final JdbcParquetArgs args = JdbcParquetArgs.create( - CONN_ARGS, 10000, "lz4", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + CONN_ARGS, 10000, "lz4", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), + "typed_first_row"); Assert.assertEquals(CompressionCodecName.LZ4_RAW, args.getCompressionCodecName()); } @@ -83,7 +86,8 @@ public void shouldMapLz4Codec() { @Test public void shouldMapUncompressedCodec() { final JdbcParquetArgs args = JdbcParquetArgs.create( - CONN_ARGS, 10000, "uncompressed", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + CONN_ARGS, 10000, "uncompressed", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), + "typed_first_row"); Assert.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); } @@ -91,7 +95,8 @@ public void shouldMapUncompressedCodec() { @Test public void shouldMapNoneCodec() { final JdbcParquetArgs args = JdbcParquetArgs.create( - CONN_ARGS, 10000, "none", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + CONN_ARGS, 10000, "none", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), + "typed_first_row"); Assert.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); } @@ -99,13 +104,15 @@ public void shouldMapNoneCodec() { @Test(expected = IllegalArgumentException.class) public void shouldRejectInvalidCodec() { JdbcParquetArgs.create( - CONN_ARGS, 10000, "invalid", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + 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()); + CONN_ARGS, 50000, "snappy", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), + "typed_first_row"); Assert.assertEquals(50000, args.fetchSize()); } @@ -113,7 +120,8 @@ public void shouldAcceptCustomFetchSize() { @Test public void shouldAcceptCustomRowGroupSize() { final JdbcParquetArgs args = JdbcParquetArgs.create( - CONN_ARGS, 10000, "snappy", 128 * 1024 * 1024, 1024 * 1024, Collections.emptyList()); + CONN_ARGS, 10000, "snappy", 128 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), + "typed_first_row"); Assert.assertEquals(128 * 1024 * 1024, args.rowGroupSize()); } 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 index 3904410a..0b355425 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java @@ -58,7 +58,8 @@ public static void beforeAll() throws SQLException, ClassNotFoundException { 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); + connection, QueryBuilderArgs.create("COFFEES"), Optional.empty(), false, + "typed_first_row"); Assert.assertNotNull(schema); Assert.assertEquals("COFFEES", schema.getName()); @@ -83,7 +84,8 @@ public void shouldCreateSchemaFromDatabase() throws ClassNotFoundException, SQLE 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); + connection, QueryBuilderArgs.create("COFFEES"), Optional.empty(), true, + "typed_first_row"); Assert.assertEquals(14, schema.getFieldCount()); // CREATED and UPDATED should have timestamp logical type @@ -95,7 +97,8 @@ public void shouldCreateSchemaWithLogicalTypes() throws ClassNotFoundException, 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); + connection, QueryBuilderArgs.create("COFFEES"), Optional.of("CustomSchema"), false, + "typed_first_row"); Assert.assertEquals("CustomSchema", schema.getName()); } @@ -105,7 +108,8 @@ 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); + connection, QueryBuilderArgs.create("COFFEES"), Optional.empty(), false, + "typed_first_row"); // Write records to a temp file final Path tempFile = Files.createTempFile("parquet-test-", ".parquet"); @@ -160,7 +164,8 @@ 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); + connection, QueryBuilderArgs.create("COFFEES"), Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-values-test-", ".parquet"); Files.delete(tempFile); @@ -226,7 +231,8 @@ public void shouldHandleNullValues() // 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); + connection, QueryBuilderArgs.create("COFFEES"), Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-null-test-", ".parquet"); Files.delete(tempFile); @@ -279,7 +285,8 @@ 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); + connection, QueryBuilderArgs.create("COFFEES"), Optional.empty(), false, + "typed_first_row"); final String avroSchemaJson = "{\"type\":\"record\",\"name\":\"COFFEES\"," + "\"namespace\":\"dbeam_generated\",\"fields\":[]}"; @@ -328,7 +335,8 @@ 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); + connection, QueryBuilderArgs.create("COFFEES"), Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-no-avro-test-", ".parquet"); Files.delete(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 index 760cb13b..e28cd11f 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java @@ -93,7 +93,7 @@ public void shouldConvertIntegerSqlTypeToInt32() throws SQLException { @Test public void shouldConvertIntegerWithLongColumnClassNameToInt64() throws SQLException { final Type fieldType = JdbcParquetSchema.buildParquetFieldType( - "column1", Types.INTEGER, 0, "java.lang.Long", null, false); + "column1", Types.INTEGER, 0, "java.lang.Long", null, false, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); } @@ -169,7 +169,7 @@ public void shouldConvertBitSqlTypeWithNoPrecisionToBoolean() throws SQLExceptio @Test public void shouldConvertBitSqlTypeWithPrecision2ToBinary() throws SQLException { final Type fieldType = JdbcParquetSchema.buildParquetFieldType( - "column1", Types.BIT, 2, "foobar", null, false); + "column1", Types.BIT, 2, "foobar", null, false, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); } @@ -249,7 +249,7 @@ public void shouldConvertClobSqlTypeToStringBinary() throws SQLException { @Test public void shouldConvertArraySqlTypeToListWithStringElements() throws SQLException { final Type fieldType = JdbcParquetSchema.buildParquetFieldType( - "column1", Types.ARRAY, 0, "java.sql.Array", "_text", false); + "column1", Types.ARRAY, 0, "java.sql.Array", "_text", false, "typed_first_row"); Assert.assertFalse(fieldType.isPrimitive()); Assert.assertEquals(LogicalTypeAnnotation.listType(), fieldType.getLogicalTypeAnnotation()); @@ -318,7 +318,7 @@ public void shouldResolveNullColumnTypeName() { @Test public void shouldConvertUuidWithLogicalType() throws SQLException { final Type fieldType = JdbcParquetSchema.buildParquetFieldType( - "column1", Types.OTHER, 0, "foobar", "uuid", true); + "column1", Types.OTHER, 0, "foobar", "uuid", true, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY); Assert.assertEquals(16, fieldType.asPrimitiveType().getTypeLength()); @@ -330,7 +330,7 @@ public void shouldConvertUuidWithLogicalType() throws SQLException { @Test public void shouldConvertUuidWithoutLogicalType() throws SQLException { final Type fieldType = JdbcParquetSchema.buildParquetFieldType( - "column1", Types.OTHER, 0, "foobar", "uuid", false); + "column1", Types.OTHER, 0, "foobar", "uuid", false, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); Assert.assertEquals( @@ -369,7 +369,8 @@ public void shouldCreateParquetSchemaFromMockResultSet() throws SQLException { when(resultSet.getMetaData()).thenReturn(meta); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); Assert.assertEquals("test_table", schema.getName()); Assert.assertEquals(3, schema.getFieldCount()); @@ -397,7 +398,8 @@ public void shouldUseCustomSchemaName() throws SQLException { when(resultSet.getMetaData()).thenReturn(meta); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.of("CustomName"), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.of("CustomName"), false, + "typed_first_row"); Assert.assertEquals("CustomName", schema.getName()); } @@ -415,7 +417,8 @@ public void shouldNormalizeSpecialCharactersInColumnNames() throws SQLException when(resultSet.getMetaData()).thenReturn(meta); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); Assert.assertEquals("test_table_name", schema.getName()); Assert.assertEquals("column_name_with_spaces", schema.getFields().get(0).getName()); @@ -423,7 +426,7 @@ public void shouldNormalizeSpecialCharactersInColumnNames() throws SQLException private Type buildFieldType(final int sqlType, final boolean useLogicalTypes) { return JdbcParquetSchema.buildParquetFieldType( - "column1", sqlType, 0, "foobar", null, useLogicalTypes); + "column1", sqlType, 0, "foobar", null, useLogicalTypes, "typed_first_row"); } private void assertPrimitiveType( 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 index 87c7cce2..70cc3207 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java @@ -66,7 +66,8 @@ public void shouldEncodeUuidAsString() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-uuid-test-", ".parquet"); Files.delete(tempFile); @@ -92,7 +93,8 @@ public void shouldEncodeStringAndOtherTypes() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-string-test-", ".parquet"); Files.delete(tempFile); @@ -116,7 +118,8 @@ public void shouldEncodeTimestampAsMillis() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-ts-test-", ".parquet"); Files.delete(tempFile); @@ -142,7 +145,8 @@ public void shouldHandleNullFields() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false, true); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-null-test-", ".parquet"); Files.delete(tempFile); @@ -174,7 +178,8 @@ public void shouldEncodeMultipleNumericTypes() throws SQLException, IOException when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-numeric-test-", ".parquet"); Files.delete(tempFile); @@ -199,7 +204,8 @@ public void shouldEncodeBooleanType() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-bool-test-", ".parquet"); Files.delete(tempFile); @@ -222,7 +228,8 @@ public void shouldEncodeArrayAsParquetList() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); // Verify schema has LIST type Assert.assertFalse(schema.getFields().get(0).isPrimitive()); @@ -258,7 +265,8 @@ public void shouldEncodeIntegerArrayAsTypedList() throws SQLException, IOExcepti when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-int-array-test-", ".parquet"); Files.delete(tempFile); @@ -288,7 +296,8 @@ public void shouldEncodeBigintArrayAsTypedList() throws SQLException, IOExceptio when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-long-array-test-", ".parquet"); Files.delete(tempFile); @@ -318,7 +327,8 @@ public void shouldHandleNullArrayElements() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-null-array-test-", ".parquet"); Files.delete(tempFile); @@ -349,7 +359,8 @@ public void shouldHandleNullSqlArray() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(true); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false); + JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, + "typed_first_row"); final Path tempFile = Files.createTempFile("parquet-null-sql-array-test-", ".parquet"); Files.delete(tempFile); From 4897b01fb63026fe9ff5a32cbe74c11836d8a5a0 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 9 Apr 2026 14:23:24 +0200 Subject: [PATCH 58/80] Extract unified JdbcMetering from JdbcAvroMetering and JdbcParquetMetering The two metering classes were 100% identical except for the log prefix ("jdbcavroio" vs "jdbcparquetio"). Introduce JdbcMetering parameterized by a format name string and update all callers. JdbcAvroMetering and JdbcParquetMetering are retained as deprecated thin subclasses for backwards compatibility. --- .../com/spotify/dbeam/avro/JdbcAvroIO.java | 4 +- .../spotify/dbeam/avro/JdbcAvroMetering.java | 84 +------------ .../com/spotify/dbeam/avro/JdbcMetering.java | 110 ++++++++++++++++++ .../com/spotify/dbeam/jobs/JdbcAvroJob.java | 4 +- .../spotify/dbeam/jobs/JdbcAvroJobTest.java | 4 +- .../spotify/dbeam/parquet/JdbcParquetIO.java | 5 +- .../spotify/dbeam/parquet/JdbcParquetJob.java | 3 +- .../dbeam/parquet/JdbcParquetMetering.java | 81 ++----------- 8 files changed, 134 insertions(+), 161 deletions(-) create mode 100644 dbeam-core/src/main/java/com/spotify/dbeam/avro/JdbcMetering.java 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..40aa3b9d 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 @@ -137,7 +137,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 +147,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() { 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/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/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/test/java/com/spotify/dbeam/jobs/JdbcAvroJobTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/jobs/JdbcAvroJobTest.java index 0a8174d8..bf0b695e 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; @@ -286,7 +286,7 @@ public void shouldFailOnEmptyInput() throws IOException, ClassNotFoundException @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-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java index 839fcbcd..bc88f025 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java @@ -23,6 +23,7 @@ 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; @@ -143,7 +144,7 @@ private static class JdbcParquetWriter extends FileBasedSink.Writer parquetWriter; private Connection connection; - private JdbcParquetMetering metering; + private JdbcMetering metering; private ChannelOutputFile channelOutputFile; JdbcParquetWriter( @@ -155,7 +156,7 @@ private static class JdbcParquetWriter extends FileBasedSink.Writer 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("jdbcparquetio : Execute query took {} seconds", elapsedMs / 1000.0); - this.executeQueryElapsedMs.inc(elapsedMs); - } - - public void exposeWrittenBytes(final long count) { - this.bytesWritten.inc(count); - } } From a6c86f14f43ad022a9435126d7e4540ca844960a Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Thu, 9 Apr 2026 14:25:05 +0200 Subject: [PATCH 59/80] Extract unified BenchJdbcJob from BenchJdbcAvroJob and BenchJdbcParquetJob The two benchmark classes were ~98% identical, differing only in which job class they instantiate. Introduce BenchJdbcJob parameterized by a JobRunner functional interface and a job name string. BenchJdbcAvroJob and BenchJdbcParquetJob are retained as thin wrappers that supply the format-specific job factory, preserving their main() entry points and test compatibility. --- .../spotify/dbeam/jobs/BenchJdbcAvroJob.java | 108 +------------ .../com/spotify/dbeam/jobs/BenchJdbcJob.java | 145 ++++++++++++++++++ .../dbeam/parquet/BenchJdbcParquetJob.java | 110 +------------ 3 files changed, 160 insertions(+), 203 deletions(-) create mode 100644 dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java 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..6040b440 --- /dev/null +++ b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java @@ -0,0 +1,145 @@ +/*- + * -\-\- + * 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.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"))); + } +} 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 index 934289aa..62920841 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BenchJdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BenchJdbcParquetJob.java @@ -20,114 +20,20 @@ package com.spotify.dbeam.parquet; -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.jobs.BenchJdbcJob; import com.spotify.dbeam.jobs.ExceptionHandling; -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 BenchJdbcParquetJob { - public interface BenchJdbcParquetOptions extends PipelineOptions { - @Description("Number of benchmark executions.") - @Default.Integer(3) - int getExecutions(); - - void setExecutions(int value); - } - - private final PipelineOptions pipelineOptions; - private List> metrics = newArrayList(); - - public BenchJdbcParquetJob(final PipelineOptions pipelineOptions) { - this.pipelineOptions = pipelineOptions; - } - - public static BenchJdbcParquetJob create(final String[] cmdLineArgs) { - PipelineOptionsFactory.register(BenchJdbcParquetOptions.class); - PipelineOptions options = JdbcParquetJob.buildPipelineOptions(cmdLineArgs); - return new BenchJdbcParquetJob(options); - } - - public void run() throws Exception { - int executions = pipelineOptions.as(BenchJdbcParquetOptions.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 = - JdbcParquetJob.create(pipelineOptions, output).runExport(); - this.metrics.add(MetricsHelper.getMetrics(pipelineResult)); - } - System.out.println("Summary for BenchJdbcParquetJob"); - 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( + "BenchJdbcParquetJob", + cmdLineArgs, + JdbcParquetJob::buildPipelineOptions, + (opts, output) -> JdbcParquetJob.create(opts, output).runExport()); } public static void main(String[] cmdLineArgs) { From a2a372e8e3e04747581d6590ac5eb31ea94ce411 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 11:59:08 +0200 Subject: [PATCH 60/80] Fix UUID write to match FIXED_LEN_BYTE_ARRAY(16) schema When useLogicalTypes is enabled, the Parquet schema declares UUID columns as FIXED_LEN_BYTE_ARRAY(16) with UUID logical type, but the writer was emitting a variable-length UTF-8 string. This mismatch would produce corrupt files or write failures. Now checks the field's logical type annotation: if it is UUID, writes the 16-byte big-endian binary representation (MSB first, LSB second). Without logical types, UUID is still written as a STRING. --- .../parquet/JdbcParquetWriteSupport.java | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) 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 index 96c52aa2..a93e8fdc 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -46,6 +46,7 @@ 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; @@ -54,9 +55,12 @@ 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. @@ -91,7 +95,7 @@ public static JdbcParquetWriteSupport create(ResultSet resultSet, MessageType sc final ColumnWriter[] writers = new ColumnWriter[columnCount + 1]; for (int i = 1; i <= columnCount; i++) { - writers[i] = computeColumnWriter(meta, i, arrayMode); + writers[i] = computeColumnWriter(meta, i, arrayMode, schema.getType(i - 1)); } return new JdbcParquetWriteSupport(schema, writers, columnCount); @@ -113,7 +117,7 @@ public void writeRecord(RecordConsumer consumer, ResultSet resultSet) throws SQL } static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int column, - final String arrayMode) + final String arrayMode, final Type fieldType) throws SQLException { final int columnType = meta.getColumnType(column); final int fieldIndex = column - 1; @@ -274,6 +278,24 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int }; 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()) { From f847d9f965d53d7c6bd72a013d029193ff9c7017 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 11:59:47 +0200 Subject: [PATCH 61/80] Make Calendar thread-safe in JdbcParquetWriteSupport GregorianCalendar is not thread-safe. With --queryParallelism > 1, concurrent rs.getTimestamp(col, CALENDAR) calls could corrupt timestamp values. Use a ThreadLocal to give each thread its own Calendar instance. --- .../com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 index a93e8fdc..6ae65574 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -69,7 +69,8 @@ */ public class JdbcParquetWriteSupport { - private static final Calendar CALENDAR = new GregorianCalendar(TimeZone.getTimeZone("UTC")); + private static final ThreadLocal CALENDAR = + ThreadLocal.withInitial(() -> new GregorianCalendar(TimeZone.getTimeZone("UTC"))); @FunctionalInterface interface ColumnWriter { @@ -175,7 +176,7 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int case TIME: case TIME_WITH_TIMEZONE: return (consumer, rs) -> { - final Timestamp timestamp = rs.getTimestamp(column, CALENDAR); + final Timestamp timestamp = rs.getTimestamp(column, CALENDAR.get()); if (timestamp != null && !rs.wasNull()) { consumer.startField(normalizedName, fieldIndex); consumer.addLong(timestamp.getTime()); From 8c17ba5d7e3360d717f9e2aa5148f69ff1b533c3 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 12:00:23 +0200 Subject: [PATCH 62/80] Close InputStream in BeamJdbcParquetSchema.parseInputParquetSchemaFile The InputStream from FileSystems.open() and the Scanner were never closed, leaking a file handle. Use try-with-resources. --- .../com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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 index 1dab6153..91c03dc5 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java @@ -102,9 +102,10 @@ public static Optional parseOptionalInputParquetSchemaFile(final St public static MessageType parseInputParquetSchemaFile(final String filename) throws IOException { final MatchResult.Metadata m = FileSystems.matchSingleFileSpec(filename); - final InputStream inputStream = Channels.newInputStream(FileSystems.open(m.resourceId())); - final String schemaText = - new Scanner(inputStream, "UTF-8").useDelimiter("\\A").next(); - return MessageTypeParser.parseMessageType(schemaText); + try (InputStream inputStream = Channels.newInputStream(FileSystems.open(m.resourceId())); + Scanner scanner = new Scanner(inputStream, "UTF-8").useDelimiter("\\A")) { + final String schemaText = scanner.next(); + return MessageTypeParser.parseMessageType(schemaText); + } } } From d3ed59856eb34dcb1bd786ec0558f3b556225e91 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 12:01:08 +0200 Subject: [PATCH 63/80] Fix NPE in BenchJdbcJob when metric key is missing from map Optional.of(null) throws NPE if the metric key (e.g. KbWritePerSec) is absent from the map. Use Optional.ofNullable() instead. --- .../src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index 6040b440..ea8e9214 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java @@ -109,7 +109,7 @@ String tsvMetrics() { c -> String.format( "% 10d", - Optional.of(this.metrics.get(i).get(c)).orElse(0L))) + Optional.ofNullable(this.metrics.get(i).get(c)).orElse(0L))) .collect(tabJoining))); final List stats = columns.stream() @@ -117,7 +117,7 @@ String tsvMetrics() { c -> Stats.of( (Iterable) - this.metrics.stream().map(m -> Optional.of(m.get(c)).orElse(0L)) + this.metrics.stream().map(m -> Optional.ofNullable(m.get(c)).orElse(0L)) ::iterator)) .collect(Collectors.toList()); final Map> relevantStats = From feb47048e8c0478a42810e37339c95ef780730f4 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 12:02:45 +0200 Subject: [PATCH 64/80] Align default rowGroupSize to 128MB in JdbcParquetArgs The API default (JdbcParquetArgs.create()) was 64MB while the CLI default (ParquetPipelineOptions) was 128MB. Standardize on 128MB which matches Parquet's own default and the CLI behavior. --- .../main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java | 2 +- .../java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 index bc51a347..4521e34a 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java @@ -121,7 +121,7 @@ public static JdbcParquetArgs create(final JdbcConnectionArgs jdbcConnectionArgs jdbcConnectionArgs, 10000, "snappy", - 64 * 1024 * 1024, + 128 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), "typed_first_row"); 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 index 3ea087ef..c8a38ba2 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java @@ -44,7 +44,7 @@ public void shouldCreateWithDefaults() { Assert.assertEquals(10000, args.fetchSize()); Assert.assertEquals("snappy", args.parquetCodec()); - Assert.assertEquals(64 * 1024 * 1024, args.rowGroupSize()); + Assert.assertEquals(128 * 1024 * 1024, args.rowGroupSize()); Assert.assertEquals(1024 * 1024, args.pageSize()); Assert.assertEquals(Collections.emptyList(), args.preCommand()); } From 3ffd23387f5e69d0c4bc8fd96e27cd1f9714aa2a Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 12:04:14 +0200 Subject: [PATCH 65/80] Query database once for both Parquet and Avro schemas prepareExport() was issuing two separate queries to generate the Parquet MessageType and Avro Schema. Now executes a single sqlQueryWithLimitOne and passes the ResultSet to both JdbcParquetSchema.createParquetSchema() and JdbcAvroSchema.createAvroSchema(). When a --parquetSchemaFilePath is provided, only the Avro schema still needs a DB query since the Parquet schema comes from the file. --- .../dbeam/parquet/BeamJdbcParquetSchema.java | 21 ++++--- .../spotify/dbeam/parquet/JdbcParquetJob.java | 56 +++++++++++++------ 2 files changed, 52 insertions(+), 25 deletions(-) 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 index 91c03dc5..7422a4b7 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java @@ -61,24 +61,31 @@ public static MessageType createSchema( throws Exception { final long startTime = System.nanoTime(); final MessageType generatedSchema = generateParquetSchema(args, connection, arrayMode); - final long elapsedTimeSchema = (System.nanoTime() - startTime) / 1000000; - LOGGER.info("Elapsed time to schema {} seconds", elapsedTimeSchema / 1000.0); + final long elapsedNanos = System.nanoTime() - startTime; + exposeSchemaMetrics(pipeline, elapsedNanos); + return generatedSchema; + } - final Counter cnt = - Metrics.counter(BeamJdbcParquetSchema.class.getCanonicalName(), "schemaElapsedTimeMs"); + 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())) + Create.of(Collections.singletonList(0)) + .withType(TypeDescriptors.integers())) .apply( "ExposeSchemaCounters", MapElements.into(TypeDescriptors.integers()) .via( v -> { - cnt.inc(elapsedTimeSchema); + cnt.inc(elapsedMs); return v; })); - return generatedSchema; } private static MessageType generateParquetSchema( 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 index a60fe669..d8fe5056 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -35,6 +35,8 @@ 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; @@ -122,11 +124,42 @@ public void prepareExport() throws Exception { final Schema avroSchema; final String arrayMode = jdbcExportArgs.jdbcAvroOptions().arrayMode(); try (Connection connection = jdbcExportArgs.createConnection()) { - generatedSchema = createSchema(connection, arrayMode); - avroSchema = generateAvroSchema(connection); + final String schemaFilePath = + pipelineOptions.as(ParquetPipelineOptions.class) + .getParquetSchemaFilePath(); + final java.util.Optional inputSchema = + BeamJdbcParquetSchema + .parseOptionalInputParquetSchemaFile(schemaFilePath); + if (inputSchema.isPresent()) { + generatedSchema = inputSchema.get(); + avroSchema = generateAvroSchema(connection); + } 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(); + final String tableName = + pipelineOptions.as(DBeamPipelineOptions.class).getTable(); JobNameConfiguration.configureJobName( pipeline.getOptions(), connection.getCatalog(), tableName); } @@ -190,7 +223,8 @@ 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)); + .orElseGet(() -> String.format( + "Generate schema from JDBC ResultSet from %s", dbUrl)); return JdbcAvroSchema.createSchemaByReadingOneRow( connection, jdbcExportArgs.queryBuilderArgs(), @@ -202,20 +236,6 @@ private Schema generateAvroSchema(final Connection connection) throws Exception jdbcExportArgs.jdbcAvroOptions().nullableArrayItems()); } - private MessageType createSchema(final Connection connection, final String arrayMode) - throws Exception { - final String schemaFilePath = - pipelineOptions.as(ParquetPipelineOptions.class).getParquetSchemaFilePath(); - final java.util.Optional inputSchema = - BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(schemaFilePath); - if (inputSchema.isPresent()) { - return inputSchema.get(); - } else { - return BeamJdbcParquetSchema.createSchema( - this.pipeline, jdbcExportArgs, connection, arrayMode); - } - } - public Pipeline getPipeline() { return pipeline; } From f9d393df097a68663fa6417a578c9fda087fd87b Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 12:04:55 +0200 Subject: [PATCH 66/80] Rename normalizeForAvro to normalizeFieldName in JdbcParquetSchema The method replaces non-alphanumeric characters with underscores and is not Avro-specific. The old name was misleading in a Parquet module. --- .../java/com/spotify/dbeam/parquet/JdbcParquetSchema.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 index 5b9dd9d1..a28fcbdc 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -115,7 +115,7 @@ public static MessageType createParquetSchema( fields.add( buildParquetFieldType( - normalizeForAvro(columnName), + normalizeFieldName(columnName), columnType, precision, columnClassName, @@ -132,7 +132,7 @@ static String getDatabaseTableName(final ResultSetMetaData meta) throws SQLExcep for (int i = 1; i <= meta.getColumnCount(); i++) { String metaTableName = meta.getTableName(i); if (metaTableName != null && !metaTableName.isEmpty()) { - return normalizeForAvro(metaTableName); + return normalizeFieldName(metaTableName); } } return defaultTableName; @@ -288,7 +288,7 @@ static String resolveArrayElementTypeName(final String columnTypeName) { return "text"; } - private static String normalizeForAvro(final String input) { + private static String normalizeFieldName(final String input) { return input.replaceAll("[^A-Za-z0-9_]", "_"); } } From 6f70f95d34a450814d441b25329c8b66c33a03fa Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 12:05:34 +0200 Subject: [PATCH 67/80] Make LOGGER fields final in JdbcParquetJob and BeamJdbcParquetSchema --- .../java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java | 2 +- .../src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 index 7422a4b7..5be5d639 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java @@ -44,7 +44,7 @@ public class BeamJdbcParquetSchema { - private static Logger LOGGER = LoggerFactory.getLogger(BeamJdbcParquetSchema.class); + private static final Logger LOGGER = LoggerFactory.getLogger(BeamJdbcParquetSchema.class); /** * Generate Parquet schema by reading one row. Expose Beam metrics via a Beam PTransform. 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 index d8fe5056..f5e0f627 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -52,7 +52,7 @@ public class JdbcParquetJob { - private static Logger LOGGER = LoggerFactory.getLogger(JdbcParquetJob.class); + private static final Logger LOGGER = LoggerFactory.getLogger(JdbcParquetJob.class); private final PipelineOptions pipelineOptions; private final Pipeline pipeline; From a2f2582d27e804c8bcd5237e2dbb037ab2f4b99a Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 12:05:50 +0200 Subject: [PATCH 68/80] Document DECIMAL/NUMERIC type conversion in Parquet docs These SQL types fall through to the default STRING representation. Document this explicitly so users are not surprised by the lack of fixed-point DECIMAL Parquet type. --- docs/parquet-type-conversion.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/parquet-type-conversion.md b/docs/parquet-type-conversion.md index 00cb3096..c0a6d294 100644 --- a/docs/parquet-type-conversion.md +++ b/docs/parquet-type-conversion.md @@ -31,6 +31,7 @@ All columns are represented as `optional` (nullable) fields in the Parquet schem | 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) | | From dff137d474d1b5a655a019ac49b9f8a96e5103f3 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 15 Apr 2026 12:06:11 +0200 Subject: [PATCH 69/80] Use portable file size command in e2e.sh Replace platform-specific stat flags (macOS -f%z vs Linux -c%s) with wc -c which works on both platforms without fallback logic. --- e2e/e2e.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/e2e/e2e.sh b/e2e/e2e.sh index 1d57d57c..f50c0b93 100755 --- a/e2e/e2e.sh +++ b/e2e/e2e.sh @@ -126,7 +126,7 @@ runDBeamParquetDockerCon() { "--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 ($(stat -f%z "$OUTPUT_FILE" 2>/dev/null || stat -c%s "$OUTPUT_FILE") bytes)" + 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" From 1d642bc501537e02d05d957bdb536136765b6ef3 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Fri, 17 Apr 2026 08:37:33 +0200 Subject: [PATCH 70/80] Change rowGroupSize and pageSize from int to long Parquet's withRowGroupSize() accepts long, allowing row groups larger than 2GB for wide tables. The int type would overflow for such values. --- .../com/spotify/dbeam/parquet/JdbcParquetArgs.java | 12 ++++++------ .../com/spotify/dbeam/parquet/JdbcParquetIO.java | 2 +- .../dbeam/parquet/ParquetPipelineOptions.java | 12 ++++++------ .../spotify/dbeam/parquet/JdbcParquetArgsTest.java | 6 +++--- 4 files changed, 16 insertions(+), 16 deletions(-) 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 index 4521e34a..feca87f6 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetArgs.java @@ -44,9 +44,9 @@ public abstract class JdbcParquetArgs implements Serializable { public abstract String parquetCodec(); - public abstract int rowGroupSize(); + public abstract long rowGroupSize(); - public abstract int pageSize(); + public abstract long pageSize(); public abstract List preCommand(); @@ -83,9 +83,9 @@ abstract static class Builder { abstract Builder setParquetCodec(String parquetCodec); - abstract Builder setRowGroupSize(int rowGroupSize); + abstract Builder setRowGroupSize(long rowGroupSize); - abstract Builder setPageSize(int pageSize); + abstract Builder setPageSize(long pageSize); abstract Builder setPreCommand(List preCommand); @@ -98,8 +98,8 @@ public static JdbcParquetArgs create( final JdbcConnectionArgs jdbcConnectionArgs, final int fetchSize, final String parquetCodec, - final int rowGroupSize, - final int pageSize, + final long rowGroupSize, + final long pageSize, final List preCommand, final String arrayMode) { Preconditions.checkArgument( 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 index bc88f025..aa3b835a 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java @@ -174,7 +174,7 @@ protected void prepareWrite(final WritableByteChannel channel) throws Exception channelOutputFile, schema, avroSchemaJson, jdbcParquetArgs.arrayMode()) .withCompressionCodec(jdbcParquetArgs.getCompressionCodecName()) .withRowGroupSize(jdbcParquetArgs.rowGroupSize()) - .withPageSize(jdbcParquetArgs.pageSize()) + .withPageSize((int) jdbcParquetArgs.pageSize()) .withWriteMode(ParquetFileWriter.Mode.CREATE) .build(); LOGGER.info("jdbcparquetio : Write prepared"); 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 index 4fcb9054..0c0157b4 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java @@ -41,15 +41,15 @@ public interface ParquetPipelineOptions extends PipelineOptions { @Description("Parquet row group size in bytes. Larger values improve read performance " + "but use more memory during writes.") - @Default.Integer(134217728) - Integer getRowGroupSize(); + @Default.Long(134217728) + Long getRowGroupSize(); - void setRowGroupSize(Integer value); + void setRowGroupSize(Long value); @Description("Parquet page size in bytes. Controls the granularity of encoding " + "and compression within a column chunk.") - @Default.Integer(1048576) - Integer getPageSize(); + @Default.Long(1048576) + Long getPageSize(); - void setPageSize(Integer value); + void setPageSize(Long value); } 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 index c8a38ba2..665274a3 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java @@ -44,8 +44,8 @@ public void shouldCreateWithDefaults() { Assert.assertEquals(10000, args.fetchSize()); Assert.assertEquals("snappy", args.parquetCodec()); - Assert.assertEquals(128 * 1024 * 1024, args.rowGroupSize()); - Assert.assertEquals(1024 * 1024, args.pageSize()); + Assert.assertEquals(128L * 1024 * 1024, args.rowGroupSize()); + Assert.assertEquals(1024L * 1024, args.pageSize()); Assert.assertEquals(Collections.emptyList(), args.preCommand()); } @@ -123,6 +123,6 @@ public void shouldAcceptCustomRowGroupSize() { CONN_ARGS, 10000, "snappy", 128 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), "typed_first_row"); - Assert.assertEquals(128 * 1024 * 1024, args.rowGroupSize()); + Assert.assertEquals(128L * 1024 * 1024, args.rowGroupSize()); } } From 616db9f327a6d19dca3e5969fc3008781055a4b6 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Fri, 17 Apr 2026 08:38:15 +0200 Subject: [PATCH 71/80] Share normalizeFieldName between JdbcParquetSchema and WriteSupport Both classes independently applied the same regex to normalize column names. If one changed without the other, schema and data would mismatch. Now WriteSupport calls JdbcParquetSchema.normalizeFieldName. --- .../main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java | 2 +- .../java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 index a28fcbdc..43882399 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -288,7 +288,7 @@ static String resolveArrayElementTypeName(final String columnTypeName) { return "text"; } - private static String normalizeFieldName(final String input) { + 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 index 6ae65574..95b0e50b 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -124,7 +124,7 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int final int fieldIndex = column - 1; final String fieldName = meta.getColumnName(column).isEmpty() ? meta.getColumnLabel(column) : meta.getColumnName(column); - final String normalizedName = fieldName.replaceAll("[^A-Za-z0-9_]", "_"); + final String normalizedName = JdbcParquetSchema.normalizeFieldName(fieldName); switch (columnType) { case VARCHAR: From e8da56fc8d14e577d20003efb086bb8fbbc2061a Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Fri, 17 Apr 2026 08:38:52 +0200 Subject: [PATCH 72/80] Add null guard in mapAvroCodecToParquetCodec If getAvroCodec() returns null, the method would NPE on avroCodec.equals(). Default to snappy in that case. --- .../main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java | 3 +++ 1 file changed, 3 insertions(+) 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 index f5e0f627..1891e6f2 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -258,6 +258,9 @@ static String resolveParquetCodec(final PipelineOptions options) { } static String mapAvroCodecToParquetCodec(final String avroCodec) { + if (avroCodec == null) { + return "snappy"; + } if (avroCodec.equals("snappy")) { return "snappy"; } else if (avroCodec.startsWith("deflate")) { From 37fc3fe73a5ed4581c2113fcb1598757f550e1ea Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Fri, 17 Apr 2026 08:39:30 +0200 Subject: [PATCH 73/80] Expose schema metrics on file-override path When --parquetSchemaFilePath was provided, the Avro schema generation time was not tracked in Beam metrics. Now exposeSchemaMetrics is called on both code paths. --- .../main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java | 3 +++ 1 file changed, 3 insertions(+) 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 index 1891e6f2..ac767d01 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -131,8 +131,11 @@ public void prepareExport() throws Exception { 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(); From 3c25ebec5e9d066ba394100d7289f8433644e702 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Fri, 17 Apr 2026 08:55:37 +0200 Subject: [PATCH 74/80] Add comments noting DATE is stored as TIMESTAMP(MILLIS) SQL DATE columns are stored as epoch millis (not Parquet DATE) for consistency with the Avro export path. Document this in both the schema builder and writer to prevent future confusion. --- .../main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java | 2 +- .../java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 index 43882399..b5fc14d3 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -161,7 +161,7 @@ static Type buildParquetFieldType( .named(columnName); } case TIMESTAMP: - case DATE: + case DATE: // stored as TIMESTAMP(MILLIS) for Avro-path compatibility case TIME: case TIME_WITH_TIMEZONE: if (useLogicalTypes) { 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 index 95b0e50b..374275e9 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -172,7 +172,7 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int } }; case TIMESTAMP: - case DATE: + case DATE: // written as epoch millis for Avro-path compatibility case TIME: case TIME_WITH_TIMEZONE: return (consumer, rs) -> { From 81b5677bf03422ec852209525e08103de28d983f Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Fri, 17 Apr 2026 09:49:03 +0200 Subject: [PATCH 75/80] Remove unused BeamJdbcParquetSchema.createSchema JdbcParquetJob.prepareExport queries the DB once for both Parquet and Avro schemas and calls exposeSchemaMetrics directly, leaving createSchema and generateParquetSchema referenced only from their own test. Drop the dead helpers and the corresponding test case. --- .../dbeam/parquet/BeamJdbcParquetSchema.java | 34 -------------- .../parquet/BeamJdbcParquetSchemaTest.java | 47 ------------------- 2 files changed, 81 deletions(-) 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 index 5be5d639..3ec037bc 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java @@ -20,12 +20,9 @@ package com.spotify.dbeam.parquet; -import com.spotify.dbeam.args.JdbcExportArgs; import java.io.IOException; import java.io.InputStream; import java.nio.channels.Channels; -import java.sql.Connection; -import java.sql.SQLException; import java.util.Collections; import java.util.Optional; import java.util.Scanner; @@ -46,26 +43,6 @@ public class BeamJdbcParquetSchema { private static final Logger LOGGER = LoggerFactory.getLogger(BeamJdbcParquetSchema.class); - /** - * Generate Parquet schema by reading one row. Expose Beam metrics via a Beam PTransform. - * - * @param pipeline Beam SDK pipeline, to expose metrics - * @param args JdbcExportArgs with DBeam configuration - * @param connection JDBC connection to query input schema - * @return Parquet MessageType schema based on configuration - * @throws Exception in case of failure to query database - */ - public static MessageType createSchema( - final Pipeline pipeline, final JdbcExportArgs args, final Connection connection, - final String arrayMode) - throws Exception { - final long startTime = System.nanoTime(); - final MessageType generatedSchema = generateParquetSchema(args, connection, arrayMode); - final long elapsedNanos = System.nanoTime() - startTime; - exposeSchemaMetrics(pipeline, elapsedNanos); - return generatedSchema; - } - public static void exposeSchemaMetrics( final Pipeline pipeline, final long elapsedNanos) { final long elapsedMs = elapsedNanos / 1000000; @@ -88,17 +65,6 @@ public static void exposeSchemaMetrics( })); } - private static MessageType generateParquetSchema( - final JdbcExportArgs args, final Connection connection, final String arrayMode) - throws SQLException { - return JdbcParquetSchema.createSchemaByReadingOneRow( - connection, - args.queryBuilderArgs(), - args.avroSchemaName(), - args.useAvroLogicalTypes(), - arrayMode); - } - public static Optional parseOptionalInputParquetSchemaFile(final String filename) throws IOException { if (filename == null || filename.isEmpty()) { 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 index 4c404e23..c81a4ec3 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java @@ -20,63 +20,16 @@ package com.spotify.dbeam.parquet; -import com.spotify.dbeam.DbTestHelper; -import com.spotify.dbeam.args.JdbcAvroArgs; -import com.spotify.dbeam.args.JdbcConnectionArgs; -import com.spotify.dbeam.args.JdbcExportArgs; -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.SQLException; -import java.time.Duration; import java.util.Optional; -import org.apache.beam.sdk.testing.TestPipeline; import org.apache.parquet.schema.MessageType; import org.junit.Assert; -import org.junit.BeforeClass; -import org.junit.Rule; import org.junit.Test; public class BeamJdbcParquetSchemaTest { - private static final String CONNECTION_URL = - "jdbc:h2:mem:testbeamschema;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; - - @Rule public final transient TestPipeline pipeline = TestPipeline.create(); - - @BeforeClass - public static void beforeAll() throws SQLException, ClassNotFoundException { - DbTestHelper.createFixtures(CONNECTION_URL); - } - - private static JdbcExportArgs createArgs() throws ClassNotFoundException { - return JdbcExportArgs.create( - JdbcAvroArgs.create(JdbcConnectionArgs.create(CONNECTION_URL)), - QueryBuilderArgs.create("COFFEES"), - "dbeam_generated", - Optional.empty(), - Optional.empty(), - false, - Duration.ofMinutes(1), - Optional.empty()); - } - - @Test - public void shouldCreateSchemaAndExposeMetrics() throws Exception { - final JdbcExportArgs args = createArgs(); - try (Connection connection = args.createConnection()) { - final MessageType schema = - BeamJdbcParquetSchema.createSchema(pipeline, args, connection, "typed_first_row"); - - Assert.assertNotNull(schema); - Assert.assertEquals("COFFEES", schema.getName()); - Assert.assertEquals(14, schema.getFieldCount()); - } - pipeline.run().waitUntilFinish(); - } - @Test public void shouldParseInputSchemaFile() throws IOException { final Path schemaFile = Files.createTempFile("test-schema-", ".parquet.txt"); From 6791a13b77e423116dabcea318e919ca3cbf0b3e Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Fri, 17 Apr 2026 09:49:28 +0200 Subject: [PATCH 76/80] Read Parquet schema file via CharStreams with explicit UTF-8 Scanner with a stringly-named "UTF-8" charset and an "\\A" delimiter trick is opaque. Switch to CharStreams.toString with an InputStreamReader and StandardCharsets.UTF_8 for clearer intent and a typed charset. --- .../com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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 index 3ec037bc..f2a67423 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java @@ -20,12 +20,14 @@ 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 java.util.Scanner; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.io.FileSystems; import org.apache.beam.sdk.io.fs.MatchResult; @@ -76,9 +78,8 @@ public static Optional parseOptionalInputParquetSchemaFile(final St 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())); - Scanner scanner = new Scanner(inputStream, "UTF-8").useDelimiter("\\A")) { - final String schemaText = scanner.next(); - return MessageTypeParser.parseMessageType(schemaText); + InputStreamReader reader = new InputStreamReader(inputStream, StandardCharsets.UTF_8)) { + return MessageTypeParser.parseMessageType(CharStreams.toString(reader)); } } } From 5dad9ac14a1f403bf54e03af41862b28ac345d2c Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Fri, 17 Apr 2026 13:29:13 +0200 Subject: [PATCH 77/80] fmt --- .../com/spotify/dbeam/args/JdbcAvroArgs.java | 4 +- .../com/spotify/dbeam/avro/JdbcAvroIO.java | 9 +- .../spotify/dbeam/avro/JdbcAvroRecord.java | 6 +- .../dbeam/avro/JdbcAvroRecordConverter.java | 16 +- .../spotify/dbeam/avro/JdbcAvroSchema.java | 92 ++-- .../com/spotify/dbeam/jobs/BenchJdbcJob.java | 4 +- .../dbeam/options/ArrayHandlingMode.java | 3 +- .../dbeam/options/JdbcExportArgsFactory.java | 3 +- .../java/com/spotify/dbeam/TestHelper.java | 28 +- .../dbeam/avro/JdbcAvroRecordTest.java | 24 +- .../dbeam/avro/JdbcAvroSchemaTest.java | 34 +- .../dbeam/avro/PostgresJdbcAvroTest.java | 401 +++++++++++++----- .../spotify/dbeam/beam/BeamHelperTest.java | 7 +- .../dbeam/options/ArrayHandlingModeTest.java | 20 +- .../dbeam/parquet/BeamJdbcParquetSchema.java | 13 +- .../dbeam/parquet/ChannelOutputFile.java | 4 +- .../spotify/dbeam/parquet/JdbcParquetIO.java | 18 +- .../spotify/dbeam/parquet/JdbcParquetJob.java | 83 ++-- .../dbeam/parquet/JdbcParquetSchema.java | 64 +-- .../parquet/JdbcParquetWriteSupport.java | 44 +- .../dbeam/parquet/ParquetPipelineOptions.java | 10 +- .../spotify/dbeam/parquet/PsqlParquetJob.java | 3 +- .../parquet/BeamJdbcParquetSchemaTest.java | 12 +- .../parquet/BenchJdbcParquetJobTest.java | 43 +- .../dbeam/parquet/JdbcParquetArgsTest.java | 91 +++- .../dbeam/parquet/JdbcParquetJobTest.java | 17 +- .../dbeam/parquet/JdbcParquetRecordTest.java | 133 +++--- .../dbeam/parquet/JdbcParquetSchemaTest.java | 58 +-- .../parquet/PostgresJdbcParquetTest.java | 242 ++++++----- 29 files changed, 890 insertions(+), 596 deletions(-) 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 40aa3b9d..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; @@ -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/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/jobs/BenchJdbcJob.java b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java index ea8e9214..b3728981 100644 --- a/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java +++ b/dbeam-core/src/main/java/com/spotify/dbeam/jobs/BenchJdbcJob.java @@ -62,9 +62,7 @@ public interface BenchJdbcOptions extends PipelineOptions { private List> metrics = newArrayList(); public BenchJdbcJob( - final String jobName, - final PipelineOptions pipelineOptions, - final JobRunner jobRunner) { + final String jobName, final PipelineOptions pipelineOptions, final JobRunner jobRunner) { this.jobName = jobName; this.pipelineOptions = pipelineOptions; this.jobRunner = jobRunner; 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/avro/JdbcAvroRecordTest.java b/dbeam-core/src/test/java/com/spotify/dbeam/avro/JdbcAvroRecordTest.java index f9906186..facd94f5 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,15 +48,11 @@ 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.mockito.Mockito; -import org.postgresql.jdbc.PgArray; public class JdbcAvroRecordTest { @@ -83,7 +74,9 @@ 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()); @@ -161,7 +154,9 @@ 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( @@ -178,7 +173,9 @@ 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()); } @@ -195,8 +192,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)); 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..ce8c5be9 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 @@ -125,23 +125,27 @@ public void shouldConvertBitSqlTypeWithPrecision2ToBytes() throws SQLException { @Test public void shouldThrowOnNonSupportedTypes() throws SQLException { final ResultSet resultSet = buildMockResultSet(Types.STRUCT); - RuntimeException thrown = Assert.assertThrows(RuntimeException.class, - () -> createAvroSchemaForSingleField(resultSet, false)); + RuntimeException thrown = + Assert.assertThrows( + RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet, false)); Assert.assertEquals("STRUCT type is not supported", thrown.getMessage()); final ResultSet resultSet2 = buildMockResultSet(Types.REF); - RuntimeException thrown2 = Assert.assertThrows(RuntimeException.class, - () -> createAvroSchemaForSingleField(resultSet2, false)); + RuntimeException thrown2 = + Assert.assertThrows( + RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet2, false)); Assert.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)); + RuntimeException thrown3 = + Assert.assertThrows( + RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet3, false)); Assert.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)); + RuntimeException thrown4 = + Assert.assertThrows( + RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet4, false)); Assert.assertEquals("DATALINK type is not supported", thrown4.getMessage()); } @@ -195,8 +199,14 @@ 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..c5f4a147 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 @@ -51,8 +51,7 @@ 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,8 +82,8 @@ 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); @@ -105,18 +104,19 @@ 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())); assertGenericRecordArrayField(actualRecord, "array_field", new Utf8(uuidExpected.toString())); @@ -127,30 +127,52 @@ 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]; + 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")); assertGenericRecordArrayField(actualRecord, "array_field1", "some_text_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)); + Assert.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]; + Assert.assertArrayEquals( + expectedValue, ((java.nio.ByteBuffer) actualRecord.get("array_field")).array()); } @Test @@ -204,30 +236,82 @@ 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()); + 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")); @@ -237,8 +321,7 @@ public void shouldHandleArrayWithNullsWithoutReadingFirstRow() throws SQLExcepti Assert.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 = + Assert.assertThrows( + RuntimeException.class, () -> converter.convertResultSetIntoAvroBytes()); + Assert.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 = + Assert.assertThrows( + RuntimeException.class, () -> converter.convertResultSetIntoAvroBytes()); + Assert.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 = + 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 '_'", thrown.getMessage()); } @@ -354,16 +518,33 @@ 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)); + RuntimeException thrown = + Assert.assertThrows( + RuntimeException.class, + () -> + JdbcAvroSchema.createAvroSchema( + resultSet, + "ns", + "conn_url", + Optional.empty(), + "doc", + true, + arrayMode, + nullableArrayItems)); Assert.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..a4081f98 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 @@ -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( + "Exception should indicate GCS file matching error", + e.getMessage().contains("Error matching file spec") + && e.getMessage().contains("gs://does-not-exist-1")); } } } 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..9d10fc81 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 @@ -26,20 +26,24 @@ public class ArrayHandlingModeTest { @Test public void validValuesShouldPass() { - Assert.assertEquals(ArrayHandlingMode.Bytes, - ArrayHandlingMode.validateValue(ArrayHandlingMode.Bytes)); - Assert.assertEquals(ArrayHandlingMode.TypedMetaFromFirstRow, + Assert.assertEquals( + ArrayHandlingMode.Bytes, ArrayHandlingMode.validateValue(ArrayHandlingMode.Bytes)); + Assert.assertEquals( + ArrayHandlingMode.TypedMetaFromFirstRow, ArrayHandlingMode.validateValue(ArrayHandlingMode.TypedMetaFromFirstRow)); - Assert.assertEquals(ArrayHandlingMode.TypedMetaPostgres, + Assert.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 = + Assert.assertThrows( + RuntimeException.class, () -> ArrayHandlingMode.validateValue("invalid")); + Assert.assertEquals( + "Invalid value 'invalid' for array handling mode. Allowed values: " + + "[bytes, typed_first_row, typed_postgres]", thrown.getMessage()); } } 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 index f2a67423..8b70a4f3 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchema.java @@ -45,18 +45,15 @@ public class BeamJdbcParquetSchema { private static final Logger LOGGER = LoggerFactory.getLogger(BeamJdbcParquetSchema.class); - public static void exposeSchemaMetrics( - final Pipeline pipeline, final long elapsedNanos) { + 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"); + final Counter cnt = + Metrics.counter(BeamJdbcParquetSchema.class.getCanonicalName(), "schemaElapsedTimeMs"); pipeline .apply( "ExposeSchemaCountersSeed", - Create.of(Collections.singletonList(0)) - .withType(TypeDescriptors.integers())) + Create.of(Collections.singletonList(0)).withType(TypeDescriptors.integers())) .apply( "ExposeSchemaCounters", MapElements.into(TypeDescriptors.integers()) @@ -78,7 +75,7 @@ public static Optional parseOptionalInputParquetSchemaFile(final St 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)) { + 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/ChannelOutputFile.java b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java index ad101953..5f8bfb16 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ChannelOutputFile.java @@ -28,8 +28,8 @@ import org.apache.parquet.io.PositionOutputStream; /** - * An {@link OutputFile} implementation backed by a {@link WritableByteChannel}, - * enabling Parquet writing without Hadoop dependencies. + * An {@link OutputFile} implementation backed by a {@link WritableByteChannel}, enabling Parquet + * writing without Hadoop dependencies. */ public class ChannelOutputFile implements OutputFile { 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 index aa3b835a..cae8cc03 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetIO.java @@ -170,13 +170,14 @@ protected void prepareWrite(final WritableByteChannel channel) throws Exception 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(); + 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"); } @@ -190,8 +191,7 @@ private ResultSet executeQuery(final String query) throws Exception { jdbcParquetArgs.statementPreparator().setParameters(statement); } - if (jdbcParquetArgs.preCommand() != null - && !jdbcParquetArgs.preCommand().isEmpty()) { + if (jdbcParquetArgs.preCommand() != null && !jdbcParquetArgs.preCommand().isEmpty()) { final Statement stmt = connection.createStatement(); for (String command : jdbcParquetArgs.preCommand()) { stmt.execute(command); 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 index ac767d01..084e8c3f 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetJob.java @@ -125,44 +125,49 @@ public void prepareExport() throws Exception { final String arrayMode = jdbcExportArgs.jdbcAvroOptions().arrayMode(); try (Connection connection = jdbcExportArgs.createConnection()) { final String schemaFilePath = - pipelineOptions.as(ParquetPipelineOptions.class) - .getParquetSchemaFilePath(); + pipelineOptions.as(ParquetPipelineOptions.class).getParquetSchemaFilePath(); final java.util.Optional inputSchema = - BeamJdbcParquetSchema - .parseOptionalInputParquetSchemaFile(schemaFilePath); + BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(schemaFilePath); if (inputSchema.isPresent()) { final long startTime = System.nanoTime(); generatedSchema = inputSchema.get(); avroSchema = generateAvroSchema(connection); - BeamJdbcParquetSchema.exposeSchemaMetrics( - pipeline, System.nanoTime() - startTime); + 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()); + final ResultSet rs = + stmt.executeQuery(jdbcExportArgs.queryBuilderArgs().sqlQueryWithLimitOne()); rs.next(); - generatedSchema = JdbcParquetSchema - .createParquetSchema(rs, jdbcExportArgs.avroSchemaName(), - jdbcExportArgs.useAvroLogicalTypes(), arrayMode); + 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()); + 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); + BeamJdbcParquetSchema.exposeSchemaMetrics(pipeline, System.nanoTime() - startTime); } queries = jdbcExportArgs.queryBuilderArgs().buildQueries(connection); - final String tableName = - pipelineOptions.as(DBeamPipelineOptions.class).getTable(); + final String tableName = pipelineOptions.as(DBeamPipelineOptions.class).getTable(); JobNameConfiguration.configureJobName( pipeline.getOptions(), connection.getCatalog(), tableName); } @@ -177,16 +182,16 @@ public void prepareExport() throws Exception { 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); + 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)) @@ -201,8 +206,7 @@ private void checkMetrics(PipelineResult pipelineResult) throws FailedValidation if (!this.dataOnly) { BeamHelper.saveMetrics(metrics, output); } - final Long recordCount = - metrics.getOrDefault(JdbcMetering.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( @@ -225,9 +229,9 @@ public PipelineResult runExport() throws Exception { 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)); + jdbcExportArgs + .avroDoc() + .orElseGet(() -> String.format("Generate schema from JDBC ResultSet from %s", dbUrl)); return JdbcAvroSchema.createSchemaByReadingOneRow( connection, jdbcExportArgs.queryBuilderArgs(), @@ -256,8 +260,7 @@ static String resolveParquetCodec(final PipelineOptions options) { if (parquetCodec != null && !parquetCodec.isEmpty()) { return parquetCodec; } - return mapAvroCodecToParquetCodec( - options.as(JdbcExportPipelineOptions.class).getAvroCodec()); + return mapAvroCodecToParquetCodec(options.as(JdbcExportPipelineOptions.class).getAvroCodec()); } static String mapAvroCodecToParquetCodec(final String avroCodec) { 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 index b5fc14d3..92a03184 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetSchema.java @@ -148,17 +148,14 @@ static Type buildParquetFieldType( final String arrayMode) { switch (columnType) { case BIGINT: - return Types.optional(PrimitiveType.PrimitiveTypeName.INT64) - .named(columnName); + 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); + return Types.optional(PrimitiveType.PrimitiveTypeName.INT64).named(columnName); } else { - return Types.optional(PrimitiveType.PrimitiveTypeName.INT32) - .named(columnName); + return Types.optional(PrimitiveType.PrimitiveTypeName.INT32).named(columnName); } case TIMESTAMP: case DATE: // stored as TIMESTAMP(MILLIS) for Avro-path compatibility @@ -166,41 +163,32 @@ static Type buildParquetFieldType( case TIME_WITH_TIMEZONE: if (useLogicalTypes) { return Types.optional(PrimitiveType.PrimitiveTypeName.INT64) - .as(LogicalTypeAnnotation.timestampType(true, - LogicalTypeAnnotation.TimeUnit.MILLIS)) + .as(LogicalTypeAnnotation.timestampType(true, LogicalTypeAnnotation.TimeUnit.MILLIS)) .named(columnName); } else { - return Types.optional(PrimitiveType.PrimitiveTypeName.INT64) - .named(columnName); + return Types.optional(PrimitiveType.PrimitiveTypeName.INT64).named(columnName); } case BOOLEAN: - return Types.optional(PrimitiveType.PrimitiveTypeName.BOOLEAN) - .named(columnName); + return Types.optional(PrimitiveType.PrimitiveTypeName.BOOLEAN).named(columnName); case BIT: if (precision <= 1) { - return Types.optional(PrimitiveType.PrimitiveTypeName.BOOLEAN) - .named(columnName); + return Types.optional(PrimitiveType.PrimitiveTypeName.BOOLEAN).named(columnName); } else { - return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) - .named(columnName); + return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY).named(columnName); } case BINARY: case VARBINARY: case LONGVARBINARY: case BLOB: - return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) - .named(columnName); + return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY).named(columnName); case DOUBLE: - return Types.optional(PrimitiveType.PrimitiveTypeName.DOUBLE) - .named(columnName); + return Types.optional(PrimitiveType.PrimitiveTypeName.DOUBLE).named(columnName); case FLOAT: case REAL: - return Types.optional(PrimitiveType.PrimitiveTypeName.FLOAT) - .named(columnName); + return Types.optional(PrimitiveType.PrimitiveTypeName.FLOAT).named(columnName); case ARRAY: if ("bytes".equals(arrayMode)) { - return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) - .named(columnName); + 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). @@ -214,7 +202,7 @@ static Type buildParquetFieldType( .as(LogicalTypeAnnotation.uuidType()) .named(columnName); } - // fall through to string + // fall through to string case VARCHAR: case CHAR: case CLOB: @@ -229,9 +217,9 @@ static Type buildParquetFieldType( } /** - * 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. + * 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); @@ -239,20 +227,15 @@ static Type buildArrayElementType(final String columnTypeName) { case "int": case "int4": case "int2": - return Types.optional(PrimitiveType.PrimitiveTypeName.INT32) - .named("element"); + return Types.optional(PrimitiveType.PrimitiveTypeName.INT32).named("element"); case "int8": - return Types.optional(PrimitiveType.PrimitiveTypeName.INT64) - .named("element"); + return Types.optional(PrimitiveType.PrimitiveTypeName.INT64).named("element"); case "float4": - return Types.optional(PrimitiveType.PrimitiveTypeName.FLOAT) - .named("element"); + return Types.optional(PrimitiveType.PrimitiveTypeName.FLOAT).named("element"); case "float8": - return Types.optional(PrimitiveType.PrimitiveTypeName.DOUBLE) - .named("element"); + return Types.optional(PrimitiveType.PrimitiveTypeName.DOUBLE).named("element"); case "bool": - return Types.optional(PrimitiveType.PrimitiveTypeName.BOOLEAN) - .named("element"); + return Types.optional(PrimitiveType.PrimitiveTypeName.BOOLEAN).named("element"); default: return Types.optional(PrimitiveType.PrimitiveTypeName.BINARY) .as(LogicalTypeAnnotation.stringType()) @@ -261,9 +244,8 @@ static Type buildArrayElementType(final String columnTypeName) { } /** - * 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. + * 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) { 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 index 374275e9..bbf0da05 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/JdbcParquetWriteSupport.java @@ -63,9 +63,9 @@ 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). + * 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 { @@ -81,16 +81,15 @@ interface ColumnWriter { private final ColumnWriter[] columnWriters; private final int columnCount; - public JdbcParquetWriteSupport(MessageType schema, ColumnWriter[] columnWriters, - 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 { + 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]; @@ -106,9 +105,7 @@ public MessageType getSchema() { return schema; } - /** - * Write the current row of the ResultSet to the RecordConsumer. - */ + /** 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++) { @@ -117,13 +114,15 @@ public void writeRecord(RecordConsumer consumer, ResultSet resultSet) throws SQL consumer.endMessage(); } - static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int column, - final String arrayMode, final Type fieldType) + 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 fieldName = + meta.getColumnName(column).isEmpty() + ? meta.getColumnLabel(column) + : meta.getColumnName(column); final String normalizedName = JdbcParquetSchema.normalizeFieldName(fieldName); switch (columnType) { @@ -279,15 +278,15 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int }; case OTHER: if (Objects.equals(meta.getColumnTypeName(column), "uuid")) { - final boolean isUuidLogicalType = fieldType.getLogicalTypeAnnotation() != null - && fieldType.getLogicalTypeAnnotation() - .equals(LogicalTypeAnnotation.uuidType()); + 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 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()); @@ -306,7 +305,7 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int } }; } - // fall through to string + // fall through to string default: return (consumer, rs) -> { final String val = rs.getString(column); @@ -319,8 +318,7 @@ static ColumnWriter computeColumnWriter(final ResultSetMetaData meta, final int } } - static void writeArrayElement( - RecordConsumer consumer, Object item, String elementType) { + static void writeArrayElement(RecordConsumer consumer, Object item, String elementType) { switch (elementType) { case "int": case "int4": 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 index 0c0157b4..8a176537 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/ParquetPipelineOptions.java @@ -39,15 +39,17 @@ public interface ParquetPipelineOptions extends PipelineOptions { void setParquetCodec(String value); - @Description("Parquet row group size in bytes. Larger values improve read performance " - + "but use more memory during writes.") + @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.") + @Description( + "Parquet page size in bytes. Controls the granularity of encoding " + + "and compression within a column chunk.") @Default.Long(1048576) Long getPageSize(); 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 index 6cf531b0..964854ad 100644 --- a/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/PsqlParquetJob.java +++ b/dbeam-parquet/src/main/java/com/spotify/dbeam/parquet/PsqlParquetJob.java @@ -29,8 +29,7 @@ public class PsqlParquetJob { private final JdbcParquetJob job; private final PsqlReplicationCheck psqlReplicationCheck; - public PsqlParquetJob( - final JdbcParquetJob job, final PsqlReplicationCheck psqlReplicationCheck) { + public PsqlParquetJob(final JdbcParquetJob job, final PsqlReplicationCheck psqlReplicationCheck) { this.job = job; this.psqlReplicationCheck = psqlReplicationCheck; } 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 index c81a4ec3..bec6b51f 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java @@ -33,11 +33,13 @@ 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()); + Files.write( + schemaFile, + ("message test_table {\n" + + " optional int64 id;\n" + + " optional binary name (STRING);\n" + + "}") + .getBytes()); final Optional schema = BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(schemaFile.toString()); 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 index e360fc57..0e5ef26c 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java @@ -27,7 +27,6 @@ import com.spotify.dbeam.DbTestHelper; import com.spotify.dbeam.TestHelper; import java.io.ByteArrayOutputStream; -import java.io.IOException; import java.io.PrintStream; import java.nio.file.Path; import java.sql.SQLException; @@ -49,16 +48,17 @@ public static void beforeAll() throws SQLException, ClassNotFoundException { 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(); + 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")); } @@ -70,16 +70,17 @@ public void shouldOutputMetricsTsvSummary() throws Exception { 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(); + new String[] { + "--targetParallelism=1", + "--skipPartitionCheck", + "--connectionUrl=" + CONNECTION_URL, + "--username=", + "--table=COFFEES", + "--output=" + benchDir.toString(), + "--avroCodec=snappy", + "--executions=2" + }) + .run(); } finally { System.setOut(originalOut); } 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 index 665274a3..3b31de8d 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java @@ -58,45 +58,75 @@ public void shouldMapSnappyCodec() { @Test public void shouldMapGzipCodec() { - final JdbcParquetArgs args = JdbcParquetArgs.create( - CONN_ARGS, 10000, "gzip", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), - "typed_first_row"); + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "gzip", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); Assert.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"); + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "zstd", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); Assert.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"); + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "lz4", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); Assert.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"); + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "uncompressed", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); Assert.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"); + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "none", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); Assert.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); } @@ -104,24 +134,41 @@ public void shouldMapNoneCodec() { @Test(expected = IllegalArgumentException.class) public void shouldRejectInvalidCodec() { JdbcParquetArgs.create( - CONN_ARGS, 10000, "invalid", 64 * 1024 * 1024, 1024 * 1024, Collections.emptyList(), + 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"); + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 50000, + "snappy", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); Assert.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"); + final JdbcParquetArgs args = + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "snappy", + 128 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row"); Assert.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 index fcfdf8d9..c979b09e 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java @@ -116,8 +116,7 @@ public void shouldRunJdbcParquetJobDataOnly() throws Exception { .runExport(); assertThat( - TestHelper.listDir(outputPath.toFile()), - containsInAnyOrder("part-00000-of-00001.parquet")); + TestHelper.listDir(outputPath.toFile()), containsInAnyOrder("part-00000-of-00001.parquet")); } @Test @@ -223,12 +222,14 @@ public void shouldRunParquetJobWithMinRows() throws Exception { 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()); + 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[] { 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 index 0b355425..02708677 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java @@ -57,9 +57,13 @@ public static void beforeAll() throws SQLException, ClassNotFoundException { @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"); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.empty(), + false, + "typed_first_row"); Assert.assertNotNull(schema); Assert.assertEquals("COFFEES", schema.getName()); @@ -83,9 +87,13 @@ public void shouldCreateSchemaFromDatabase() throws ClassNotFoundException, SQLE @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"); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.empty(), + true, + "typed_first_row"); Assert.assertEquals(14, schema.getFieldCount()); // CREATED and UPDATED should have timestamp logical type @@ -96,9 +104,13 @@ public void shouldCreateSchemaWithLogicalTypes() throws ClassNotFoundException, @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"); + final MessageType schema = + JdbcParquetSchema.createSchemaByReadingOneRow( + connection, + QueryBuilderArgs.create("COFFEES"), + Optional.of("CustomSchema"), + false, + "typed_first_row"); Assert.assertEquals("CustomSchema", schema.getName()); } @@ -107,9 +119,13 @@ public void shouldCreateSchemaWithCustomName() throws ClassNotFoundException, SQ 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"); + 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"); @@ -135,12 +151,9 @@ public void shouldWriteAndReadBackParquetRecords() // Read back and verify final Configuration conf = new Configuration(); - final org.apache.hadoop.fs.Path hadoopPath = - new org.apache.hadoop.fs.Path(tempFile.toUri()); + 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()) { + ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { Group record1 = reader.read(); Assert.assertNotNull(record1); // Verify first record has expected fields @@ -163,15 +176,19 @@ public void shouldWriteAndReadBackParquetRecords() 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 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 ResultSet rs = + connection.createStatement().executeQuery("SELECT * FROM COFFEES ORDER BY COF_NAME"); final OutputFile outputFile = new ChannelOutputFile( @@ -190,12 +207,9 @@ public void shouldWriteCorrectFieldValues() } final Configuration conf = new Configuration(); - final org.apache.hadoop.fs.Path hadoopPath = - new org.apache.hadoop.fs.Path(tempFile.toUri()); + 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()) { + ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { Group record = reader.read(); Assert.assertNotNull(record); @@ -226,19 +240,21 @@ public void shouldWriteCorrectFieldValues() } @Test - public void shouldHandleNullValues() - throws ClassNotFoundException, SQLException, IOException { + 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 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 ResultSet rs = connection.createStatement().executeQuery("SELECT * FROM COFFEES LIMIT 1"); final OutputFile outputFile = new ChannelOutputFile( @@ -257,12 +273,9 @@ public void shouldHandleNullValues() } final Configuration conf = new Configuration(); - final org.apache.hadoop.fs.Path hadoopPath = - new org.apache.hadoop.fs.Path(tempFile.toUri()); + 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()) { + ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { Group record = reader.read(); Assert.assertNotNull(record); @@ -284,18 +297,22 @@ public void shouldHandleNullValues() 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 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 ResultSet rs = connection.createStatement().executeQuery("SELECT * FROM COFFEES LIMIT 1"); final OutputFile outputFile = new ChannelOutputFile( @@ -315,15 +332,13 @@ public void shouldWriteAvroSchemaInParquetFooter() // 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()); + 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(); Assert.assertTrue(keyValueMetaData.containsKey(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY)); - final String actualAvroSchema = - keyValueMetaData.get(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY); + final String actualAvroSchema = keyValueMetaData.get(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY); Assert.assertEquals(avroSchemaJson, actualAvroSchema); } @@ -334,15 +349,18 @@ public void shouldWriteAvroSchemaInParquetFooter() 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 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 ResultSet rs = connection.createStatement().executeQuery("SELECT * FROM COFFEES LIMIT 1"); final OutputFile outputFile = new ChannelOutputFile( @@ -362,8 +380,7 @@ public void shouldOmitAvroSchemaWhenNotProvided() } final Configuration conf = new Configuration(); - final org.apache.hadoop.fs.Path hadoopPath = - new org.apache.hadoop.fs.Path(tempFile.toUri()); + 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 = 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 index e28cd11f..1f6c2ad5 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java @@ -92,8 +92,9 @@ public void shouldConvertIntegerSqlTypeToInt32() throws SQLException { @Test public void shouldConvertIntegerWithLongColumnClassNameToInt64() throws SQLException { - final Type fieldType = JdbcParquetSchema.buildParquetFieldType( - "column1", Types.INTEGER, 0, "java.lang.Long", null, false, "typed_first_row"); + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.INTEGER, 0, "java.lang.Long", null, false, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); } @@ -168,8 +169,9 @@ public void shouldConvertBitSqlTypeWithNoPrecisionToBoolean() throws SQLExceptio @Test public void shouldConvertBitSqlTypeWithPrecision2ToBinary() throws SQLException { - final Type fieldType = JdbcParquetSchema.buildParquetFieldType( - "column1", Types.BIT, 2, "foobar", null, false, "typed_first_row"); + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.BIT, 2, "foobar", null, false, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); } @@ -222,8 +224,7 @@ public void shouldConvertVarcharSqlTypeToStringBinary() throws SQLException { assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); Assert.assertEquals( - LogicalTypeAnnotation.stringType(), - fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @Test @@ -232,8 +233,7 @@ public void shouldConvertCharSqlTypeToStringBinary() throws SQLException { assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); Assert.assertEquals( - LogicalTypeAnnotation.stringType(), - fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @Test @@ -242,14 +242,14 @@ public void shouldConvertClobSqlTypeToStringBinary() throws SQLException { assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); Assert.assertEquals( - LogicalTypeAnnotation.stringType(), - fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + 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"); + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.ARRAY, 0, "java.sql.Array", "_text", false, "typed_first_row"); Assert.assertFalse(fieldType.isPrimitive()); Assert.assertEquals(LogicalTypeAnnotation.listType(), fieldType.getLogicalTypeAnnotation()); @@ -296,7 +296,8 @@ public void shouldConvertTextArrayToListOfString() { final Type elementType = JdbcParquetSchema.buildArrayElementType("_text"); assertPrimitiveType(elementType, PrimitiveType.PrimitiveTypeName.BINARY); - Assert.assertEquals(LogicalTypeAnnotation.stringType(), + Assert.assertEquals( + LogicalTypeAnnotation.stringType(), elementType.asPrimitiveType().getLogicalTypeAnnotation()); } @@ -317,25 +318,25 @@ public void shouldResolveNullColumnTypeName() { @Test public void shouldConvertUuidWithLogicalType() throws SQLException { - final Type fieldType = JdbcParquetSchema.buildParquetFieldType( - "column1", Types.OTHER, 0, "foobar", "uuid", true, "typed_first_row"); + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.OTHER, 0, "foobar", "uuid", true, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY); Assert.assertEquals(16, fieldType.asPrimitiveType().getTypeLength()); Assert.assertEquals( - LogicalTypeAnnotation.uuidType(), - fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + 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"); + final Type fieldType = + JdbcParquetSchema.buildParquetFieldType( + "column1", Types.OTHER, 0, "foobar", "uuid", false, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); Assert.assertEquals( - LogicalTypeAnnotation.stringType(), - fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @Test @@ -344,8 +345,7 @@ public void shouldDefaultConversionToStringBinary() throws SQLException { assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); Assert.assertEquals( - LogicalTypeAnnotation.stringType(), - fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @Test @@ -369,8 +369,8 @@ public void shouldCreateParquetSchemaFromMockResultSet() throws SQLException { when(resultSet.getMetaData()).thenReturn(meta); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); Assert.assertEquals("test_table", schema.getName()); Assert.assertEquals(3, schema.getFieldCount()); @@ -398,8 +398,8 @@ public void shouldUseCustomSchemaName() throws SQLException { when(resultSet.getMetaData()).thenReturn(meta); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.of("CustomName"), false, - "typed_first_row"); + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.of("CustomName"), false, "typed_first_row"); Assert.assertEquals("CustomName", schema.getName()); } @@ -417,8 +417,8 @@ public void shouldNormalizeSpecialCharactersInColumnNames() throws SQLException when(resultSet.getMetaData()).thenReturn(meta); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); Assert.assertEquals("test_table_name", schema.getName()); Assert.assertEquals("column_name_with_spaces", schema.getFields().get(0).getName()); 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 index 70cc3207..56d947c9 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java @@ -66,15 +66,19 @@ public void shouldEncodeUuidAsString() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + 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); - Assert.assertEquals(uuidExpected.toString(), actualUuid); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + String actualUuid = record.getString("uuid_field", 0); + Assert.assertEquals(uuidExpected.toString(), actualUuid); + }); } @Test @@ -93,15 +97,19 @@ public void shouldEncodeStringAndOtherTypes() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + 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 -> { - Assert.assertEquals("some_text_42", record.getString("text_field", 0)); - Assert.assertEquals("some_other_42", record.getString("other_field", 0)); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Assert.assertEquals("some_text_42", record.getString("text_field", 0)); + Assert.assertEquals("some_other_42", record.getString("other_field", 0)); + }); } @Test @@ -118,15 +126,19 @@ public void shouldEncodeTimestampAsMillis() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + 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); - Assert.assertEquals(1488300933000L, actualTs); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + long actualTs = record.getLong("ts_field", 0); + Assert.assertEquals(1488300933000L, actualTs); + }); } @Test @@ -145,15 +157,19 @@ public void shouldHandleNullFields() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false, true); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + 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 -> { - Assert.assertEquals("alice", record.getString("name", 0)); - Assert.assertEquals(0, record.getFieldRepetitionCount("age")); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Assert.assertEquals("alice", record.getString("name", 0)); + Assert.assertEquals(0, record.getFieldRepetitionCount("age")); + }); } @Test @@ -167,8 +183,7 @@ public void shouldEncodeMultipleNumericTypes() throws SQLException, IOException 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"); + TestHelper.mockResultSetMeta(meta, 4, Types.DOUBLE, "double_col", "java.lang.Double", "float8"); final ResultSet resultSet = buildMockResultSet(meta); when(resultSet.getInt(1)).thenReturn(42); @@ -178,17 +193,21 @@ public void shouldEncodeMultipleNumericTypes() throws SQLException, IOException when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + 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 -> { - Assert.assertEquals(42, record.getInteger("int_col", 0)); - Assert.assertEquals(9999999999L, record.getLong("long_col", 0)); - Assert.assertEquals(3.14f, record.getFloat("float_col", 0), 0.001f); - Assert.assertEquals(2.71828, record.getDouble("double_col", 0), 0.00001); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Assert.assertEquals(42, record.getInteger("int_col", 0)); + Assert.assertEquals(9999999999L, record.getLong("long_col", 0)); + Assert.assertEquals(3.14f, record.getFloat("float_col", 0), 0.001f); + Assert.assertEquals(2.71828, record.getDouble("double_col", 0), 0.00001); + }); } @Test @@ -196,22 +215,25 @@ 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"); + 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"); + 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 -> { - Assert.assertTrue(record.getBoolean("is_active", 0)); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Assert.assertTrue(record.getBoolean("is_active", 0)); + }); } @Test @@ -228,8 +250,8 @@ public void shouldEncodeArrayAsParquetList() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + JdbcParquetSchema.createParquetSchema( + resultSet, Optional.empty(), false, "typed_first_row"); // Verify schema has LIST type Assert.assertFalse(schema.getFields().get(0).isPrimitive()); @@ -239,16 +261,17 @@ public void shouldEncodeArrayAsParquetList() throws SQLException, IOException { final Path tempFile = Files.createTempFile("parquet-array-test-", ".parquet"); Files.delete(tempFile); - writeAndVerify(schema, resultSet, tempFile, record -> { - Group tagsList = record.getGroup("tags", 0); - Assert.assertEquals(3, tagsList.getFieldRepetitionCount("list")); - Assert.assertEquals("rock", - tagsList.getGroup("list", 0).getString("element", 0)); - Assert.assertEquals("jazz", - tagsList.getGroup("list", 1).getString("element", 0)); - Assert.assertEquals("blues", - tagsList.getGroup("list", 2).getString("element", 0)); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Group tagsList = record.getGroup("tags", 0); + Assert.assertEquals(3, tagsList.getFieldRepetitionCount("list")); + Assert.assertEquals("rock", tagsList.getGroup("list", 0).getString("element", 0)); + Assert.assertEquals("jazz", tagsList.getGroup("list", 1).getString("element", 0)); + Assert.assertEquals("blues", tagsList.getGroup("list", 2).getString("element", 0)); + }); } @Test @@ -265,21 +288,22 @@ public void shouldEncodeIntegerArrayAsTypedList() throws SQLException, IOExcepti when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + 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); - Assert.assertEquals(3, scoresList.getFieldRepetitionCount("list")); - Assert.assertEquals(10, - scoresList.getGroup("list", 0).getInteger("element", 0)); - Assert.assertEquals(20, - scoresList.getGroup("list", 1).getInteger("element", 0)); - Assert.assertEquals(30, - scoresList.getGroup("list", 2).getInteger("element", 0)); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Group scoresList = record.getGroup("scores", 0); + Assert.assertEquals(3, scoresList.getFieldRepetitionCount("list")); + Assert.assertEquals(10, scoresList.getGroup("list", 0).getInteger("element", 0)); + Assert.assertEquals(20, scoresList.getGroup("list", 1).getInteger("element", 0)); + Assert.assertEquals(30, scoresList.getGroup("list", 2).getInteger("element", 0)); + }); } @Test @@ -296,21 +320,22 @@ public void shouldEncodeBigintArrayAsTypedList() throws SQLException, IOExceptio when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + 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); - Assert.assertEquals(3, idsList.getFieldRepetitionCount("list")); - Assert.assertEquals(100L, - idsList.getGroup("list", 0).getLong("element", 0)); - Assert.assertEquals(200L, - idsList.getGroup("list", 1).getLong("element", 0)); - Assert.assertEquals(300L, - idsList.getGroup("list", 2).getLong("element", 0)); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Group idsList = record.getGroup("ids", 0); + Assert.assertEquals(3, idsList.getFieldRepetitionCount("list")); + Assert.assertEquals(100L, idsList.getGroup("list", 0).getLong("element", 0)); + Assert.assertEquals(200L, idsList.getGroup("list", 1).getLong("element", 0)); + Assert.assertEquals(300L, idsList.getGroup("list", 2).getLong("element", 0)); + }); } @Test @@ -327,24 +352,25 @@ public void shouldHandleNullArrayElements() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(false); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + 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); - Assert.assertEquals(3, tagsList.getFieldRepetitionCount("list")); - // First element: present - Assert.assertEquals("first", - tagsList.getGroup("list", 0).getString("element", 0)); - // Second element: null (element field has 0 repetitions) - Assert.assertEquals(0, - tagsList.getGroup("list", 1).getFieldRepetitionCount("element")); - // Third element: present - Assert.assertEquals("third", - tagsList.getGroup("list", 2).getString("element", 0)); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + Group tagsList = record.getGroup("tags", 0); + Assert.assertEquals(3, tagsList.getFieldRepetitionCount("list")); + // First element: present + Assert.assertEquals("first", tagsList.getGroup("list", 0).getString("element", 0)); + // Second element: null (element field has 0 repetitions) + Assert.assertEquals(0, tagsList.getGroup("list", 1).getFieldRepetitionCount("element")); + // Third element: present + Assert.assertEquals("third", tagsList.getGroup("list", 2).getString("element", 0)); + }); } @Test @@ -359,15 +385,19 @@ public void shouldHandleNullSqlArray() throws SQLException, IOException { when(resultSet.wasNull()).thenReturn(true); final MessageType schema = - JdbcParquetSchema.createParquetSchema(resultSet, Optional.empty(), false, - "typed_first_row"); + 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) - Assert.assertEquals(0, record.getFieldRepetitionCount("tags")); - }); + writeAndVerify( + schema, + resultSet, + tempFile, + record -> { + // Entire array field is null (optional, 0 repetitions) + Assert.assertEquals(0, record.getFieldRepetitionCount("tags")); + }); } @FunctionalInterface @@ -375,13 +405,15 @@ 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)); + 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) From c592db8540f74ee803995fdd18eef6e3cffdbd9d Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 13 May 2026 08:58:12 +0200 Subject: [PATCH 78/80] Migrate parent POM from JUnit 4 to JUnit 5 Replace junit:junit 4.13.2 dependency management with JUnit 5 BOM (junit-bom 5.11.4) and add mockito-junit-jupiter for extension support. --- pom.xml | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/pom.xml b/pom.xml index 8605aa92..0432d0a5 100644 --- a/pom.xml +++ b/pom.xml @@ -142,7 +142,7 @@ 1.84 - 4.13.2 + 5.11.4 8.4.0 3.5.8 42.7.8 @@ -433,10 +433,11 @@ - junit - junit - ${junit.version} - test + org.junit + junit-bom + ${junit-jupiter.version} + pom + import org.hamcrest @@ -456,6 +457,12 @@ 5.20.0 test + + org.mockito + mockito-junit-jupiter + 5.20.0 + test + com.h2database h2 From 05cef428e5fc6561c6a16b939321ec5efc915c1f Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 13 May 2026 09:04:07 +0200 Subject: [PATCH 79/80] Migrate dbeam-core tests from JUnit 4 to JUnit 5 Replace JUnit 4 annotations and assertions with JUnit 5 equivalents: - @BeforeClass/@AfterClass -> @BeforeAll/@AfterAll - @Ignore -> @Disabled - @Test(expected=...) -> assertThrows() - Assert -> Assertions (with swapped message parameter order) - org.junit.* imports -> org.junit.jupiter.api.* --- dbeam-core/pom.xml | 4 +- .../dbeam/args/JdbcExportOptionsTest.java | 241 ++++++++++-------- .../dbeam/args/ParallelQueryBuilderTest.java | 6 +- .../dbeam/args/QueryBuilderArgsTest.java | 55 ++-- .../spotify/dbeam/args/QueryBuilderTest.java | 30 +-- .../dbeam/avro/JdbcAvroRecordTest.java | 81 +++--- .../dbeam/avro/JdbcAvroSchemaTest.java | 56 ++-- .../dbeam/avro/PostgresJdbcAvroTest.java | 46 ++-- .../spotify/dbeam/beam/BeamHelperTest.java | 10 +- .../spotify/dbeam/jobs/BeamHelperTest.java | 20 +- .../dbeam/jobs/BenchJdbcAvroJobTest.java | 6 +- .../spotify/dbeam/jobs/JdbcAvroJobTest.java | 68 ++--- .../spotify/dbeam/jobs/PsqlAvroJobTest.java | 2 +- .../dbeam/jobs/PsqlReplicationCheckTest.java | 42 +-- .../dbeam/options/ArrayHandlingModeTest.java | 14 +- .../dbeam/options/InputAvroSchemaTest.java | 54 ++-- .../options/JobNameConfigurationTest.java | 12 +- .../dbeam/options/PasswordReaderTest.java | 14 +- 18 files changed, 409 insertions(+), 352 deletions(-) diff --git a/dbeam-core/pom.xml b/dbeam-core/pom.xml index e7a6b22b..c5fff8e3 100644 --- a/dbeam-core/pom.xml +++ b/dbeam-core/pom.xml @@ -97,8 +97,8 @@ - junit - junit + org.junit.jupiter + junit-jupiter test 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 facd94f5..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 @@ -49,9 +49,9 @@ 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.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; public class JdbcAvroRecordTest { @@ -59,7 +59,7 @@ 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); } @@ -78,14 +78,14 @@ public void shouldCreateSchema() throws ClassNotFoundException, SQLException { 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", @@ -103,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 @@ -158,8 +159,8 @@ public void shouldCreateSchemaWithLogicalTypes() throws ClassNotFoundException, 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")); } @@ -177,7 +178,7 @@ public void shouldCreateSchemaWithCustomSchemaName() throws ClassNotFoundExcepti ArrayHandlingMode.TypedMetaFromFirstRow, false); - Assert.assertEquals("CustomSchemaName", actual.getName()); + Assertions.assertEquals("CustomSchemaName", actual.getName()); } @Test @@ -213,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()); @@ -241,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 @@ -263,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 ce8c5be9..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,34 +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( + Assertions.assertThrows( RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet, false)); - Assert.assertEquals("STRUCT type is not supported", thrown.getMessage()); + Assertions.assertEquals("STRUCT type is not supported", thrown.getMessage()); final ResultSet resultSet2 = buildMockResultSet(Types.REF); RuntimeException thrown2 = - Assert.assertThrows( + Assertions.assertThrows( RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet2, false)); - Assert.assertEquals("REF and REF_CURSOR type are not supported", thrown2.getMessage()); + Assertions.assertEquals("REF and REF_CURSOR type are not supported", thrown2.getMessage()); final ResultSet resultSet3 = buildMockResultSet(Types.REF_CURSOR); RuntimeException thrown3 = - Assert.assertThrows( + Assertions.assertThrows( RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet3, false)); - Assert.assertEquals("REF and REF_CURSOR type are not supported", thrown3.getMessage()); + Assertions.assertEquals("REF and REF_CURSOR type are not supported", thrown3.getMessage()); final ResultSet resultSet4 = buildMockResultSet(Types.DATALINK); RuntimeException thrown4 = - Assert.assertThrows( + Assertions.assertThrows( RuntimeException.class, () -> createAvroSchemaForSingleField(resultSet4, false)); - Assert.assertEquals("DATALINK type is not supported", thrown4.getMessage()); + Assertions.assertEquals("DATALINK type is not supported", thrown4.getMessage()); } @Test @@ -156,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 @@ -165,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 @@ -174,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 @@ -182,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 @@ -191,8 +191,8 @@ 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( 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 c5f4a147..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,8 +45,8 @@ 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 { @@ -87,10 +87,10 @@ public void assertGenericRecordArrayField( 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)); } } @@ -118,7 +118,7 @@ public void shouldEncodeUUIDValues() throws SQLException, IOException { 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())); } @@ -173,8 +173,8 @@ public void shouldEncodeStringAndOtherValues() throws SQLException, IOException 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")); + 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"); @@ -190,7 +190,7 @@ public void shouldThrowOnArrayWithNulls() throws SQLException, IOException { when(resultSet.getArray(1)).thenReturn(null); when(resultSet.isFirst()).thenReturn(true); - Assert.assertThrows( + Assertions.assertThrows( RuntimeException.class, () -> JdbcAvroSchema.createAvroSchema( @@ -225,7 +225,7 @@ public void shouldHandleArrayWithNullsUsingArrayAsBytes() throws SQLException, I GenericRecord actualRecord = bytesToGenericRecords(schema, converter.convertResultSetIntoAvroBytes())[0]; - Assert.assertArrayEquals( + Assertions.assertArrayEquals( expectedValue, ((java.nio.ByteBuffer) actualRecord.get("array_field")).array()); } @@ -313,12 +313,12 @@ public void shouldHandleArrayWithNullsWithoutReadingFirstRow() throws SQLExcepti 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")); + 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()); @@ -434,9 +434,9 @@ public void shouldThrowOnArrayWithNullsIfDisabled() throws SQLException { final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(resultSet, arrayMode, nullableArrayItems); RuntimeException thrown = - Assert.assertThrows( + Assertions.assertThrows( RuntimeException.class, () -> converter.convertResultSetIntoAvroBytes()); - Assert.assertEquals( + Assertions.assertEquals( "Array item is null in column 'array_field_varchar', use " + "--nullableArrayItems", thrown.getMessage()); } @@ -465,9 +465,9 @@ public void shouldThrowOnUnknownType() throws SQLException { final JdbcAvroRecordConverter converter = JdbcAvroRecordConverter.create(resultSet, arrayMode, nullableArrayItems); RuntimeException thrown = - Assert.assertThrows( + Assertions.assertThrows( RuntimeException.class, () -> converter.convertResultSetIntoAvroBytes()); - Assert.assertEquals( + Assertions.assertEquals( "Value of type class java.io.File in column 'invalid_array' is not " + "supported", thrown.getMessage()); } @@ -494,7 +494,7 @@ public void shouldThrowOnInvalidArrayColumnTypeName() throws SQLException { boolean nullableArrayItems = true; RuntimeException thrown = - Assert.assertThrows( + Assertions.assertThrows( RuntimeException.class, () -> JdbcAvroSchema.createAvroSchema( @@ -506,7 +506,7 @@ public void shouldThrowOnInvalidArrayColumnTypeName() throws SQLException { true, arrayMode, nullableArrayItems)); - Assert.assertEquals( + Assertions.assertEquals( "columnName=array_field_text columnTypeName=text should start with '_'", thrown.getMessage()); } @@ -533,7 +533,7 @@ public void shouldThrowOnNotSupportedArrayColumnTypeName() throws SQLException { boolean nullableArrayItems = true; RuntimeException thrown = - Assert.assertThrows( + Assertions.assertThrows( RuntimeException.class, () -> JdbcAvroSchema.createAvroSchema( @@ -545,7 +545,7 @@ public void shouldThrowOnNotSupportedArrayColumnTypeName() throws SQLException { true, arrayMode, nullableArrayItems)); - Assert.assertEquals( + 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 a4081f98..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 { @@ -41,9 +41,9 @@ public void shouldWorkWithGcs() throws IOException { // 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")); + && 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 bf0b695e..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 @@ -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,33 +258,38 @@ 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 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 9d10fc81..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,18 +20,18 @@ 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( + Assertions.assertEquals( ArrayHandlingMode.Bytes, ArrayHandlingMode.validateValue(ArrayHandlingMode.Bytes)); - Assert.assertEquals( + Assertions.assertEquals( ArrayHandlingMode.TypedMetaFromFirstRow, ArrayHandlingMode.validateValue(ArrayHandlingMode.TypedMetaFromFirstRow)); - Assert.assertEquals( + Assertions.assertEquals( ArrayHandlingMode.TypedMetaPostgres, ArrayHandlingMode.validateValue(ArrayHandlingMode.TypedMetaPostgres)); } @@ -39,9 +39,9 @@ public void validValuesShouldPass() { @Test public void invalidValueShouldThrow() { RuntimeException thrown = - Assert.assertThrows( + Assertions.assertThrows( RuntimeException.class, () -> ArrayHandlingMode.validateValue("invalid")); - Assert.assertEquals( + 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); } } From 6276a05c978eb8e05c46c1d176bdbc5e4a4ba1f6 Mon Sep 17 00:00:00 2001 From: Luis Bianchin Date: Wed, 13 May 2026 09:07:11 +0200 Subject: [PATCH 80/80] Migrate dbeam-parquet tests from JUnit 4 to JUnit 5 Replace JUnit 4 annotations and assertions with JUnit 5 equivalents: - @BeforeClass -> @BeforeAll - @Test(expected=...) -> assertThrows() - Assert -> Assertions - org.junit.* imports -> org.junit.jupiter.api.* --- dbeam-parquet/pom.xml | 4 +- .../parquet/BeamJdbcParquetSchemaTest.java | 14 +-- .../parquet/BenchJdbcParquetJobTest.java | 30 +++--- .../dbeam/parquet/ChannelOutputFileTest.java | 30 +++--- .../dbeam/parquet/JdbcParquetArgsTest.java | 51 +++++----- .../dbeam/parquet/JdbcParquetJobTest.java | 36 ++++---- .../dbeam/parquet/JdbcParquetRecordTest.java | 92 ++++++++++--------- .../dbeam/parquet/JdbcParquetSchemaTest.java | 72 +++++++-------- .../parquet/PostgresJdbcParquetTest.java | 67 +++++++------- .../dbeam/parquet/PsqlParquetJobTest.java | 57 +++++++----- 10 files changed, 233 insertions(+), 220 deletions(-) diff --git a/dbeam-parquet/pom.xml b/dbeam-parquet/pom.xml index 9c72de4f..adfe8e3f 100644 --- a/dbeam-parquet/pom.xml +++ b/dbeam-parquet/pom.xml @@ -72,8 +72,8 @@ test - junit - junit + org.junit.jupiter + junit-jupiter test 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 index bec6b51f..31fe33ac 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BeamJdbcParquetSchemaTest.java @@ -25,8 +25,8 @@ import java.nio.file.Path; import java.util.Optional; import org.apache.parquet.schema.MessageType; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; public class BeamJdbcParquetSchemaTest { @@ -44,9 +44,9 @@ public void shouldParseInputSchemaFile() throws IOException { final Optional schema = BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(schemaFile.toString()); - Assert.assertTrue(schema.isPresent()); - Assert.assertEquals("test_table", schema.get().getName()); - Assert.assertEquals(2, schema.get().getFieldCount()); + Assertions.assertTrue(schema.isPresent()); + Assertions.assertEquals("test_table", schema.get().getName()); + Assertions.assertEquals(2, schema.get().getFieldCount()); Files.deleteIfExists(schemaFile); } @@ -56,7 +56,7 @@ public void shouldReturnEmptyForNullFilename() throws IOException { final Optional schema = BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(null); - Assert.assertFalse(schema.isPresent()); + Assertions.assertFalse(schema.isPresent()); } @Test @@ -64,6 +64,6 @@ public void shouldReturnEmptyForEmptyFilename() throws IOException { final Optional schema = BeamJdbcParquetSchema.parseOptionalInputParquetSchemaFile(""); - Assert.assertFalse(schema.isPresent()); + 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 index 0e5ef26c..f2425260 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/BenchJdbcParquetJobTest.java @@ -30,16 +30,16 @@ import java.io.PrintStream; import java.nio.file.Path; import java.sql.SQLException; -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 BenchJdbcParquetJobTest { private static final String CONNECTION_URL = "jdbc:h2:mem:testbench;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; - @BeforeClass + @BeforeAll public static void beforeAll() throws SQLException, ClassNotFoundException { DbTestHelper.createFixtures(CONNECTION_URL); } @@ -86,17 +86,17 @@ public void shouldOutputMetricsTsvSummary() throws Exception { } final String output = capturedOut.toString(); - Assert.assertTrue(output.contains("Summary for BenchJdbcParquetJob")); - Assert.assertTrue(output.contains("recordCount")); - Assert.assertTrue(output.contains("writeElapsedMs")); - Assert.assertTrue(output.contains("bytesWritten")); - Assert.assertTrue(output.contains("KbWritePerSec")); - Assert.assertTrue(output.contains("run_00")); - Assert.assertTrue(output.contains("run_01")); - Assert.assertTrue(output.contains("max")); - Assert.assertTrue(output.contains("mean")); - Assert.assertTrue(output.contains("min")); - Assert.assertTrue(output.contains("stddev")); + 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 index e684d1f7..b82d94a9 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/ChannelOutputFileTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/ChannelOutputFileTest.java @@ -25,8 +25,8 @@ import java.nio.channels.Channels; import java.nio.channels.WritableByteChannel; import org.apache.parquet.io.PositionOutputStream; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; public class ChannelOutputFileTest { @@ -37,19 +37,19 @@ public void shouldTrackPosition() throws IOException { final ChannelOutputFile outputFile = new ChannelOutputFile(channel); try (PositionOutputStream out = outputFile.create(0)) { - Assert.assertEquals(0, out.getPos()); + Assertions.assertEquals(0, out.getPos()); out.write(42); - Assert.assertEquals(1, out.getPos()); + Assertions.assertEquals(1, out.getPos()); out.write(new byte[] {1, 2, 3, 4, 5}); - Assert.assertEquals(6, out.getPos()); + Assertions.assertEquals(6, out.getPos()); out.write(new byte[] {0, 1, 2, 3, 4, 5, 6, 7}, 2, 3); - Assert.assertEquals(9, out.getPos()); + Assertions.assertEquals(9, out.getPos()); } - Assert.assertEquals(9, baos.size()); + Assertions.assertEquals(9, baos.size()); } @Test @@ -64,10 +64,10 @@ public void shouldWriteCorrectBytes() throws IOException { } byte[] result = baos.toByteArray(); - Assert.assertEquals(3, result.length); - Assert.assertEquals(65, result[0]); - Assert.assertEquals(66, result[1]); - Assert.assertEquals(67, result[2]); + Assertions.assertEquals(3, result.length); + Assertions.assertEquals(65, result[0]); + Assertions.assertEquals(66, result[1]); + Assertions.assertEquals(67, result[2]); } @Test @@ -76,8 +76,8 @@ public void shouldNotSupportBlockSize() throws IOException { final WritableByteChannel channel = Channels.newChannel(baos); final ChannelOutputFile outputFile = new ChannelOutputFile(channel); - Assert.assertFalse(outputFile.supportsBlockSize()); - Assert.assertEquals(0, outputFile.defaultBlockSize()); + Assertions.assertFalse(outputFile.supportsBlockSize()); + Assertions.assertEquals(0, outputFile.defaultBlockSize()); } @Test @@ -87,8 +87,8 @@ public void shouldReturnSameStreamFromCreateOrOverwrite() throws IOException { final ChannelOutputFile outputFile = new ChannelOutputFile(channel); try (PositionOutputStream out = outputFile.createOrOverwrite(0)) { - Assert.assertNotNull(out); - Assert.assertEquals(0, out.getPos()); + 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 index 3b31de8d..2ee9a33a 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetArgsTest.java @@ -23,8 +23,8 @@ import com.spotify.dbeam.args.JdbcConnectionArgs; import java.util.Collections; import org.apache.parquet.hadoop.metadata.CompressionCodecName; -import org.junit.Assert; -import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; public class JdbcParquetArgsTest { @@ -42,18 +42,18 @@ private static JdbcConnectionArgs connArgs() { public void shouldCreateWithDefaults() { final JdbcParquetArgs args = JdbcParquetArgs.create(CONN_ARGS); - Assert.assertEquals(10000, args.fetchSize()); - Assert.assertEquals("snappy", args.parquetCodec()); - Assert.assertEquals(128L * 1024 * 1024, args.rowGroupSize()); - Assert.assertEquals(1024L * 1024, args.pageSize()); - Assert.assertEquals(Collections.emptyList(), args.preCommand()); + 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); - Assert.assertEquals(CompressionCodecName.SNAPPY, args.getCompressionCodecName()); + Assertions.assertEquals(CompressionCodecName.SNAPPY, args.getCompressionCodecName()); } @Test @@ -68,7 +68,7 @@ public void shouldMapGzipCodec() { Collections.emptyList(), "typed_first_row"); - Assert.assertEquals(CompressionCodecName.GZIP, args.getCompressionCodecName()); + Assertions.assertEquals(CompressionCodecName.GZIP, args.getCompressionCodecName()); } @Test @@ -83,7 +83,7 @@ public void shouldMapZstdCodec() { Collections.emptyList(), "typed_first_row"); - Assert.assertEquals(CompressionCodecName.ZSTD, args.getCompressionCodecName()); + Assertions.assertEquals(CompressionCodecName.ZSTD, args.getCompressionCodecName()); } @Test @@ -98,7 +98,7 @@ public void shouldMapLz4Codec() { Collections.emptyList(), "typed_first_row"); - Assert.assertEquals(CompressionCodecName.LZ4_RAW, args.getCompressionCodecName()); + Assertions.assertEquals(CompressionCodecName.LZ4_RAW, args.getCompressionCodecName()); } @Test @@ -113,7 +113,7 @@ public void shouldMapUncompressedCodec() { Collections.emptyList(), "typed_first_row"); - Assert.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); + Assertions.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); } @Test @@ -128,19 +128,22 @@ public void shouldMapNoneCodec() { Collections.emptyList(), "typed_first_row"); - Assert.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); + Assertions.assertEquals(CompressionCodecName.UNCOMPRESSED, args.getCompressionCodecName()); } - @Test(expected = IllegalArgumentException.class) + @Test public void shouldRejectInvalidCodec() { - JdbcParquetArgs.create( - CONN_ARGS, - 10000, - "invalid", - 64 * 1024 * 1024, - 1024 * 1024, - Collections.emptyList(), - "typed_first_row"); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + JdbcParquetArgs.create( + CONN_ARGS, + 10000, + "invalid", + 64 * 1024 * 1024, + 1024 * 1024, + Collections.emptyList(), + "typed_first_row")); } @Test @@ -155,7 +158,7 @@ public void shouldAcceptCustomFetchSize() { Collections.emptyList(), "typed_first_row"); - Assert.assertEquals(50000, args.fetchSize()); + Assertions.assertEquals(50000, args.fetchSize()); } @Test @@ -170,6 +173,6 @@ public void shouldAcceptCustomRowGroupSize() { Collections.emptyList(), "typed_first_row"); - Assert.assertEquals(128L * 1024 * 1024, args.rowGroupSize()); + 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 index c979b09e..8405a26f 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetJobTest.java @@ -31,9 +31,9 @@ import java.nio.file.Files; import java.nio.file.Path; import java.sql.SQLException; -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 JdbcParquetJobTest { @@ -43,7 +43,7 @@ public class JdbcParquetJobTest { private static Path passwordPath; private static Path sqlPath; - @BeforeClass + @BeforeAll public static void beforeAll() throws SQLException, ClassNotFoundException, IOException { testDir = TestHelper.createTmpDirPath("jdbc-parquet-test-"); passwordPath = testDir.resolve(".password"); @@ -88,7 +88,7 @@ public void shouldRunJdbcParquetJob() throws Exception { // Verify parquet schema file was written final String schemaJson = new String(Files.readAllBytes(outputPath.resolve("_PARQUET_SCHEMA.json"))); - Assert.assertTrue(schemaJson.contains("COFFEES")); + 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(); @@ -151,9 +151,9 @@ public void shouldRunJdbcParquetJobSqlFile() throws Exception { // Verify schema only has the 3 selected columns final String schemaJson = new String(Files.readAllBytes(outputPath.resolve("_PARQUET_SCHEMA.json"))); - Assert.assertTrue(schemaJson.contains("COF_NAME")); - Assert.assertTrue(schemaJson.contains("SIZE")); - Assert.assertTrue(schemaJson.contains("TOTAL")); + 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)); @@ -191,7 +191,7 @@ public void shouldRunParquetJobWithLogicalTypes() throws Exception { final String schemaJson = new String(Files.readAllBytes(outputPath.resolve("_PARQUET_SCHEMA.json"))); - Assert.assertTrue(schemaJson.contains("TIMESTAMP")); + Assertions.assertTrue(schemaJson.contains("TIMESTAMP")); } @Test @@ -249,21 +249,21 @@ public void shouldRunParquetJobWithInputSchemaFile() throws Exception { final String schemaJson = new String(Files.readAllBytes(outputPath.resolve("_PARQUET_SCHEMA.json"))); - Assert.assertTrue(schemaJson.contains("COF_NAME")); - Assert.assertTrue(schemaJson.contains("SIZE")); - Assert.assertTrue(schemaJson.contains("TOTAL")); + 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() { - Assert.assertEquals("snappy", JdbcParquetJob.mapAvroCodecToParquetCodec("snappy")); - Assert.assertEquals("gzip", JdbcParquetJob.mapAvroCodecToParquetCodec("deflate1")); - Assert.assertEquals("gzip", JdbcParquetJob.mapAvroCodecToParquetCodec("deflate6")); - Assert.assertEquals("gzip", JdbcParquetJob.mapAvroCodecToParquetCodec("deflate9")); - Assert.assertEquals("zstd", JdbcParquetJob.mapAvroCodecToParquetCodec("zstandard1")); - Assert.assertEquals("zstd", JdbcParquetJob.mapAvroCodecToParquetCodec("zstandard9")); + 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 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 index 02708677..7ee5c69e 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetRecordTest.java @@ -40,16 +40,16 @@ import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.parquet.io.OutputFile; import org.apache.parquet.schema.MessageType; -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 JdbcParquetRecordTest { private static final String CONNECTION_URL = "jdbc:h2:mem:testrecord;MODE=PostgreSQL;DATABASE_TO_UPPER=false;DB_CLOSE_DELAY=-1"; - @BeforeClass + @BeforeAll public static void beforeAll() throws SQLException, ClassNotFoundException { DbTestHelper.createFixtures(CONNECTION_URL); } @@ -65,23 +65,23 @@ public void shouldCreateSchemaFromDatabase() throws ClassNotFoundException, SQLE false, "typed_first_row"); - Assert.assertNotNull(schema); - Assert.assertEquals("COFFEES", schema.getName()); - Assert.assertEquals(14, schema.getFieldCount()); - Assert.assertEquals("COF_NAME", schema.getFields().get(0).getName()); - Assert.assertEquals("SUP_ID", schema.getFields().get(1).getName()); - Assert.assertEquals("PRICE", schema.getFields().get(2).getName()); - Assert.assertEquals("TEMPERATURE", schema.getFields().get(3).getName()); - Assert.assertEquals("SIZE", schema.getFields().get(4).getName()); - Assert.assertEquals("IS_ARABIC", schema.getFields().get(5).getName()); - Assert.assertEquals("SALES", schema.getFields().get(6).getName()); - Assert.assertEquals("TOTAL", schema.getFields().get(7).getName()); - Assert.assertEquals("CREATED", schema.getFields().get(8).getName()); - Assert.assertEquals("UPDATED", schema.getFields().get(9).getName()); - Assert.assertEquals("UID", schema.getFields().get(10).getName()); - Assert.assertEquals("ROWNUM", schema.getFields().get(11).getName()); - Assert.assertEquals("INT_ARR", schema.getFields().get(12).getName()); - Assert.assertEquals("TEXT_ARR", schema.getFields().get(13).getName()); + 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 @@ -95,10 +95,12 @@ public void shouldCreateSchemaWithLogicalTypes() throws ClassNotFoundException, true, "typed_first_row"); - Assert.assertEquals(14, schema.getFieldCount()); + Assertions.assertEquals(14, schema.getFieldCount()); // CREATED and UPDATED should have timestamp logical type - Assert.assertNotNull(schema.getFields().get(8).asPrimitiveType().getLogicalTypeAnnotation()); - Assert.assertNotNull(schema.getFields().get(9).asPrimitiveType().getLogicalTypeAnnotation()); + Assertions.assertNotNull( + schema.getFields().get(8).asPrimitiveType().getLogicalTypeAnnotation()); + Assertions.assertNotNull( + schema.getFields().get(9).asPrimitiveType().getLogicalTypeAnnotation()); } @Test @@ -112,7 +114,7 @@ public void shouldCreateSchemaWithCustomName() throws ClassNotFoundException, SQ false, "typed_first_row"); - Assert.assertEquals("CustomSchema", schema.getName()); + Assertions.assertEquals("CustomSchema", schema.getName()); } @Test @@ -155,17 +157,17 @@ public void shouldWriteAndReadBackParquetRecords() try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { Group record1 = reader.read(); - Assert.assertNotNull(record1); + Assertions.assertNotNull(record1); // Verify first record has expected fields String cofName = record1.getString("COF_NAME", 0); - Assert.assertNotNull(cofName); + Assertions.assertNotNull(cofName); Group record2 = reader.read(); - Assert.assertNotNull(record2); + Assertions.assertNotNull(record2); // Should only be 2 records Group record3 = reader.read(); - Assert.assertNull(record3); + Assertions.assertNull(record3); } // Cleanup @@ -211,29 +213,29 @@ public void shouldWriteCorrectFieldValues() try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { Group record = reader.read(); - Assert.assertNotNull(record); + Assertions.assertNotNull(record); // Verify typed values for the first record (colombian caffee, sorted) String cofName = record.getString("COF_NAME", 0); - Assert.assertEquals("colombian caffee", cofName); + Assertions.assertEquals("colombian caffee", cofName); boolean isArabic = record.getBoolean("IS_ARABIC", 0); - Assert.assertTrue(isArabic); + Assertions.assertTrue(isArabic); int sales = record.getInteger("SALES", 0); - Assert.assertEquals(13, sales); + Assertions.assertEquals(13, sales); long total = record.getLong("TOTAL", 0); - Assert.assertEquals(201L, total); + Assertions.assertEquals(201L, total); float temperature = record.getFloat("TEMPERATURE", 0); - Assert.assertEquals(87.5f, temperature, 0.01f); + Assertions.assertEquals(87.5f, temperature, 0.01f); double size = record.getDouble("SIZE", 0); - Assert.assertEquals(230.7, size, 0.01); + Assertions.assertEquals(230.7, size, 0.01); long rownum = record.getLong("ROWNUM", 0); - Assert.assertEquals(2L, rownum); + Assertions.assertEquals(2L, rownum); } Files.deleteIfExists(tempFile); @@ -277,17 +279,17 @@ public void shouldHandleNullValues() throws ClassNotFoundException, SQLException try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { Group record = reader.read(); - Assert.assertNotNull(record); + Assertions.assertNotNull(record); // SUP_ID is null - should have 0 repetitions in Parquet - Assert.assertEquals(0, record.getFieldRepetitionCount("SUP_ID")); + Assertions.assertEquals(0, record.getFieldRepetitionCount("SUP_ID")); // UPDATED is null - Assert.assertEquals(0, record.getFieldRepetitionCount("UPDATED")); + Assertions.assertEquals(0, record.getFieldRepetitionCount("UPDATED")); // Non-null fields should have 1 repetition - Assert.assertEquals(1, record.getFieldRepetitionCount("COF_NAME")); - Assert.assertEquals(1, record.getFieldRepetitionCount("IS_ARABIC")); + Assertions.assertEquals(1, record.getFieldRepetitionCount("COF_NAME")); + Assertions.assertEquals(1, record.getFieldRepetitionCount("IS_ARABIC")); } Files.deleteIfExists(tempFile); @@ -337,9 +339,9 @@ public void shouldWriteAvroSchemaInParquetFooter() ParquetFileReader.open(HadoopInputFile.fromPath(hadoopPath, conf))) { final Map keyValueMetaData = fileReader.getFooter().getFileMetaData().getKeyValueMetaData(); - Assert.assertTrue(keyValueMetaData.containsKey(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY)); + Assertions.assertTrue(keyValueMetaData.containsKey(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY)); final String actualAvroSchema = keyValueMetaData.get(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY); - Assert.assertEquals(avroSchemaJson, actualAvroSchema); + Assertions.assertEquals(avroSchemaJson, actualAvroSchema); } Files.deleteIfExists(tempFile); @@ -385,7 +387,7 @@ public void shouldOmitAvroSchemaWhenNotProvided() ParquetFileReader.open(HadoopInputFile.fromPath(hadoopPath, conf))) { final Map keyValueMetaData = fileReader.getFooter().getFileMetaData().getKeyValueMetaData(); - Assert.assertFalse(keyValueMetaData.containsKey(JdbcParquetIO.PARQUET_AVRO_SCHEMA_KEY)); + 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 index 1f6c2ad5..9979d934 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/JdbcParquetSchemaTest.java @@ -31,8 +31,8 @@ import org.apache.parquet.schema.MessageType; import org.apache.parquet.schema.PrimitiveType; import org.apache.parquet.schema.Type; -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 JdbcParquetSchemaTest { @@ -45,7 +45,7 @@ public void shouldGetDatabaseTableNameFromMetaData() throws SQLException { when(meta.getColumnCount()).thenReturn(1); when(meta.getTableName(1)).thenReturn("test_table"); - Assert.assertEquals("test_table", JdbcParquetSchema.getDatabaseTableName(meta)); + Assertions.assertEquals("test_table", JdbcParquetSchema.getDatabaseTableName(meta)); } @Test @@ -54,7 +54,7 @@ public void shouldDefaultTableNameWhenMetaDataHasEmptyTableName() throws SQLExce when(meta.getColumnCount()).thenReturn(1); when(meta.getTableName(1)).thenReturn(""); - Assert.assertEquals("no_table_name", JdbcParquetSchema.getDatabaseTableName(meta)); + Assertions.assertEquals("no_table_name", JdbcParquetSchema.getDatabaseTableName(meta)); } @Test @@ -63,7 +63,7 @@ public void shouldDefaultTableNameWhenMetaDataHasNullTableName() throws SQLExcep when(meta.getColumnCount()).thenReturn(1); when(meta.getTableName(1)).thenReturn(null); - Assert.assertEquals("no_table_name", JdbcParquetSchema.getDatabaseTableName(meta)); + Assertions.assertEquals("no_table_name", JdbcParquetSchema.getDatabaseTableName(meta)); } @Test @@ -73,7 +73,7 @@ public void shouldGetDatabaseTableNameFromFirstNonNullMetaData() throws SQLExcep when(meta.getTableName(1)).thenReturn(""); when(meta.getTableName(2)).thenReturn("test_table"); - Assert.assertEquals("test_table", JdbcParquetSchema.getDatabaseTableName(meta)); + Assertions.assertEquals("test_table", JdbcParquetSchema.getDatabaseTableName(meta)); } @Test @@ -118,7 +118,7 @@ public void shouldConvertTimestampSqlTypeToInt64() throws SQLException { final Type fieldType = buildFieldType(Types.TIMESTAMP, false); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); - Assert.assertNull(fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + Assertions.assertNull(fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @Test @@ -126,7 +126,7 @@ public void shouldConvertTimestampSqlTypeWithLogicalType() throws SQLException { final Type fieldType = buildFieldType(Types.TIMESTAMP, true); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); - Assert.assertEquals( + Assertions.assertEquals( LogicalTypeAnnotation.timestampType(true, LogicalTypeAnnotation.TimeUnit.MILLIS), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @@ -143,7 +143,7 @@ public void shouldConvertDateSqlTypeWithLogicalType() throws SQLException { final Type fieldType = buildFieldType(Types.DATE, true); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.INT64); - Assert.assertNotNull(fieldType.asPrimitiveType().getLogicalTypeAnnotation()); + Assertions.assertNotNull(fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @Test @@ -223,7 +223,7 @@ public void shouldConvertVarcharSqlTypeToStringBinary() throws SQLException { final Type fieldType = buildFieldType(Types.VARCHAR, false); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); - Assert.assertEquals( + Assertions.assertEquals( LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @@ -232,7 +232,7 @@ public void shouldConvertCharSqlTypeToStringBinary() throws SQLException { final Type fieldType = buildFieldType(Types.CHAR, false); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); - Assert.assertEquals( + Assertions.assertEquals( LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @@ -241,7 +241,7 @@ public void shouldConvertClobSqlTypeToStringBinary() throws SQLException { final Type fieldType = buildFieldType(Types.CLOB, false); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); - Assert.assertEquals( + Assertions.assertEquals( LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @@ -251,9 +251,9 @@ public void shouldConvertArraySqlTypeToListWithStringElements() throws SQLExcept JdbcParquetSchema.buildParquetFieldType( "column1", Types.ARRAY, 0, "java.sql.Array", "_text", false, "typed_first_row"); - Assert.assertFalse(fieldType.isPrimitive()); - Assert.assertEquals(LogicalTypeAnnotation.listType(), fieldType.getLogicalTypeAnnotation()); - Assert.assertEquals(Type.Repetition.OPTIONAL, fieldType.getRepetition()); + Assertions.assertFalse(fieldType.isPrimitive()); + Assertions.assertEquals(LogicalTypeAnnotation.listType(), fieldType.getLogicalTypeAnnotation()); + Assertions.assertEquals(Type.Repetition.OPTIONAL, fieldType.getRepetition()); } @Test @@ -296,24 +296,24 @@ public void shouldConvertTextArrayToListOfString() { final Type elementType = JdbcParquetSchema.buildArrayElementType("_text"); assertPrimitiveType(elementType, PrimitiveType.PrimitiveTypeName.BINARY); - Assert.assertEquals( + Assertions.assertEquals( LogicalTypeAnnotation.stringType(), elementType.asPrimitiveType().getLogicalTypeAnnotation()); } @Test public void shouldResolveH2IntegerArrayType() { - Assert.assertEquals("int4", JdbcParquetSchema.resolveArrayElementTypeName("INTEGER ARRAY")); + Assertions.assertEquals("int4", JdbcParquetSchema.resolveArrayElementTypeName("INTEGER ARRAY")); } @Test public void shouldResolveH2VarcharArrayType() { - Assert.assertEquals("text", JdbcParquetSchema.resolveArrayElementTypeName("VARCHAR ARRAY")); + Assertions.assertEquals("text", JdbcParquetSchema.resolveArrayElementTypeName("VARCHAR ARRAY")); } @Test public void shouldResolveNullColumnTypeName() { - Assert.assertEquals("text", JdbcParquetSchema.resolveArrayElementTypeName(null)); + Assertions.assertEquals("text", JdbcParquetSchema.resolveArrayElementTypeName(null)); } @Test @@ -323,8 +323,8 @@ public void shouldConvertUuidWithLogicalType() throws SQLException { "column1", Types.OTHER, 0, "foobar", "uuid", true, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY); - Assert.assertEquals(16, fieldType.asPrimitiveType().getTypeLength()); - Assert.assertEquals( + Assertions.assertEquals(16, fieldType.asPrimitiveType().getTypeLength()); + Assertions.assertEquals( LogicalTypeAnnotation.uuidType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @@ -335,7 +335,7 @@ public void shouldConvertUuidWithoutLogicalType() throws SQLException { "column1", Types.OTHER, 0, "foobar", "uuid", false, "typed_first_row"); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); - Assert.assertEquals( + Assertions.assertEquals( LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @@ -344,7 +344,7 @@ public void shouldDefaultConversionToStringBinary() throws SQLException { final Type fieldType = buildFieldType(Types.SQLXML, false); assertPrimitiveType(fieldType, PrimitiveType.PrimitiveTypeName.BINARY); - Assert.assertEquals( + Assertions.assertEquals( LogicalTypeAnnotation.stringType(), fieldType.asPrimitiveType().getLogicalTypeAnnotation()); } @@ -372,17 +372,17 @@ public void shouldCreateParquetSchemaFromMockResultSet() throws SQLException { JdbcParquetSchema.createParquetSchema( resultSet, Optional.empty(), false, "typed_first_row"); - Assert.assertEquals("test_table", schema.getName()); - Assert.assertEquals(3, schema.getFieldCount()); - Assert.assertEquals("id", schema.getFields().get(0).getName()); - Assert.assertEquals("name", schema.getFields().get(1).getName()); - Assert.assertEquals("active", schema.getFields().get(2).getName()); + 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); - Assert.assertTrue(schema.getFields().get(0).isRepetition(Type.Repetition.OPTIONAL)); - Assert.assertTrue(schema.getFields().get(1).isRepetition(Type.Repetition.OPTIONAL)); - Assert.assertTrue(schema.getFields().get(2).isRepetition(Type.Repetition.OPTIONAL)); + 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 @@ -401,7 +401,7 @@ public void shouldUseCustomSchemaName() throws SQLException { JdbcParquetSchema.createParquetSchema( resultSet, Optional.of("CustomName"), false, "typed_first_row"); - Assert.assertEquals("CustomName", schema.getName()); + Assertions.assertEquals("CustomName", schema.getName()); } @Test @@ -420,8 +420,8 @@ public void shouldNormalizeSpecialCharactersInColumnNames() throws SQLException JdbcParquetSchema.createParquetSchema( resultSet, Optional.empty(), false, "typed_first_row"); - Assert.assertEquals("test_table_name", schema.getName()); - Assert.assertEquals("column_name_with_spaces", schema.getFields().get(0).getName()); + 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) { @@ -431,7 +431,7 @@ private Type buildFieldType(final int sqlType, final boolean useLogicalTypes) { private void assertPrimitiveType( final Type actual, final PrimitiveType.PrimitiveTypeName expected) { - Assert.assertTrue(actual.isPrimitive()); - Assert.assertEquals(expected, actual.asPrimitiveType().getPrimitiveTypeName()); + 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 index 56d947c9..a2d89a40 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PostgresJdbcParquetTest.java @@ -41,8 +41,8 @@ import org.apache.parquet.hadoop.example.GroupReadSupport; import org.apache.parquet.io.OutputFile; import org.apache.parquet.schema.MessageType; -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 PostgresJdbcParquetTest { @@ -77,7 +77,7 @@ public void shouldEncodeUuidAsString() throws SQLException, IOException { tempFile, record -> { String actualUuid = record.getString("uuid_field", 0); - Assert.assertEquals(uuidExpected.toString(), actualUuid); + Assertions.assertEquals(uuidExpected.toString(), actualUuid); }); } @@ -107,8 +107,8 @@ public void shouldEncodeStringAndOtherTypes() throws SQLException, IOException { resultSet, tempFile, record -> { - Assert.assertEquals("some_text_42", record.getString("text_field", 0)); - Assert.assertEquals("some_other_42", record.getString("other_field", 0)); + Assertions.assertEquals("some_text_42", record.getString("text_field", 0)); + Assertions.assertEquals("some_other_42", record.getString("other_field", 0)); }); } @@ -137,7 +137,7 @@ public void shouldEncodeTimestampAsMillis() throws SQLException, IOException { tempFile, record -> { long actualTs = record.getLong("ts_field", 0); - Assert.assertEquals(1488300933000L, actualTs); + Assertions.assertEquals(1488300933000L, actualTs); }); } @@ -167,8 +167,8 @@ public void shouldHandleNullFields() throws SQLException, IOException { resultSet, tempFile, record -> { - Assert.assertEquals("alice", record.getString("name", 0)); - Assert.assertEquals(0, record.getFieldRepetitionCount("age")); + Assertions.assertEquals("alice", record.getString("name", 0)); + Assertions.assertEquals(0, record.getFieldRepetitionCount("age")); }); } @@ -203,10 +203,10 @@ public void shouldEncodeMultipleNumericTypes() throws SQLException, IOException resultSet, tempFile, record -> { - Assert.assertEquals(42, record.getInteger("int_col", 0)); - Assert.assertEquals(9999999999L, record.getLong("long_col", 0)); - Assert.assertEquals(3.14f, record.getFloat("float_col", 0), 0.001f); - Assert.assertEquals(2.71828, record.getDouble("double_col", 0), 0.00001); + 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); }); } @@ -232,7 +232,7 @@ public void shouldEncodeBooleanType() throws SQLException, IOException { resultSet, tempFile, record -> { - Assert.assertTrue(record.getBoolean("is_active", 0)); + Assertions.assertTrue(record.getBoolean("is_active", 0)); }); } @@ -254,8 +254,8 @@ public void shouldEncodeArrayAsParquetList() throws SQLException, IOException { resultSet, Optional.empty(), false, "typed_first_row"); // Verify schema has LIST type - Assert.assertFalse(schema.getFields().get(0).isPrimitive()); - Assert.assertEquals( + Assertions.assertFalse(schema.getFields().get(0).isPrimitive()); + Assertions.assertEquals( org.apache.parquet.schema.LogicalTypeAnnotation.listType(), schema.getFields().get(0).getLogicalTypeAnnotation()); @@ -267,10 +267,10 @@ public void shouldEncodeArrayAsParquetList() throws SQLException, IOException { tempFile, record -> { Group tagsList = record.getGroup("tags", 0); - Assert.assertEquals(3, tagsList.getFieldRepetitionCount("list")); - Assert.assertEquals("rock", tagsList.getGroup("list", 0).getString("element", 0)); - Assert.assertEquals("jazz", tagsList.getGroup("list", 1).getString("element", 0)); - Assert.assertEquals("blues", tagsList.getGroup("list", 2).getString("element", 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)); }); } @@ -299,10 +299,10 @@ public void shouldEncodeIntegerArrayAsTypedList() throws SQLException, IOExcepti tempFile, record -> { Group scoresList = record.getGroup("scores", 0); - Assert.assertEquals(3, scoresList.getFieldRepetitionCount("list")); - Assert.assertEquals(10, scoresList.getGroup("list", 0).getInteger("element", 0)); - Assert.assertEquals(20, scoresList.getGroup("list", 1).getInteger("element", 0)); - Assert.assertEquals(30, scoresList.getGroup("list", 2).getInteger("element", 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)); }); } @@ -331,10 +331,10 @@ public void shouldEncodeBigintArrayAsTypedList() throws SQLException, IOExceptio tempFile, record -> { Group idsList = record.getGroup("ids", 0); - Assert.assertEquals(3, idsList.getFieldRepetitionCount("list")); - Assert.assertEquals(100L, idsList.getGroup("list", 0).getLong("element", 0)); - Assert.assertEquals(200L, idsList.getGroup("list", 1).getLong("element", 0)); - Assert.assertEquals(300L, idsList.getGroup("list", 2).getLong("element", 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)); }); } @@ -363,13 +363,14 @@ public void shouldHandleNullArrayElements() throws SQLException, IOException { tempFile, record -> { Group tagsList = record.getGroup("tags", 0); - Assert.assertEquals(3, tagsList.getFieldRepetitionCount("list")); + Assertions.assertEquals(3, tagsList.getFieldRepetitionCount("list")); // First element: present - Assert.assertEquals("first", tagsList.getGroup("list", 0).getString("element", 0)); + Assertions.assertEquals("first", tagsList.getGroup("list", 0).getString("element", 0)); // Second element: null (element field has 0 repetitions) - Assert.assertEquals(0, tagsList.getGroup("list", 1).getFieldRepetitionCount("element")); + Assertions.assertEquals( + 0, tagsList.getGroup("list", 1).getFieldRepetitionCount("element")); // Third element: present - Assert.assertEquals("third", tagsList.getGroup("list", 2).getString("element", 0)); + Assertions.assertEquals("third", tagsList.getGroup("list", 2).getString("element", 0)); }); } @@ -396,7 +397,7 @@ public void shouldHandleNullSqlArray() throws SQLException, IOException { tempFile, record -> { // Entire array field is null (optional, 0 repetitions) - Assert.assertEquals(0, record.getFieldRepetitionCount("tags")); + Assertions.assertEquals(0, record.getFieldRepetitionCount("tags")); }); } @@ -427,7 +428,7 @@ private void writeAndVerify( try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), hadoopPath).withConf(conf).build()) { Group record = reader.read(); - Assert.assertNotNull(record); + Assertions.assertNotNull(record); assertion.assertRecord(record); } 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 index 6bee80e7..0abe565f 100644 --- a/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PsqlParquetJobTest.java +++ b/dbeam-parquet/src/test/java/com/spotify/dbeam/parquet/PsqlParquetJobTest.java @@ -23,46 +23,53 @@ import com.spotify.dbeam.TestHelper; import java.io.IOException; import java.nio.file.Path; -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 PsqlParquetJobTest { private static Path testDir; private static Path passwordPath; - @BeforeClass + @BeforeAll public static void beforeAll() throws IOException { testDir = TestHelper.createTmpDirPath("psql-parquet-test-"); passwordPath = testDir.resolve(".password"); passwordPath.toFile().createNewFile(); } - @Test(expected = IllegalArgumentException.class) - public void shouldFailOnNonPostgresConnection() throws Exception { + @Test + public void shouldFailOnNonPostgresConnection() { final Path outputPath = testDir.resolve("shouldFailOnNonPostgres"); - PsqlParquetJob.create( - new String[] { - "--connectionUrl=jdbc:h2:mem:test", - "--username=", - "--passwordFile=" + passwordPath.toString(), - "--table=COFFEES", - "--output=" + outputPath, - "--partition=2025-02-28" - }); + 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(expected = IllegalArgumentException.class) - public void shouldFailOnMissingPartition() throws Exception { + @Test + public void shouldFailOnMissingPartition() { final Path outputPath = testDir.resolve("shouldFailOnMissingPartition"); - PsqlParquetJob.create( - new String[] { - "--connectionUrl=jdbc:postgresql://localhost/test", - "--username=", - "--passwordFile=" + passwordPath.toString(), - "--table=COFFEES", - "--output=" + outputPath, - "--skipPartitionCheck" - }); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + PsqlParquetJob.create( + new String[] { + "--connectionUrl=jdbc:postgresql://localhost/test", + "--username=", + "--passwordFile=" + passwordPath.toString(), + "--table=COFFEES", + "--output=" + outputPath, + "--skipPartitionCheck" + })); } }