From 8c0ef225ae5602ad078808239b2d826c4ae4dabe Mon Sep 17 00:00:00 2001 From: Siddharth Agrawal Date: Tue, 7 Nov 2023 15:12:32 -0800 Subject: [PATCH] Ensure traceId is set for write --- CHANGES.md | 1 + .../google/cloud/spark/bigquery/SparkBigQueryConfig.java | 3 ++- .../cloud/spark/bigquery/SparkBigQueryConfigTest.java | 8 +++++--- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 1e58b9eac..0f38add79 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -4,6 +4,7 @@ * PR #1117: Make read session caching duration configurable * PR #1118: Improve read session caching key +* PR #1122: Set traceId on write * PR #1127: Fix job labeling for mixed case Dataproc job names ## 0.34.0 - 2023-10-31 diff --git a/spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/SparkBigQueryConfig.java b/spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/SparkBigQueryConfig.java index ec0d2eb2c..29765d69a 100644 --- a/spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/SparkBigQueryConfig.java +++ b/spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/SparkBigQueryConfig.java @@ -526,7 +526,8 @@ public static SparkBigQueryConfig from( } com.google.common.base.Optional traceApplicationNameParam = - getAnyOption(globalOptions, options, "traceApplicationName"); + getAnyOption(globalOptions, options, "traceApplicationName") + .or(com.google.common.base.Optional.fromNullable("traceApplicationName")); config.traceId = traceApplicationNameParam.transform( traceApplicationName -> { diff --git a/spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/SparkBigQueryConfigTest.java b/spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/SparkBigQueryConfigTest.java index 5649e2128..41d0d8faf 100644 --- a/spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/SparkBigQueryConfigTest.java +++ b/spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/SparkBigQueryConfigTest.java @@ -124,7 +124,8 @@ public void testDefaults() { assertThat(config.getWriteMethod()).isEqualTo(SparkBigQueryConfig.WriteMethod.INDIRECT); assertThat(config.getCacheExpirationTimeInMinutes()) .isEqualTo(SparkBigQueryConfig.DEFAULT_CACHE_EXPIRATION_IN_MINUTES); - assertThat(config.getTraceId().isPresent()).isFalse(); + assertThat(config.getTraceId().isPresent()).isTrue(); + assertThat(config.getTraceId().get().startsWith("Spark:traceApplicationName:")); assertThat(config.getBigQueryJobLabels()).isEmpty(); assertThat(config.getEnableModeCheckForSchemaFields()).isTrue(); assertThat(config.getDatetimeZoneId()).isEqualTo(ZoneId.of("UTC")); @@ -171,7 +172,7 @@ public void testConfigFromOptions() { .put("writeMethod", "direct") .put("cacheExpirationTimeInMinutes", "100") .put("traceJobId", "traceJobId") - .put("traceApplicationName", "traceApplicationName") + .put("traceApplicationName", "traceApplicationNameTest") .put("bigQueryJobLabel.foo", "bar") .put("enableModeCheckForSchemaFields", "false") .put("datetimeZoneId", "Asia/Jerusalem") @@ -223,7 +224,8 @@ public void testConfigFromOptions() { assertThat(config.getArrowCompressionCodec()).isEqualTo(CompressionCodec.ZSTD); assertThat(config.getWriteMethod()).isEqualTo(SparkBigQueryConfig.WriteMethod.DIRECT); assertThat(config.getCacheExpirationTimeInMinutes()).isEqualTo(100); - assertThat(config.getTraceId()).isEqualTo(Optional.of("Spark:traceApplicationName:traceJobId")); + assertThat(config.getTraceId()) + .isEqualTo(Optional.of("Spark:traceApplicationNameTest:traceJobId")); assertThat(config.getBigQueryJobLabels()).hasSize(1); assertThat(config.getBigQueryJobLabels()).containsEntry("foo", "bar"); assertThat(config.getEnableModeCheckForSchemaFields()).isFalse();