From 68528ef3b91aec1321d8ef277b85cb70ab6e0a32 Mon Sep 17 00:00:00 2001
From: chngpe <chngpe@amazon.com>
Date: Tue, 21 May 2024 22:44:23 +0000
Subject: [PATCH] Revert "Com.google.cloud google cloud bigquery version update
 (#1975)"

This reverts commit ee6cba1a4e5e90d4103cdfadf88d0a10407229e7.
---
 athena-google-bigquery/pom.xml                |  2 +-
 .../bigquery/BigQueryMetadataHandlerTest.java | 25 +++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/athena-google-bigquery/pom.xml b/athena-google-bigquery/pom.xml
index 5b7043b86b..6bae32b005 100644
--- a/athena-google-bigquery/pom.xml
+++ b/athena-google-bigquery/pom.xml
@@ -55,7 +55,7 @@
         <dependency>
             <groupId>com.google.cloud</groupId>
             <artifactId>google-cloud-bigquery</artifactId>
-            <version>2.40.1</version>
+            <version>2.37.0</version>
             <exclusions>
                 <exclusion>
                     <groupId>com.google.api.grpc</groupId>
diff --git a/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryMetadataHandlerTest.java b/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryMetadataHandlerTest.java
index d3062d24d3..3b1dd1b044 100644
--- a/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryMetadataHandlerTest.java
+++ b/athena-google-bigquery/src/test/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryMetadataHandlerTest.java
@@ -33,15 +33,22 @@
 import com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest;
 import com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse;
 import com.amazonaws.athena.connector.lambda.security.FederatedIdentity;
+import com.google.api.gax.paging.Page;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.Dataset;
 import com.google.cloud.bigquery.DatasetId;
+import com.google.cloud.bigquery.Field;
+import com.google.cloud.bigquery.FieldList;
+import com.google.cloud.bigquery.FieldValue;
+import com.google.cloud.bigquery.FieldValueList;
 import com.google.cloud.bigquery.Job;
 import com.google.cloud.bigquery.JobStatus;
+import com.google.cloud.bigquery.LegacySQLTypeName;
 import com.google.cloud.bigquery.Schema;
 import com.google.cloud.bigquery.StandardTableDefinition;
 import com.google.cloud.bigquery.Table;
 import com.google.cloud.bigquery.TableId;
+import com.google.cloud.bigquery.TableResult;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -53,8 +60,10 @@
 import org.mockito.junit.MockitoJUnitRunner;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT;
@@ -219,10 +228,26 @@ public void testDoGetTable() throws IOException
     @Test
     public void testDoGetSplits() throws Exception
     {
+
+//        mockedStatic.when(() -> BigQueryUtils.fixCaseForDatasetName(any(String.class), any(String.class), any(BigQuery.class))).thenReturn("testDataset");
+//        mockedStatic.when(() -> BigQueryUtils.fixCaseForTableName(any(String.class), any(String.class), any(String.class), any(BigQuery.class))).thenReturn("testTable");
         BlockAllocator blockAllocator = new BlockAllocatorImpl();
         GetSplitsRequest request = new GetSplitsRequest(federatedIdentity,
                 QUERY_ID, CATALOG, TABLE_NAME,
                 mock(Block.class), Collections.<String>emptyList(), new Constraints(new HashMap<>(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT), null);
+        // added schema with integer column countCol
+        List<Field> testSchemaFields = Arrays.asList(Field.of("countCol", LegacySQLTypeName.INTEGER));
+        Schema tableSchema = Schema.of(testSchemaFields);
+
+        // mocked table row count as 15
+        List<FieldValue> bigQueryRowValue = Arrays.asList(FieldValue.of(FieldValue.Attribute.PRIMITIVE, "15"));
+        FieldValueList fieldValueList = FieldValueList.of(bigQueryRowValue,
+                FieldList.of(testSchemaFields));
+        List<FieldValueList> tableRows = Arrays.asList(fieldValueList);
+
+        Page<FieldValueList> pageNoSchema = new BigQueryPage<>(tableRows);
+        TableResult result = new TableResult(tableSchema, tableRows.size(), pageNoSchema);
+//        when(job.getQueryResults()).thenReturn(result);
 
         GetSplitsResponse response = bigQueryMetadataHandler.doGetSplits(blockAllocator, request);