extract_test.go

Documentation: cloud.google.com/go/bigquery

     1  // Copyright 2015 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package bigquery
    16  
    17  import (
    18  	"testing"
    19  	"time"
    20  
    21  	"cloud.google.com/go/internal/testutil"
    22  	"github.com/google/go-cmp/cmp"
    23  	bq "google.golang.org/api/bigquery/v2"
    24  )
    25  
    26  func defaultExtractJob() *bq.Job {
    27  	return &bq.Job{
    28  		JobReference: &bq.JobReference{JobId: "RANDOM", ProjectId: "client-project-id"},
    29  		Configuration: &bq.JobConfiguration{
    30  			Extract: &bq.JobConfigurationExtract{
    31  				SourceTable: &bq.TableReference{
    32  					ProjectId: "client-project-id",
    33  					DatasetId: "dataset-id",
    34  					TableId:   "table-id",
    35  				},
    36  				DestinationUris: []string{"uri"},
    37  			},
    38  		},
    39  	}
    40  }
    41  
    42  func defaultExtractModelJob() *bq.Job {
    43  	return &bq.Job{
    44  		JobReference: &bq.JobReference{JobId: "RANDOM", ProjectId: "client-project-id"},
    45  		Configuration: &bq.JobConfiguration{
    46  			Extract: &bq.JobConfigurationExtract{
    47  				SourceModel: &bq.ModelReference{
    48  					ProjectId: "client-project-id",
    49  					DatasetId: "dataset-id",
    50  					ModelId:   "model-id",
    51  				},
    52  				DestinationUris: []string{"uri"},
    53  			},
    54  		},
    55  	}
    56  }
    57  
    58  func defaultGCS() *GCSReference {
    59  	return &GCSReference{
    60  		URIs: []string{"uri"},
    61  	}
    62  }
    63  
    64  func TestExtract(t *testing.T) {
    65  	defer fixRandomID("RANDOM")()
    66  	c := &Client{
    67  		projectID: "client-project-id",
    68  	}
    69  
    70  	testCases := []struct {
    71  		dst    *GCSReference
    72  		src    *Table
    73  		config ExtractConfig
    74  		want   *bq.Job
    75  	}{
    76  		{
    77  			dst:  defaultGCS(),
    78  			src:  c.Dataset("dataset-id").Table("table-id"),
    79  			want: defaultExtractJob(),
    80  		},
    81  		{
    82  			dst: defaultGCS(),
    83  			src: c.Dataset("dataset-id").Table("table-id"),
    84  			config: ExtractConfig{
    85  				DisableHeader: true,
    86  				Labels:        map[string]string{"a": "b"},
    87  				JobTimeout:    8 * time.Second,
    88  			},
    89  			want: func() *bq.Job {
    90  				j := defaultExtractJob()
    91  				j.Configuration.Labels = map[string]string{"a": "b"}
    92  				j.Configuration.JobTimeoutMs = 8000
    93  				f := false
    94  				j.Configuration.Extract.PrintHeader = &f
    95  				return j
    96  			}(),
    97  		},
    98  		{
    99  			dst: func() *GCSReference {
   100  				g := NewGCSReference("uri")
   101  				g.Compression = Gzip
   102  				g.DestinationFormat = JSON
   103  				g.FieldDelimiter = "\t"
   104  				return g
   105  			}(),
   106  			src: c.Dataset("dataset-id").Table("table-id"),
   107  			want: func() *bq.Job {
   108  				j := defaultExtractJob()
   109  				j.Configuration.Extract.Compression = "GZIP"
   110  				j.Configuration.Extract.DestinationFormat = "NEWLINE_DELIMITED_JSON"
   111  				j.Configuration.Extract.FieldDelimiter = "\t"
   112  				return j
   113  			}(),
   114  		},
   115  		{
   116  			dst: func() *GCSReference {
   117  				g := NewGCSReference("uri")
   118  				g.DestinationFormat = Avro
   119  				g.Compression = Snappy
   120  				return g
   121  			}(),
   122  			src: c.Dataset("dataset-id").Table("table-id"),
   123  			config: ExtractConfig{
   124  				UseAvroLogicalTypes: true,
   125  			},
   126  			want: func() *bq.Job {
   127  				j := defaultExtractJob()
   128  				j.Configuration.Extract.UseAvroLogicalTypes = true
   129  				j.Configuration.Extract.DestinationFormat = "AVRO"
   130  				j.Configuration.Extract.Compression = "SNAPPY"
   131  				return j
   132  			}(),
   133  		},
   134  	}
   135  
   136  	for i, tc := range testCases {
   137  		ext := tc.src.ExtractorTo(tc.dst)
   138  		tc.config.Src = ext.Src
   139  		tc.config.Dst = ext.Dst
   140  		ext.ExtractConfig = tc.config
   141  		got := ext.newJob()
   142  		checkJob(t, i, got, tc.want)
   143  
   144  		jc, err := bqToJobConfig(got.Configuration, c)
   145  		if err != nil {
   146  			t.Fatalf("#%d: %v", i, err)
   147  		}
   148  		diff := testutil.Diff(jc, &ext.ExtractConfig,
   149  			cmp.AllowUnexported(Table{}, Client{}))
   150  		if diff != "" {
   151  			t.Errorf("#%d: (got=-, want=+:\n%s", i, diff)
   152  		}
   153  	}
   154  }
   155  
   156  func TestExtractModel(t *testing.T) {
   157  	defer fixRandomID("RANDOM")()
   158  	c := &Client{
   159  		projectID: "client-project-id",
   160  	}
   161  
   162  	testCases := []struct {
   163  		dst      *GCSReference
   164  		srcModel *Model
   165  		config   ExtractConfig
   166  		want     *bq.Job
   167  	}{
   168  		{
   169  			dst:      defaultGCS(),
   170  			srcModel: c.Dataset("dataset-id").Model("model-id"),
   171  			want:     defaultExtractModelJob(),
   172  		},
   173  		{
   174  			dst: &GCSReference{
   175  				URIs:              []string{"uri"},
   176  				DestinationFormat: TFSavedModel,
   177  			},
   178  			srcModel: c.Dataset("dataset-id").Model("model-id"),
   179  			want: func() *bq.Job {
   180  				j := defaultExtractModelJob()
   181  				j.Configuration.Extract.DestinationFormat = string(TFSavedModel)
   182  				return j
   183  			}(),
   184  		},
   185  	}
   186  
   187  	for i, tc := range testCases {
   188  		ext := tc.srcModel.ExtractorTo(tc.dst)
   189  		tc.config.SrcModel = ext.SrcModel
   190  		tc.config.Dst = ext.Dst
   191  		ext.ExtractConfig = tc.config
   192  		got := ext.newJob()
   193  		checkJob(t, i, got, tc.want)
   194  
   195  		jc, err := bqToJobConfig(got.Configuration, c)
   196  		if err != nil {
   197  			t.Fatalf("#%d: %v", i, err)
   198  		}
   199  		diff := testutil.Diff(jc, &ext.ExtractConfig,
   200  			cmp.AllowUnexported(Model{}, Client{}))
   201  		if diff != "" {
   202  			t.Errorf("#%d: (got=-, want=+:\n%s", i, diff)
   203  		}
   204  	}
   205  }
   206
View as plain text