Merge branch 'main' into jdbc/refactor-connection-properties

Author: logachev

google-cloud-bigquery-jdbc/pom.xml

@@ -48,6 +48,18 @@
         <version>3.5.2</version>
         <configuration>
           <skip>${skipSurefire}</skip>
+          <systemPropertyVariables>
+            <JDBC_TESTS>true</JDBC_TESTS>
+          </systemPropertyVariables>
+        </configuration>
+      </plugin>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-failsafe-plugin</artifactId>
+        <configuration>
+          <systemPropertyVariables>
+            <JDBC_TESTS>true</JDBC_TESTS>
+          </systemPropertyVariables>
         </configuration>
       </plugin>
       <plugin>

google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java

@@ -32,6 +32,7 @@
 import com.google.cloud.bigquery.Job;
 import com.google.cloud.bigquery.JobInfo;
 import com.google.cloud.bigquery.QueryJobConfiguration;
+import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.cloud.bigquery.exception.BigQueryJdbcException;
 import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
 import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException;
@@ -934,8 +935,12 @@ private BigQuery getBigQueryConnection() {
       bigQueryOptions.setTransportOptions(this.httpTransportOptions);
     }
 
-    BigQueryOptions options = bigQueryOptions.setHeaderProvider(this.headerProvider).build();
-    options.setQueryPreviewEnabled(String.valueOf(this.useStatelessQueryMode));
+    BigQueryOptions options = bigQueryOptions.setHeaderProvider(HEADER_PROVIDER).build();
+    options.setDefaultJobCreationMode(
+        this.useStatelessQueryMode
+            ? JobCreationMode.JOB_CREATION_OPTIONAL
+            : JobCreationMode.JOB_CREATION_REQUIRED);
+
     return options.getService();
   }
 

google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionProperty.java

@@ -17,12 +17,14 @@
 package com.google.cloud.bigquery.jdbc;
 
 import java.util.List;
+import java.util.function.Supplier;
 
 class BigQueryConnectionProperty {
 
   private final String name;
   private final String description;
   private final String defaultValue;
+  private final Supplier<String> defaultValueSupplier;
   private final List<String> validValues;
 
   public String getName() {
@@ -34,6 +36,9 @@ public String getDescription() {
   }
 
   public String getDefaultValue() {
+    if (defaultValueSupplier != null) {
+      return defaultValueSupplier.get();
+    }
     return defaultValue;
   }
 
@@ -43,6 +48,7 @@ public List<String> getValidValues() {
 
   BigQueryConnectionProperty(Builder builder) {
     this.name = builder.name;
+    this.defaultValueSupplier = builder.defaultValueSupplier;
     this.defaultValue = builder.defaultValue;
     this.description = builder.description;
     this.validValues = builder.validValues;
@@ -79,6 +85,7 @@ static final class Builder {
     private String name;
     private String description;
     private String defaultValue;
+    private Supplier<String> defaultValueSupplier = null;
     private List<String> validValues;
 
     private Builder(BigQueryConnectionProperty bigQueryConnectionProperty) {
@@ -105,6 +112,11 @@ Builder setDefaultValue(String defaultValue) {
       return this;
     }
 
+    Builder setLazyDefaultValue(Supplier<String> defaultValueSupplier) {
+      this.defaultValueSupplier = defaultValueSupplier;
+      return this;
+    }
+
     Builder setValidValues(List<String> validValues) {
       this.validValues = validValues;
       return this;

google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java

@@ -38,6 +38,7 @@
 import java.awt.Desktop;
 import java.io.BufferedReader;
 import java.io.ByteArrayInputStream;
+import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
@@ -299,11 +300,11 @@ private static boolean isFileExists(String filename) {
     }
   }
 
-  private static boolean isJson(String value) {
+  private static boolean isJson(byte[] value) {
     try {
       // This is done this way to ensure strict Json parsing
       // https://github.com/google/gson/issues/1208#issuecomment-2120764686
-      InputStream stream = new ByteArrayInputStream(value.getBytes());
+      InputStream stream = new ByteArrayInputStream(value);
       InputStreamReader reader = new InputStreamReader(stream);
       JsonReader jsonReader = new JsonReader(reader);
       jsonReader.setStrictness(Strictness.STRICT);
@@ -333,17 +334,24 @@ private static GoogleCredentials getGoogleServiceAccountCredentials(
 
       final String keyPath = pvtKeyPath != null ? pvtKeyPath : pvtKey;
       PrivateKey key = null;
-      InputStream stream = null;
+      byte[] keyBytes = pvtKey != null ? pvtKey.getBytes() : null;
 
       if (isFileExists(keyPath)) {
-        key = privateKeyFromP12File(keyPath, p12Password);
-        if (key == null) {
-          stream = Files.newInputStream(Paths.get(keyPath));
+        try (InputStream stream = new FileInputStream(keyPath)) {
+          int bufferSize = 1024 * 1024;
+          byte[] buffer = new byte[bufferSize];
+          stream.read(buffer, 0, bufferSize);
+          keyBytes = buffer;
         }
-      } else if (isJson(pvtKey)) {
-        stream = new ByteArrayInputStream(pvtKey.getBytes());
+      }
+
+      InputStream stream = null;
+      if (isJson(keyBytes)) {
+        stream = new ByteArrayInputStream(keyBytes);
       } else if (pvtKey != null) {
         key = privateKeyFromPkcs8(pvtKey);
+      } else if (keyBytes != null) {
+        key = privateKeyFromP12Bytes(keyBytes, p12Password);
       }
 
       if (stream != null) {
@@ -672,9 +680,9 @@ private static GoogleCredentials getServiceAccountImpersonatedCredentials(
         impersonationLifetimeInt);
   }
 
-  static PrivateKey privateKeyFromP12File(String privateKeyFile, String password) {
+  static PrivateKey privateKeyFromP12Bytes(byte[] privateKey, String password) {
     try {
-      InputStream stream = Files.newInputStream(Paths.get(privateKeyFile));
+      InputStream stream = new ByteArrayInputStream(privateKey);
       return SecurityUtils.loadPrivateKeyFromKeyStore(
           SecurityUtils.getPkcs12KeyStore(), stream, "notasecret", "privatekey", password);
     } catch (IOException | GeneralSecurityException e) {

google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java

@@ -43,6 +43,7 @@ class BigQueryJdbcRootLogger {
   private static final Logger logger = Logger.getLogger("com.google.cloud.bigquery");
 
   private static final Logger storageLogger = Logger.getLogger("com.google.cloud.bigquery.storage");
+  private static final boolean isTest = Boolean.getBoolean("JDBC_TESTS");
 
   private static Handler fileHandler = null;
   private static Path currentLogPath = null;
@@ -51,6 +52,12 @@ class BigQueryJdbcRootLogger {
   static {
     logger.setUseParentHandlers(false);
     storageLogger.setUseParentHandlers(true);
+    if (isTest) {
+      ConsoleHandler consoleHandler = new ConsoleHandler();
+      consoleHandler.setLevel(Level.SEVERE);
+      consoleHandler.setFormatter(getFormatter());
+      logger.addHandler(consoleHandler);
+    }
   }
 
   public static Formatter getFormatter() {
@@ -111,13 +118,12 @@ private static void setHandler() throws IOException {
     // If File handler exists, use it. Else create new one.
     for (Handler h : logger.getHandlers()) {
       if (h instanceof ConsoleHandler) {
-        h.close();
-        logger.removeHandler(h);
-        break;
-      }
-      if (h instanceof FileHandler) {
+        if (!isTest) {
+          h.close();
+          logger.removeHandler(h);
+        }
+      } else if (h instanceof FileHandler) {
         fileHandler = h;
-        break;
       }
     }
 

google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java

@@ -32,6 +32,7 @@
 import java.util.Map.Entry;
 import java.util.Properties;
 import java.util.Set;
+import java.util.function.Supplier;
 import java.util.logging.Level;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -358,7 +359,7 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
                   BigQueryConnectionProperty.newBuilder()
                       .setName(PROJECT_ID_PROPERTY_NAME)
                       .setDescription("A globally unique identifier for your project.")
-                      .setDefaultValue(BigQueryOptions.getDefaultProjectId())
+                      .setLazyDefaultValue(() -> BigQueryOptions.getDefaultProjectId())
                       .build(),
                   BigQueryConnectionProperty.newBuilder()
                       .setName(LOG_PATH_PROPERTY_NAME)

google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java

@@ -21,6 +21,8 @@
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.rpc.HeaderProvider;
 import com.google.api.gax.rpc.TransportChannelProvider;
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.cloud.bigquery.exception.BigQueryJdbcException;
 import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
 import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
@@ -35,6 +37,9 @@ public class BigQueryConnectionTest {
 
   private static final String DEFAULT_VERSION = "0.0.0";
   private static final String DEFAULT_JDBC_TOKEN_VALUE = "Google-BigQuery-JDBC-Driver";
+  private static final String BASE_URL =
+      "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
+          + "OAuthType=2;OAuthAccessToken=redacted;ProjectId=project;";
   private String expectedVersion;
 
   @Before
@@ -378,4 +383,34 @@ public void testBigQueryReadClientKeepAliveSettings() throws SQLException, IOExc
       assertTrue(grpcProvider.getKeepAliveWithoutCalls());
     }
   }
+
+  @Test
+  public void testBigQueryJobCreationMode_required() throws Exception {
+    String url = BASE_URL + "JobCreationMode=1;";
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      BigQuery bq = connection.getBigQuery();
+      assertEquals(
+          bq.getOptions().getDefaultJobCreationMode(), JobCreationMode.JOB_CREATION_REQUIRED);
+    }
+  }
+
+  @Test
+  public void testBigQueryJobCreationMode_optional() throws Exception {
+    String url = BASE_URL + "JobCreationMode=2;";
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      BigQuery bq = connection.getBigQuery();
+      assertEquals(
+          bq.getOptions().getDefaultJobCreationMode(), JobCreationMode.JOB_CREATION_OPTIONAL);
+    }
+  }
+
+  @Test
+  public void testBigQueryJobCreationMode_default() throws Exception {
+    String url = BASE_URL;
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      BigQuery bq = connection.getBigQuery();
+      assertEquals(
+          bq.getOptions().getDefaultJobCreationMode(), JobCreationMode.JOB_CREATION_OPTIONAL);
+    }
+  }
 }

google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtilityTest.java

@@ -32,6 +32,7 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URL;
+import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.security.PrivateKey;
 import java.util.Collections;
@@ -501,31 +502,25 @@ public void testPrivateKeyFromPkcs8_wrong() {
   // keytool -genkey -alias privatekey -keyalg RSA -keysize 2048 -storepass notasecret \
   //   -keypass  notasecret -storetype pkcs12 -keystore ./fake.p12
   @Test
-  public void testPrivateKeyFromP12File() {
+  public void testPrivateKeyFromP12Bytes() {
     URL resource = BigQueryJdbcOAuthUtilityTest.class.getResource("/fake.p12");
     try {
       PrivateKey pk =
-          BigQueryJdbcOAuthUtility.privateKeyFromP12File(
-              Paths.get(resource.toURI()).toAbsolutePath().toString(), "notasecret");
+          BigQueryJdbcOAuthUtility.privateKeyFromP12Bytes(
+              Files.readAllBytes(Paths.get(resource.toURI())), "notasecret");
       assertNotNull(pk);
     } catch (Exception e) {
       assertTrue(false);
     }
   }
 
   @Test
-  public void testPrivateKeyFromP12File_missing_file() {
-    PrivateKey pk = BigQueryJdbcOAuthUtility.privateKeyFromP12File("", "");
-    assertNull(pk);
-  }
-
-  @Test
-  public void testPrivateKeyFromP12File_wrong_password() {
+  public void testPrivateKeyFromP12Bytes_wrong_password() {
     URL resource = BigQueryJdbcOAuthUtilityTest.class.getResource("/fake.p12");
     try {
       PrivateKey pk =
-          BigQueryJdbcOAuthUtility.privateKeyFromP12File(
-              Paths.get(resource.toURI()).toAbsolutePath().toString(), "fake");
+          BigQueryJdbcOAuthUtility.privateKeyFromP12Bytes(
+              Files.readAllBytes(Paths.get(resource.toURI())), "fake");
       assertNull(pk);
     } catch (Exception e) {
       assertTrue(false);