Source file src/cloud.google.com/go/bigquery/query_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  	"github.com/google/go-cmp/cmp/cmpopts"
    24  	bq "google.golang.org/api/bigquery/v2"
    25  )
    26  
    27  func defaultQueryJob() *bq.Job {
    28  	pfalse := false
    29  	return &bq.Job{
    30  		JobReference: &bq.JobReference{JobId: "RANDOM", ProjectId: "client-project-id"},
    31  		Configuration: &bq.JobConfiguration{
    32  			Query: &bq.JobConfigurationQuery{
    33  				DestinationTable: &bq.TableReference{
    34  					ProjectId: "client-project-id",
    35  					DatasetId: "dataset-id",
    36  					TableId:   "table-id",
    37  				},
    38  				Query: "query string",
    39  				DefaultDataset: &bq.DatasetReference{
    40  					ProjectId: "def-project-id",
    41  					DatasetId: "def-dataset-id",
    42  				},
    43  				UseLegacySql: &pfalse,
    44  			},
    45  		},
    46  	}
    47  }
    48  
    49  var defaultQuery = &QueryConfig{
    50  	Q:                "query string",
    51  	DefaultProjectID: "def-project-id",
    52  	DefaultDatasetID: "def-dataset-id",
    53  }
    54  
    55  func TestQuery(t *testing.T) {
    56  	defer fixRandomID("RANDOM")()
    57  	c := &Client{
    58  		projectID: "client-project-id",
    59  	}
    60  	testCases := []struct {
    61  		dst         *Table
    62  		src         *QueryConfig
    63  		jobIDConfig JobIDConfig
    64  		want        *bq.Job
    65  	}{
    66  		{
    67  			dst:  c.Dataset("dataset-id").Table("table-id"),
    68  			src:  defaultQuery,
    69  			want: defaultQueryJob(),
    70  		},
    71  		{
    72  			dst: c.Dataset("dataset-id").Table("table-id"),
    73  			src: &QueryConfig{
    74  				Q:      "query string",
    75  				Labels: map[string]string{"a": "b"},
    76  				DryRun: true,
    77  			},
    78  			want: func() *bq.Job {
    79  				j := defaultQueryJob()
    80  				j.Configuration.Labels = map[string]string{"a": "b"}
    81  				j.Configuration.DryRun = true
    82  				j.Configuration.Query.DefaultDataset = nil
    83  				return j
    84  			}(),
    85  		},
    86  		{
    87  			dst:         c.Dataset("dataset-id").Table("table-id"),
    88  			jobIDConfig: JobIDConfig{JobID: "jobID", AddJobIDSuffix: true},
    89  			src:         &QueryConfig{Q: "query string"},
    90  			want: func() *bq.Job {
    91  				j := defaultQueryJob()
    92  				j.Configuration.Query.DefaultDataset = nil
    93  				j.JobReference.JobId = "jobID-RANDOM"
    94  				return j
    95  			}(),
    96  		},
    97  		{
    98  			dst:         c.DatasetInProject("another-project", "dataset-id").Table("table-id"),
    99  			jobIDConfig: JobIDConfig{JobID: "jobID", ProjectID: "another-project"},
   100  			src:         &QueryConfig{Q: "query string"},
   101  			want: func() *bq.Job {
   102  				j := defaultQueryJob()
   103  				j.Configuration.Query.DefaultDataset = nil
   104  				j.Configuration.Query.DestinationTable.ProjectId = "another-project"
   105  				j.JobReference.JobId = "jobID"
   106  				j.JobReference.ProjectId = "another-project"
   107  				return j
   108  			}(),
   109  		},
   110  		{
   111  			dst: &Table{},
   112  			src: defaultQuery,
   113  			want: func() *bq.Job {
   114  				j := defaultQueryJob()
   115  				j.Configuration.Query.DestinationTable = nil
   116  				return j
   117  			}(),
   118  		},
   119  		{
   120  			dst: c.Dataset("dataset-id").Table("table-id"),
   121  			src: &QueryConfig{
   122  				Q: "query string",
   123  				TableDefinitions: map[string]ExternalData{
   124  					"atable": func() *GCSReference {
   125  						g := NewGCSReference("uri")
   126  						g.AllowJaggedRows = true
   127  						g.AllowQuotedNewlines = true
   128  						g.Compression = Gzip
   129  						g.Encoding = UTF_8
   130  						g.FieldDelimiter = ";"
   131  						g.IgnoreUnknownValues = true
   132  						g.MaxBadRecords = 1
   133  						g.Quote = "'"
   134  						g.SkipLeadingRows = 2
   135  						g.Schema = Schema{{Name: "name", Type: StringFieldType}}
   136  						return g
   137  					}(),
   138  				},
   139  			},
   140  			want: func() *bq.Job {
   141  				j := defaultQueryJob()
   142  				j.Configuration.Query.DefaultDataset = nil
   143  				td := make(map[string]bq.ExternalDataConfiguration)
   144  				quote := "'"
   145  				td["atable"] = bq.ExternalDataConfiguration{
   146  					Compression:         "GZIP",
   147  					IgnoreUnknownValues: true,
   148  					MaxBadRecords:       1,
   149  					SourceFormat:        "CSV", // must be explicitly set.
   150  					SourceUris:          []string{"uri"},
   151  					CsvOptions: &bq.CsvOptions{
   152  						AllowJaggedRows:     true,
   153  						AllowQuotedNewlines: true,
   154  						Encoding:            "UTF-8",
   155  						FieldDelimiter:      ";",
   156  						SkipLeadingRows:     2,
   157  						Quote:               &quote,
   158  					},
   159  					Schema: &bq.TableSchema{
   160  						Fields: []*bq.TableFieldSchema{
   161  							{Name: "name", Type: "STRING"},
   162  						},
   163  					},
   164  				}
   165  				j.Configuration.Query.TableDefinitions = td
   166  				return j
   167  			}(),
   168  		},
   169  		{
   170  			dst: &Table{
   171  				ProjectID: "project-id",
   172  				DatasetID: "dataset-id",
   173  				TableID:   "table-id",
   174  			},
   175  			src: &QueryConfig{
   176  				Q:                 "query string",
   177  				DefaultProjectID:  "def-project-id",
   178  				DefaultDatasetID:  "def-dataset-id",
   179  				CreateDisposition: CreateNever,
   180  				WriteDisposition:  WriteTruncate,
   181  			},
   182  			want: func() *bq.Job {
   183  				j := defaultQueryJob()
   184  				j.Configuration.Query.DestinationTable.ProjectId = "project-id"
   185  				j.Configuration.Query.WriteDisposition = "WRITE_TRUNCATE"
   186  				j.Configuration.Query.CreateDisposition = "CREATE_NEVER"
   187  				return j
   188  			}(),
   189  		},
   190  		{
   191  			dst: c.Dataset("dataset-id").Table("table-id"),
   192  			src: &QueryConfig{
   193  				Q:                 "query string",
   194  				DefaultProjectID:  "def-project-id",
   195  				DefaultDatasetID:  "def-dataset-id",
   196  				DisableQueryCache: true,
   197  			},
   198  			want: func() *bq.Job {
   199  				j := defaultQueryJob()
   200  				f := false
   201  				j.Configuration.Query.UseQueryCache = &f
   202  				return j
   203  			}(),
   204  		},
   205  		{
   206  			dst: c.Dataset("dataset-id").Table("table-id"),
   207  			src: &QueryConfig{
   208  				Q:                 "query string",
   209  				DefaultProjectID:  "def-project-id",
   210  				DefaultDatasetID:  "def-dataset-id",
   211  				AllowLargeResults: true,
   212  			},
   213  			want: func() *bq.Job {
   214  				j := defaultQueryJob()
   215  				j.Configuration.Query.AllowLargeResults = true
   216  				return j
   217  			}(),
   218  		},
   219  		{
   220  			dst: c.Dataset("dataset-id").Table("table-id"),
   221  			src: &QueryConfig{
   222  				Q:                       "query string",
   223  				DefaultProjectID:        "def-project-id",
   224  				DefaultDatasetID:        "def-dataset-id",
   225  				DisableFlattenedResults: true,
   226  			},
   227  			want: func() *bq.Job {
   228  				j := defaultQueryJob()
   229  				f := false
   230  				j.Configuration.Query.FlattenResults = &f
   231  				j.Configuration.Query.AllowLargeResults = true
   232  				return j
   233  			}(),
   234  		},
   235  		{
   236  			dst: c.Dataset("dataset-id").Table("table-id"),
   237  			src: &QueryConfig{
   238  				Q:                "query string",
   239  				DefaultProjectID: "def-project-id",
   240  				DefaultDatasetID: "def-dataset-id",
   241  				Priority:         QueryPriority("low"),
   242  			},
   243  			want: func() *bq.Job {
   244  				j := defaultQueryJob()
   245  				j.Configuration.Query.Priority = "low"
   246  				return j
   247  			}(),
   248  		},
   249  		{
   250  			dst: c.Dataset("dataset-id").Table("table-id"),
   251  			src: &QueryConfig{
   252  				Q:                "query string",
   253  				DefaultProjectID: "def-project-id",
   254  				DefaultDatasetID: "def-dataset-id",
   255  				MaxBillingTier:   3,
   256  				MaxBytesBilled:   5,
   257  			},
   258  			want: func() *bq.Job {
   259  				j := defaultQueryJob()
   260  				tier := int64(3)
   261  				j.Configuration.Query.MaximumBillingTier = &tier
   262  				j.Configuration.Query.MaximumBytesBilled = 5
   263  				return j
   264  			}(),
   265  		},
   266  		{
   267  			dst: c.Dataset("dataset-id").Table("table-id"),
   268  			src: &QueryConfig{
   269  				Q:                "query string",
   270  				DefaultProjectID: "def-project-id",
   271  				DefaultDatasetID: "def-dataset-id",
   272  				UseStandardSQL:   true,
   273  			},
   274  			want: defaultQueryJob(),
   275  		},
   276  		{
   277  			dst: c.Dataset("dataset-id").Table("table-id"),
   278  			src: &QueryConfig{
   279  				Q:                "query string",
   280  				DefaultProjectID: "def-project-id",
   281  				DefaultDatasetID: "def-dataset-id",
   282  				UseLegacySQL:     true,
   283  			},
   284  			want: func() *bq.Job {
   285  				j := defaultQueryJob()
   286  				ptrue := true
   287  				j.Configuration.Query.UseLegacySql = &ptrue
   288  				j.Configuration.Query.ForceSendFields = nil
   289  				return j
   290  			}(),
   291  		},
   292  		{
   293  			dst: c.Dataset("dataset-id").Table("table-id"),
   294  			src: &QueryConfig{
   295  				Q:                "query string",
   296  				DefaultProjectID: "def-project-id",
   297  				DefaultDatasetID: "def-dataset-id",
   298  				TimePartitioning: &TimePartitioning{Type: DayPartitioningType},
   299  			},
   300  			want: func() *bq.Job {
   301  				j := defaultQueryJob()
   302  				j.Configuration.Query.ForceSendFields = nil
   303  				j.Configuration.Query.TimePartitioning = &bq.TimePartitioning{
   304  					Type: "DAY",
   305  				}
   306  				return j
   307  			}(),
   308  		},
   309  		{
   310  			dst: c.Dataset("dataset-id").Table("table-id"),
   311  			src: &QueryConfig{
   312  				Q:                "query string",
   313  				DefaultProjectID: "def-project-id",
   314  				DefaultDatasetID: "def-dataset-id",
   315  				RangePartitioning: &RangePartitioning{
   316  					Field: "foo",
   317  					Range: &RangePartitioningRange{
   318  						Start:    1,
   319  						End:      2,
   320  						Interval: 3,
   321  					},
   322  				},
   323  			},
   324  			want: func() *bq.Job {
   325  				j := defaultQueryJob()
   326  				j.Configuration.Query.ForceSendFields = nil
   327  				j.Configuration.Query.RangePartitioning = &bq.RangePartitioning{
   328  					Field: "foo",
   329  					Range: &bq.RangePartitioningRange{
   330  						Start:           1,
   331  						End:             2,
   332  						Interval:        3,
   333  						ForceSendFields: []string{"Start", "End", "Interval"},
   334  					},
   335  				}
   336  				return j
   337  			}(),
   338  		},
   339  		{
   340  			dst: c.Dataset("dataset-id").Table("table-id"),
   341  			src: &QueryConfig{
   342  				Q:                "query string",
   343  				DefaultProjectID: "def-project-id",
   344  				DefaultDatasetID: "def-dataset-id",
   345  				ConnectionProperties: []*ConnectionProperty{
   346  					{Key: "key-a", Value: "value-a"},
   347  					{Key: "key-b", Value: "value-b"},
   348  				},
   349  			},
   350  			want: func() *bq.Job {
   351  				j := defaultQueryJob()
   352  				j.Configuration.Query.ForceSendFields = nil
   353  				j.Configuration.Query.ConnectionProperties = []*bq.ConnectionProperty{
   354  					{Key: "key-a", Value: "value-a"},
   355  					{Key: "key-b", Value: "value-b"},
   356  				}
   357  				return j
   358  			}(),
   359  		},
   360  	}
   361  	for i, tc := range testCases {
   362  		query := c.Query("")
   363  		query.JobIDConfig = tc.jobIDConfig
   364  		query.QueryConfig = *tc.src
   365  		query.Dst = tc.dst
   366  		got, err := query.newJob()
   367  		if err != nil {
   368  			t.Errorf("#%d: err calling query: %v", i, err)
   369  			continue
   370  		}
   371  		checkJob(t, i, got, tc.want)
   372  
   373  		// Round-trip.
   374  		jc, err := bqToJobConfig(got.Configuration, c)
   375  		if err != nil {
   376  			t.Fatalf("#%d: %v", i, err)
   377  		}
   378  		wantConfig := query.QueryConfig
   379  		// We set AllowLargeResults to true when DisableFlattenedResults is true.
   380  		if wantConfig.DisableFlattenedResults {
   381  			wantConfig.AllowLargeResults = true
   382  		}
   383  		// A QueryConfig with neither UseXXXSQL field set is equivalent
   384  		// to one where UseStandardSQL = true.
   385  		if !wantConfig.UseLegacySQL && !wantConfig.UseStandardSQL {
   386  			wantConfig.UseStandardSQL = true
   387  		}
   388  		// Treat nil and empty tables the same, and ignore the client.
   389  		tableEqual := func(t1, t2 *Table) bool {
   390  			if t1 == nil {
   391  				t1 = &Table{}
   392  			}
   393  			if t2 == nil {
   394  				t2 = &Table{}
   395  			}
   396  			return t1.ProjectID == t2.ProjectID && t1.DatasetID == t2.DatasetID && t1.TableID == t2.TableID
   397  		}
   398  		// A table definition that is a GCSReference round-trips as an ExternalDataConfig.
   399  		// TODO(jba): see if there is a way to express this with a transformer.
   400  		gcsRefToEDC := func(g *GCSReference) *ExternalDataConfig {
   401  			q := g.toBQ()
   402  			e, _ := bqToExternalDataConfig(&q)
   403  			return e
   404  		}
   405  		externalDataEqual := func(e1, e2 ExternalData) bool {
   406  			if r, ok := e1.(*GCSReference); ok {
   407  				e1 = gcsRefToEDC(r)
   408  			}
   409  			if r, ok := e2.(*GCSReference); ok {
   410  				e2 = gcsRefToEDC(r)
   411  			}
   412  			return cmp.Equal(e1, e2)
   413  		}
   414  		diff := testutil.Diff(jc.(*QueryConfig), &wantConfig,
   415  			cmp.Comparer(tableEqual),
   416  			cmp.Comparer(externalDataEqual),
   417  			cmp.AllowUnexported(QueryConfig{}),
   418  		)
   419  		if diff != "" {
   420  			t.Errorf("#%d: (got=-, want=+:\n%s", i, diff)
   421  		}
   422  	}
   423  }
   424  
   425  func TestProbeFastPath(t *testing.T) {
   426  	c := &Client{
   427  		projectID: "client-project-id",
   428  	}
   429  	pfalse := false
   430  	testCases := []struct {
   431  		inCfg   QueryConfig
   432  		wantReq *bq.QueryRequest
   433  		wantErr bool
   434  	}{
   435  		{
   436  			inCfg: QueryConfig{
   437  				Q: "foo",
   438  			},
   439  			wantReq: &bq.QueryRequest{
   440  				Query:        "foo",
   441  				UseLegacySql: &pfalse,
   442  				FormatOptions: &bq.DataFormatOptions{
   443  					UseInt64Timestamp: true,
   444  				},
   445  			},
   446  		},
   447  		{
   448  			// All things you can set and still get a successful QueryRequest
   449  			inCfg: QueryConfig{
   450  				Q:                 "foo",
   451  				DefaultProjectID:  "defproject",
   452  				DefaultDatasetID:  "defdataset",
   453  				DisableQueryCache: true,
   454  				Priority:          InteractivePriority,
   455  				MaxBytesBilled:    123,
   456  				Parameters: []QueryParameter{
   457  					{Name: "user", Value: "bob"},
   458  				},
   459  				Labels: map[string]string{
   460  					"key": "val",
   461  				},
   462  			},
   463  			wantReq: &bq.QueryRequest{
   464  				Query:          "foo",
   465  				DefaultDataset: &bq.DatasetReference{ProjectId: "defproject", DatasetId: "defdataset"},
   466  				Labels: map[string]string{
   467  					"key": "val",
   468  				},
   469  				MaximumBytesBilled: 123,
   470  				UseLegacySql:       &pfalse,
   471  				QueryParameters: []*bq.QueryParameter{
   472  					{
   473  						Name:           "user",
   474  						ParameterType:  &bq.QueryParameterType{Type: "STRING"},
   475  						ParameterValue: &bq.QueryParameterValue{Value: "bob"},
   476  					},
   477  				},
   478  				UseQueryCache: &pfalse,
   479  				FormatOptions: &bq.DataFormatOptions{
   480  					UseInt64Timestamp: true,
   481  				},
   482  			},
   483  		},
   484  		{
   485  			// fail, sets destination via API
   486  			inCfg: QueryConfig{
   487  				Q:   "foo",
   488  				Dst: &Table{},
   489  			},
   490  			wantErr: true,
   491  		},
   492  		{
   493  			// fail, sets specifies destination partitioning
   494  			inCfg: QueryConfig{
   495  				Q:                 "foo",
   496  				TimePartitioning:  &TimePartitioning{},
   497  				RangePartitioning: &RangePartitioning{},
   498  			},
   499  			wantErr: true,
   500  		},
   501  		{
   502  			// fail, sets specifies schema update options
   503  			inCfg: QueryConfig{
   504  				Q:                   "foo",
   505  				SchemaUpdateOptions: []string{"bar"},
   506  			},
   507  			wantErr: true,
   508  		},
   509  	}
   510  	for i, tc := range testCases {
   511  		in := &Query{
   512  			QueryConfig: tc.inCfg,
   513  			client:      c,
   514  		}
   515  		gotReq, err := in.probeFastPath()
   516  		if tc.wantErr && err == nil {
   517  			t.Errorf("case %d wanted error, got nil", i)
   518  		}
   519  		if diff := testutil.Diff(gotReq, tc.wantReq, cmpopts.IgnoreFields(bq.QueryRequest{}, "RequestId")); diff != "" {
   520  			t.Errorf("QueryRequest case %d: -got +want:\n%s", i, diff)
   521  		}
   522  	}
   523  }
   524  
   525  func TestConfiguringQuery(t *testing.T) {
   526  	c := &Client{
   527  		projectID: "project-id",
   528  	}
   529  
   530  	query := c.Query("q")
   531  	query.JobID = "ajob"
   532  	query.DefaultProjectID = "def-project-id"
   533  	query.DefaultDatasetID = "def-dataset-id"
   534  	query.TimePartitioning = &TimePartitioning{Expiration: 1234 * time.Second, Field: "f"}
   535  	query.Clustering = &Clustering{
   536  		Fields: []string{"cfield1"},
   537  	}
   538  	query.DestinationEncryptionConfig = &EncryptionConfig{KMSKeyName: "keyName"}
   539  	query.SchemaUpdateOptions = []string{"ALLOW_FIELD_ADDITION"}
   540  	query.JobTimeout = time.Duration(5) * time.Second
   541  
   542  	// Note: Other configuration fields are tested in other tests above.
   543  	// A lot of that can be consolidated once Client.Copy is gone.
   544  
   545  	pfalse := false
   546  	want := &bq.Job{
   547  		Configuration: &bq.JobConfiguration{
   548  			Query: &bq.JobConfigurationQuery{
   549  				Query: "q",
   550  				DefaultDataset: &bq.DatasetReference{
   551  					ProjectId: "def-project-id",
   552  					DatasetId: "def-dataset-id",
   553  				},
   554  				UseLegacySql:                       &pfalse,
   555  				TimePartitioning:                   &bq.TimePartitioning{ExpirationMs: 1234000, Field: "f", Type: "DAY"},
   556  				Clustering:                         &bq.Clustering{Fields: []string{"cfield1"}},
   557  				DestinationEncryptionConfiguration: &bq.EncryptionConfiguration{KmsKeyName: "keyName"},
   558  				SchemaUpdateOptions:                []string{"ALLOW_FIELD_ADDITION"},
   559  			},
   560  			JobTimeoutMs: 5000,
   561  		},
   562  		JobReference: &bq.JobReference{
   563  			JobId:     "ajob",
   564  			ProjectId: "project-id",
   565  		},
   566  	}
   567  
   568  	got, err := query.newJob()
   569  	if err != nil {
   570  		t.Fatalf("err calling Query.newJob: %v", err)
   571  	}
   572  	if diff := testutil.Diff(got, want); diff != "" {
   573  		t.Errorf("querying: -got +want:\n%s", diff)
   574  	}
   575  }
   576  
   577  func TestQueryLegacySQL(t *testing.T) {
   578  	c := &Client{projectID: "project-id"}
   579  	q := c.Query("q")
   580  	q.UseStandardSQL = true
   581  	q.UseLegacySQL = true
   582  	_, err := q.newJob()
   583  	if err == nil {
   584  		t.Error("UseStandardSQL and UseLegacySQL: got nil, want error")
   585  	}
   586  	q = c.Query("q")
   587  	q.Parameters = []QueryParameter{{Name: "p", Value: 3}}
   588  	q.UseLegacySQL = true
   589  	_, err = q.newJob()
   590  	if err == nil {
   591  		t.Error("Parameters and UseLegacySQL: got nil, want error")
   592  	}
   593  }
   594  

View as plain text