Addressed comments

Author: podgaietska

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/RecordTableResolver.java

@@ -1,3 +1,26 @@
+/*
+ * Copyright 2024 Copyright 2022 Aiven Oy and
+ * bigquery-connector-for-apache-kafka project contributors
+ *
+ * This software contains code derived from the Confluent BigQuery
+ * Kafka Connector, Copyright Confluent, Inc, which in turn
+ * contains code derived from the WePay BigQuery Kafka Connector,
+ * Copyright WePay, Inc.
+ *
+ * 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.wepay.kafka.connect.bigquery;
 
 import com.google.cloud.bigquery.BigQuery;
@@ -31,7 +54,6 @@ class RecordTableResolver {
   private final boolean usePartitionDecorator;
   private final boolean upsertDelete;
   private final boolean useMessageTimeDatePartitioning;
-  private final boolean useStorageApi;
 
   public RecordTableResolver(BigQuerySinkTaskConfig config, MergeBatches mergeBatches, BigQuery bigQuery) {
     this.config = config;
@@ -40,11 +62,10 @@ public RecordTableResolver(BigQuerySinkTaskConfig config, MergeBatches mergeBatc
 
     this.useMessageTimeDatePartitioning =
             config.getBoolean(BigQuerySinkConfig.BIGQUERY_MESSAGE_TIME_PARTITIONING_CONFIG);
-    this.usePartitionDecorator =
-            config.getBoolean(BigQuerySinkConfig.BIGQUERY_PARTITION_DECORATOR_CONFIG);
-    this.useStorageApi =
-            config.getBoolean(BigQuerySinkConfig.USE_STORAGE_WRITE_API_CONFIG);
-    this.upsertDelete = !useStorageApi && (config.getBoolean(BigQuerySinkConfig.UPSERT_ENABLED_CONFIG)
+    this.usePartitionDecorator = config.getBoolean(BigQuerySinkConfig.BIGQUERY_PARTITION_DECORATOR_CONFIG)
+            && !config.getBoolean(BigQuerySinkConfig.ENABLE_BATCH_MODE_CONFIG);
+    this.upsertDelete = !config.getBoolean(BigQuerySinkConfig.USE_STORAGE_WRITE_API_CONFIG)
+            && (config.getBoolean(BigQuerySinkConfig.UPSERT_ENABLED_CONFIG)
             || config.getBoolean(BigQuerySinkConfig.DELETE_ENABLED_CONFIG));
   }
 
@@ -74,9 +95,7 @@ private TableId getBaseTableId(String topic) {
     return topicToTableId.computeIfAbsent(topic, topicName -> {
       String[] datasetAndTable = TableNameUtils.getDataSetAndTableName(config, topic);
       String project = config.getString(BigQuerySinkConfig.PROJECT_CONFIG);
-      TableId baseTableId = (!useStorageApi)
-          ? TableId.of(datasetAndTable[0], datasetAndTable[1])
-          : TableId.of(project, datasetAndTable[0], datasetAndTable[1]);
+      TableId baseTableId = TableId.of(project, datasetAndTable[0], datasetAndTable[1]);
 
       if (usePartitionDecorator) {
         validatePartitioningForDecorator(baseTableId);
@@ -90,7 +109,7 @@ private void validatePartitioningForDecorator(TableId tableId) {
     StandardTableDefinition definition = retrieveTableDefinition(tableId);
     if (definition == null) {
       // If we could not find table and its definition, ignore.
-      // Table creation will potentially be handled later via SchemaManager
+      // Table creation will potentially be handled later if auto create is enabled
       return;
     }
     TimePartitioning partitioning = definition.getTimePartitioning();

kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/BigQuerySinkTaskTest.java

@@ -823,47 +823,6 @@ public void testInterruptedException() {
     );
   }
 
-  @Test
-  public void testTimePartitioningIncompatibleWithDecoratorSyntax() {
-    final String topic = "t1";
-    final String dataset = "d";
-
-    Map<String, String> properties = propertiesFactory.getProperties();
-    properties.put(BigQuerySinkConfig.BIGQUERY_PARTITION_DECORATOR_CONFIG, "true");
-    properties.put(BigQuerySinkConfig.BIGQUERY_MESSAGE_TIME_PARTITIONING_CONFIG, "true");
-    properties.put(BigQuerySinkConfig.TOPICS_CONFIG, topic);
-    properties.put(BigQuerySinkConfig.DEFAULT_DATASET_CONFIG, dataset);
-
-    Storage storage = mock(Storage.class);
-    BigQuery bigQuery = mock(BigQuery.class);
-
-    TableId tableId = TableId.of(dataset, topic);
-    StandardTableDefinition mockTableDefinition = mock(StandardTableDefinition.class);
-    when(mockTableDefinition.getTimePartitioning()).thenReturn(TimePartitioning.of(TimePartitioning.Type.HOUR));
-    Table table = mock(Table.class);
-    when(table.getDefinition()).thenReturn(mockTableDefinition);
-    when(bigQuery.getTable(tableId)).thenReturn(table);
-
-    BigQuerySinkTask testTask = new BigQuerySinkTask(
-        bigQuery,
-        null,
-        storage,
-        null,
-        mockedStorageWriteApiDefaultStream,
-        mockedBatchHandler,
-        time
-    );
-
-    SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
-    testTask.initialize(sinkTaskContext);
-    testTask.start(properties);
-
-    assertThrows(
-        ConnectException.class,
-        () -> testTask.put(Collections.singleton(spoofSinkRecord(topic, "f1", "v1", TimestampType.CREATE_TIME, 1L)))
-    );
-  }
-
   @Test
   public void testVersion() {
     assertNotNull(new BigQuerySinkTask().version());

kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/RecordTableResolverTest.java

@@ -26,7 +26,7 @@ public class RecordTableResolverTest {
     private static final String TOPIC = "topic";
     private static final String DATASET = "dataset";
     private static final String PROJECT = "project";
-    private static final TableId BASE_TABLE_ID = TableId.of("dataset", "topic");
+    private static final TableId BASE_TABLE_ID = TableId.of("project", "dataset", "topic");
     private BigQuerySinkTaskConfig mockConfig;
     private BigQuery mockBigQuery;
     private MergeBatches mockMergeBatches;
@@ -59,7 +59,6 @@ public void testResolvesTable() {
         TableId tableId = recordTableResolver.getRecordTable(mockRecord).getFullTableId();
 
         assertEquals(TOPIC, tableId.getTable());
-        assertNull(tableId.getProject());
     }
 
     @Test
@@ -77,28 +76,30 @@ public void testUpsertDeleteResolvesIntermediateTable() {
     }
 
     @Test
-    public void testResolveTableWithStorageApi() {
+    public void testResolveTableWithStorageApiIgnoresUpsertDelete() {
+        when(mockConfig.getBoolean(BigQuerySinkConfig.UPSERT_ENABLED_CONFIG)).thenReturn(true);
         when(mockConfig.getBoolean(BigQuerySinkConfig.USE_STORAGE_WRITE_API_CONFIG)).thenReturn(true);
 
+        String intermediateTable = "intermediate_topic";
+        TableId intermediateTableId = TableId.of(DATASET, intermediateTable);
+        when(mockMergeBatches.intermediateTableFor(BASE_TABLE_ID)).thenReturn(intermediateTableId);
+
         recordTableResolver = new RecordTableResolver(mockConfig, mockMergeBatches, mockBigQuery);
         TableId tableId = recordTableResolver.getRecordTable(mockRecord).getFullTableId();
 
         assertEquals(TOPIC, tableId.getTable());
-        assertEquals(PROJECT, tableId.getProject());
     }
 
     @Test
-    public void testResolveTableWithStorageApiIgnoresUpsertDelete() {
-        when(mockConfig.getBoolean(BigQuerySinkConfig.UPSERT_ENABLED_CONFIG)).thenReturn(true);
+    public void testPartitionDecoratorIgnoredWhenBatchModeEnabled() {
+        when(mockConfig.getBoolean(BigQuerySinkConfig.BIGQUERY_PARTITION_DECORATOR_CONFIG)).thenReturn(true);
+        when(mockConfig.getBoolean(BigQuerySinkConfig.ENABLE_BATCH_MODE_CONFIG)).thenReturn(true);
         when(mockConfig.getBoolean(BigQuerySinkConfig.USE_STORAGE_WRITE_API_CONFIG)).thenReturn(true);
 
-        String intermediateTable = "intermediate_topic";
-        TableId intermediateTableId = TableId.of(DATASET, intermediateTable);
-        when(mockMergeBatches.intermediateTableFor(BASE_TABLE_ID)).thenReturn(intermediateTableId);
+        mockTableWithPartitioning(null);
 
         recordTableResolver = new RecordTableResolver(mockConfig, mockMergeBatches, mockBigQuery);
         TableId tableId = recordTableResolver.getRecordTable(mockRecord).getFullTableId();
-
         assertEquals(TOPIC, tableId.getTable());
     }