github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/io/bigqueryio/helper_test.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //    http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package bigqueryio
    17  
    18  import (
    19  	"context"
    20  	"fmt"
    21  	"os/exec"
    22  	"testing"
    23  
    24  	"cloud.google.com/go/bigquery"
    25  	"github.com/apache/beam/sdks/v2/go/test/integration"
    26  )
    27  
    28  // newTempTable creates a new BigQuery table using BigQuery's Data Definition Language (DDL) and the
    29  // "bq query" console command. Reference: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language
    30  // The tables created are set to expire after a day.
    31  //
    32  // newTable takes the name of a BigQuery dataset and a DDL schema for the data,
    33  // and generates that table with a unique suffix and an expiration time of a day later.
    34  func newTempTable(t *testing.T, name string, schema string) {
    35  	t.Helper()
    36  
    37  	query := fmt.Sprintf("CREATE TABLE `%s`(%s) OPTIONS(expiration_timestamp=TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL 1 DAY))", name, schema)
    38  	cmd := exec.Command("bq", "query", "--use_legacy_sql=false", query)
    39  	_, err := cmd.CombinedOutput()
    40  	if err != nil {
    41  		t.Fatalf("error creating BigQuery table: %v", err)
    42  	}
    43  	t.Logf("Created BigQuery table %v", name)
    44  }
    45  
    46  // deleteTable deletes a BigQuery table using BigQuery's Data Definition Language (DDL) and the
    47  // "bq query" console command. Reference: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language
    48  func deleteTempTable(t *testing.T, table string) {
    49  	t.Helper()
    50  
    51  	query := fmt.Sprintf("DROP TABLE IF EXISTS `%s`", table)
    52  	cmd := exec.Command("bq", "query", "--use_legacy_sql=false", query)
    53  	t.Logf("Deleting BigQuery table %v", table)
    54  	_, err := cmd.CombinedOutput()
    55  	if err != nil {
    56  		t.Logf("Error deleting BigQuery table: %v", err)
    57  	}
    58  }
    59  
    60  func checkTableExistsAndNonEmpty(ctx context.Context, t *testing.T, project, tableID string) {
    61  	t.Helper()
    62  
    63  	client, err := bigquery.NewClient(ctx, project)
    64  	if err != nil {
    65  		t.Fatalf("error creating BigQuery client: %v", err)
    66  	}
    67  	defer client.Close()
    68  
    69  	tableRef := client.Dataset(*integration.BigQueryDataset).Table(tableID)
    70  	metadata, err := tableRef.Metadata(ctx)
    71  	if err != nil {
    72  		t.Fatalf("unable to find table: %v", err)
    73  	}
    74  	streamingBuffer := metadata.StreamingBuffer
    75  	if streamingBuffer == nil {
    76  		t.Fatalf("there's no streaming buffer for the table")
    77  	}
    78  	if streamingBuffer.EstimatedRows != inputSize {
    79  		t.Fatalf("streamingBuffer.EstimatedRows = %v, want %v", streamingBuffer.EstimatedRows, inputSize)
    80  	}
    81  }