docs(samples): add omni query on external table aws

samples/snippets/src/main/java/com/example/bigquery/CreateExternalTableAws.java

@@ -16,40 +16,54 @@
 
 package com.example.bigquery;
 
+// [START bigquery_omni_create_external_table]
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.BigQueryOptions;
 import com.google.cloud.bigquery.CsvOptions;
 import com.google.cloud.bigquery.ExternalTableDefinition;
+import com.google.cloud.bigquery.Field;
+import com.google.cloud.bigquery.Schema;
+import com.google.cloud.bigquery.StandardSQLTypeName;
 import com.google.cloud.bigquery.TableId;
 import com.google.cloud.bigquery.TableInfo;
 
+// Sample to create an external aws table
 public class CreateExternalTableAws {
 
   public static void main(String[] args) {
     // TODO(developer): Replace these variables before running the sample.
+    String projectId = "MY_PROJECT_ID";
     String datasetName = "MY_DATASET_NAME";
     String tableName = "MY_TABLE_NAME";
     // Create a aws connection
     // projects/{project_id}/locations/{location_id}/connections/{connection_id}
     String connectionId = "MY_CONNECTION_NAME";
     String sourceUri = "s3://your-bucket-name/";
     CsvOptions options = CsvOptions.newBuilder().setSkipLeadingRows(1).build();
+    Schema schema =
+        Schema.of(
+            Field.of("name", StandardSQLTypeName.STRING),
+            Field.of("post_abbr", StandardSQLTypeName.STRING));
     ExternalTableDefinition externalTableDefinition =
         ExternalTableDefinition.newBuilder(sourceUri, options)
             .setConnectionId(connectionId)
+            .setSchema(schema)
             .build();
-    createExternalTableAws(datasetName, tableName, externalTableDefinition);
+    createExternalTableAws(projectId, datasetName, tableName, externalTableDefinition);
   }
 
   public static void createExternalTableAws(
-      String datasetName, String tableName, ExternalTableDefinition externalTableDefinition) {
+      String projectId,
+      String datasetName,
+      String tableName,
+      ExternalTableDefinition externalTableDefinition) {
     try {
       // Initialize client that will be used to send requests. This client only needs to be created
       // once, and can be reused for multiple requests.
       BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();
 
-      TableId tableId = TableId.of(datasetName, tableName);
+      TableId tableId = TableId.of(projectId, datasetName, tableName);
       TableInfo tableInfo = TableInfo.newBuilder(tableId, externalTableDefinition).build();
 
       bigquery.create(tableInfo);
@@ -59,3 +73,4 @@ public static void createExternalTableAws(
     }
   }
 }
+// [END bigquery_omni_create_external_table]

samples/snippets/src/main/java/com/example/bigquery/QueryExternalTableAws.java

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * 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.example.bigquery;
+
+// [START bigquery_omni_query_external_aws_s3]
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.BigQueryException;
+import com.google.cloud.bigquery.BigQueryOptions;
+import com.google.cloud.bigquery.CsvOptions;
+import com.google.cloud.bigquery.ExternalTableDefinition;
+import com.google.cloud.bigquery.Field;
+import com.google.cloud.bigquery.QueryJobConfiguration;
+import com.google.cloud.bigquery.Schema;
+import com.google.cloud.bigquery.StandardSQLTypeName;
+import com.google.cloud.bigquery.TableId;
+import com.google.cloud.bigquery.TableInfo;
+import com.google.cloud.bigquery.TableResult;
+
+// Sample to queries an external data source aws s3 using a permanent table
+public class QueryExternalTableAws {
+
+  public static void main(String[] args) {
+    // TODO(developer): Replace these variables before running the sample.
+    String projectId = "MY_PROJECT_ID";
+    String datasetName = "MY_DATASET_NAME";
+    String tableName = "MY_TABLE_NAME";
+    // Create a aws connection
+    // projects/{project_id}/locations/{location_id}/connections/{connection_id}
+    String connectionId = "MY_CONNECTION_NAME";
+    String sourceUri = "s3://your-bucket-name/";
+    CsvOptions options = CsvOptions.newBuilder().setSkipLeadingRows(1).build();
+    Schema schema =
+        Schema.of(
+            Field.of("name", StandardSQLTypeName.STRING),
+            Field.of("post_abbr", StandardSQLTypeName.STRING));
+    String query =
+        String.format(
+            "SELECT * FROM s%:%s.%s WHERE name LIKE 'W%%'", projectId, datasetName, tableName);
+    ExternalTableDefinition externalTable =
+        ExternalTableDefinition.newBuilder(sourceUri, options)
+            .setConnectionId(connectionId)
+            .setSchema(schema)
+            .build();
+    queryExternalTableAws(projectId, datasetName, tableName, externalTable, query);
+  }
+
+  public static void queryExternalTableAws(
+      String projectId,
+      String datasetName,
+      String tableName,
+      ExternalTableDefinition externalTable,
+      String query) {
+    try {
+      // Initialize client that will be used to send requests. This client only needs to be created
+      // once, and can be reused for multiple requests.
+      BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();
+
+      TableId tableId = TableId.of(projectId, datasetName, tableName);
+      // Create a permanent table linked to the Aws file
+      bigquery.create(TableInfo.of(tableId, externalTable));
+
+      // Example query to find states starting with 'W'
+      TableResult results = bigquery.query(QueryJobConfiguration.of(query));
+
+      results
+          .iterateAll()
+          .forEach(row -> row.forEach(val -> System.out.printf("%s,", val.toString())));
+
+      System.out.println("Query on aws external permanent table performed successfully.");
+    } catch (BigQueryException | InterruptedException e) {
+      System.out.println("Query not performed \n" + e.toString());
+    }
+  }
+}
+// [END bigquery_omni_query_external_aws_s3]

samples/snippets/src/test/java/com/example/bigquery/CreateExternalTableAwsIT.java

@@ -131,7 +131,8 @@ public void testCreateExternalTableAws() {
             .setConnectionId(connectionName)
             .setSchema(schema)
             .build();
-    CreateExternalTableAws.createExternalTableAws(datasetName, tableName, externalTable);
+    CreateExternalTableAws.createExternalTableAws(
+        PROJECT_ID, datasetName, tableName, externalTable);
     assertThat(bout.toString()).contains("Aws external table created successfully");
   }
 }

samples/snippets/src/test/java/com/example/bigquery/QueryExternalTableAwsIT.java

@@ -0,0 +1,150 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * 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.example.bigquery;
+
+import static com.google.common.truth.Truth.assertThat;
+import static junit.framework.TestCase.assertNotNull;
+
+import com.google.cloud.bigquery.CsvOptions;
+import com.google.cloud.bigquery.ExternalTableDefinition;
+import com.google.cloud.bigquery.Field;
+import com.google.cloud.bigquery.Schema;
+import com.google.cloud.bigquery.StandardSQLTypeName;
+import com.google.cloud.bigquery.connection.v1.AwsCrossAccountRole;
+import com.google.cloud.bigquery.connection.v1.AwsProperties;
+import com.google.cloud.bigquery.connection.v1.Connection;
+import com.google.cloud.bigquery.connection.v1.CreateConnectionRequest;
+import com.google.cloud.bigquery.connection.v1.DeleteConnectionRequest;
+import com.google.cloud.bigquery.connection.v1.LocationName;
+import com.google.cloud.bigqueryconnection.v1.ConnectionServiceClient;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.util.UUID;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+public class QueryExternalTableAwsIT {
+
+  private static final String ID = UUID.randomUUID().toString().substring(0, 8);
+  private static final String LOCATION = "aws-us-east-1";
+  private final Logger log = Logger.getLogger(this.getClass().getName());
+  private String datasetName;
+  private String tableName;
+  private String connectionName;
+  private ByteArrayOutputStream bout;
+  private PrintStream out;
+  private PrintStream originalPrintStream;
+
+  private static final String PROJECT_ID = requireEnvVar("OMNI_PROJECT_ID");
+  private static final String AWS_ACCOUNT_ID = requireEnvVar("AWS_ACCOUNT_ID");
+  private static final String AWS_ROLE_ID = requireEnvVar("AWS_ROLE_ID");
+
+  private static String requireEnvVar(String varName) {
+    String value = System.getenv(varName);
+    assertNotNull(
+        "Environment variable " + varName + " is required to perform these tests.",
+        System.getenv(varName));
+    return value;
+  }
+
+  @BeforeClass
+  public static void checkRequirements() {
+    requireEnvVar("OMNI_PROJECT_ID");
+    requireEnvVar("AWS_ACCOUNT_ID");
+    requireEnvVar("AWS_ROLE_ID");
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    datasetName = "QUERY_EXTERNAL_TABLE_AWS_TEST_" + ID;
+    tableName = "QUERY_EXTERNAL_TABLE_AWS_TEST_" + ID;
+    connectionName = "QUERY_EXTERNAL_TABLE_AWS_TEST_" + ID;
+    bout = new ByteArrayOutputStream();
+    out = new PrintStream(bout);
+    originalPrintStream = System.out;
+    System.setOut(out);
+    // create a temporary aws connection
+    try (ConnectionServiceClient client = ConnectionServiceClient.create()) {
+      LocationName parent = LocationName.of(PROJECT_ID, LOCATION);
+      String iamRoleId = String.format("arn:aws:iam::%s:role/%s", AWS_ACCOUNT_ID, AWS_ROLE_ID);
+      AwsCrossAccountRole role = AwsCrossAccountRole.newBuilder().setIamRoleId(iamRoleId).build();
+      AwsProperties awsProperties = AwsProperties.newBuilder().setCrossAccountRole(role).build();
+      Connection connection = Connection.newBuilder().setAws(awsProperties).build();
+      CreateConnectionRequest request =
+          CreateConnectionRequest.newBuilder()
+              .setParent(parent.toString())
+              .setConnection(connection)
+              .setConnectionId(connectionName)
+              .build();
+      Connection response = client.createConnection(request);
+      connectionName = response.getName();
+      AwsCrossAccountRole accountRole = response.getAws().getCrossAccountRole();
+      System.out.println(
+          "Aws connection created successfully : Aws userId :"
+              + accountRole.getIamUserId()
+              + " Aws externalId :"
+              + accountRole.getExternalId());
+    }
+    // create a temporary dataset
+    CreateDatasetAws.createDatasetAws(PROJECT_ID, datasetName, LOCATION);
+  }
+
+  @After
+  public void tearDown() throws IOException {
+    // delete a temporary aws connection
+    try (ConnectionServiceClient client = ConnectionServiceClient.create()) {
+      DeleteConnectionRequest request =
+          DeleteConnectionRequest.newBuilder().setName(connectionName).build();
+      client.deleteConnection(request);
+      System.out.println("Connection deleted successfully");
+    }
+    // Clean up
+    DeleteTable.deleteTable(datasetName, tableName);
+    DeleteDataset.deleteDataset(PROJECT_ID, datasetName);
+    // restores print statements in the original method
+    System.out.flush();
+    System.setOut(originalPrintStream);
+    log.log(Level.INFO, bout.toString());
+  }
+
+  @Test
+  public void testQueryExternalTableAws() {
+    String sourceUri = "s3://cloud-samples-tests/us-states.csv";
+    Schema schema =
+        Schema.of(
+            Field.of("name", StandardSQLTypeName.STRING),
+            Field.of("post_abbr", StandardSQLTypeName.STRING));
+    CsvOptions options = CsvOptions.newBuilder().setSkipLeadingRows(1).build();
+    ExternalTableDefinition externalTable =
+        ExternalTableDefinition.newBuilder(sourceUri, options)
+            .setConnectionId(connectionName)
+            .setSchema(schema)
+            .build();
+    String query =
+        String.format(
+            "SELECT * FROM %s:%s.%s WHERE name LIKE 'W%%'", PROJECT_ID, datasetName, tableName);
+    QueryExternalTableAws.queryExternalTableAws(
+        PROJECT_ID, datasetName, tableName, externalTable, query);
+    assertThat(bout.toString())
+        .contains("Query on aws external permanent table performed successfully.");
+  }
+}