1425

Committed 19 Nov 2024 01:38PM UTC coverage: 94.535% (-1.4%) from 95.96%

Build # 1425

Build Type

Pull #183

travis-pro

Committed by

web-flow

Commit Message

Merge branch 'main' into dependabot/go_modules/golang.org/x/net-0.23.0

Pull Request Pull Request #183: Bump golang.org/x/net from 0.0.0-20200421231249-e086a090c8fd to 0.23.0

Run Details

2387 of 2525 relevant lines covered (94.53%)

134.15 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

31.87

/cloud/bqx/dataset.go

//  Copyright 2017 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
//
//    https://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 bqx includes generally useful abstractions for simplifying
// interactions with bigquery.
// Production extensions should go here, but test facilities should go
// in a separate bqtest package.
package bqx

import (
        "errors"
        "fmt"
        "log"
        "reflect"
        "strings"
        "time"

        "cloud.google.com/go/bigquery"
        "golang.org/x/net/context"
        "google.golang.org/api/iterator"
        "google.golang.org/api/option"
)

// Dataset provides extensions to the bigquery Dataset and Dataset
// objects to streamline common actions.
// It encapsulates the Client and Dataset to simplify methods.
type Dataset struct {
        *bigquery.Dataset // Exposes Dataset API directly.
        BqClient          *bigquery.Client
}

// NewDataset creates a Dataset for a project.
// httpClient is used to inject mocks for the bigquery client.
// if httpClient is nil, a suitable default client is used.
// Additional bigquery ClientOptions may be optionally passed as final
//   clientOpts argument.  This is useful for testing credentials.
func NewDataset(project, dataset string, clientOpts ...option.ClientOption) (Dataset, error) {
        ctx := context.Background()
        var bqClient *bigquery.Client
        var err error
        bqClient, err = bigquery.NewClient(ctx, project, clientOpts...)

        if err != nil {
                return Dataset{}, err
        }

        return Dataset{bqClient.Dataset(dataset), bqClient}, nil
}

// ResultQuery constructs a query with common QueryConfig settings for
// writing results to a table.
// Generally, may need to change WriteDisposition.
func (dsExt *Dataset) ResultQuery(query string, dryRun bool) *bigquery.Query {
        q := dsExt.BqClient.Query(query)
        q.QueryConfig.DryRun = dryRun
        if strings.HasPrefix(query, "#legacySQL") {
                q.QueryConfig.UseLegacySQL = true
        }
        // Default for unqualified table names in the query.
        q.QueryConfig.DefaultProjectID = dsExt.ProjectID
        q.QueryConfig.DefaultDatasetID = dsExt.DatasetID
        return q
}

///////////////////////////////////////////////////////////////////
// Code to execute a single query and parse single row result.
///////////////////////////////////////////////////////////////////

// QueryAndParse executes a query that should return a single row, with
// all struct fields that match query columns filled in.
// The caller must pass in the *address* of an appropriate struct.
// TODO - extend this to also handle multirow results, by passing
// slice of structs.
func (dsExt *Dataset) QueryAndParse(q string, structPtr interface{}) error {
        typeInfo := reflect.ValueOf(structPtr)

        if typeInfo.Type().Kind() != reflect.Ptr {
                return errors.New("Argument should be ptr to struct")
        }
        if reflect.Indirect(typeInfo).Kind() != reflect.Struct {
                return errors.New("Argument should be ptr to struct")
        }

        query := dsExt.ResultQuery(q, false)
        it, err := query.Read(context.Background())
        if err != nil {
                return err
        }

        // We expect a single result row, so proceed accordingly.
        err = it.Next(structPtr)
        if err != nil {
                return err
        }
        var row map[string]bigquery.Value
        // If there are more rows, then something is wrong.
        err = it.Next(&row)
        if err != iterator.Done {
                return errors.New("multiple row data")
        }
        return nil
}

// PartitionInfo provides basic information about a partition.
type PartitionInfo struct {
        PartitionID  string
        CreationTime time.Time
        LastModified time.Time
}

// GetPartitionInfo provides basic information about a partition.
func (dsExt Dataset) GetPartitionInfo(table string, partition string) (PartitionInfo, error) {
        // This uses legacy, because PARTITION_SUMMARY is not supported in standard.
        queryString := fmt.Sprintf(
                `#legacySQL
                SELECT
                  partition_id as PartitionID,
                  msec_to_timestamp(creation_time) AS CreationTime,
                  msec_to_timestamp(last_modified_time) AS LastModified
                FROM
                  [%s$__PARTITIONS_SUMMARY__]
                where partition_id = "%s" `, table, partition)
        pi := PartitionInfo{}

        err := dsExt.QueryAndParse(queryString, &pi)
        if err != nil {
                log.Println(err, ":", queryString)
                return PartitionInfo{}, err
        }
        return pi, nil
}

// DestQuery constructs a query with common Config settings for
// writing results to a table.
// If dest is nil, then this will create a DryRun query.
// TODO - should disposition be an opts... field instead?
func (dsExt *Dataset) DestQuery(query string, dest *bigquery.Table, disposition bigquery.TableWriteDisposition) *bigquery.Query {
        q := dsExt.BqClient.Query(query)
        if dest != nil {
                q.QueryConfig.Dst = dest
        } else {
                q.QueryConfig.DryRun = true
        }
        q.QueryConfig.WriteDisposition = disposition
        q.QueryConfig.AllowLargeResults = true
        // Default for unqualified table names in the query.
        q.QueryConfig.DefaultProjectID = dsExt.ProjectID
        q.QueryConfig.DefaultDatasetID = dsExt.DatasetID
        q.QueryConfig.DisableFlattenedResults = true
        return q
}

// ExecDestQuery executes a destination or dryrun query, and returns status or error.
func (dsExt *Dataset) ExecDestQuery(q *bigquery.Query) (*bigquery.JobStatus, error) {
        if q.QueryConfig.Dst == nil && q.QueryConfig.DryRun == false {
                return nil, errors.New("query must be a destination or dry run")
        }
        job, err := q.Run(context.Background())
        if err != nil {
                return nil, err
        }
        log.Println("JobID:", job.ID())
        status, err := job.Wait(context.Background())
        if err != nil {
                return status, err
        }
        return status, nil
}

1	// Copyright 2017 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	// https://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 bqx includes generally useful abstractions for simplifying
16	// interactions with bigquery.
17	// Production extensions should go here, but test facilities should go
18	// in a separate bqtest package.
19	package bqx
20
21	import (
22	"errors"
23	"fmt"
24	"log"
25	"reflect"
26	"strings"
27	"time"
28
29	"cloud.google.com/go/bigquery"
30	"golang.org/x/net/context"
31	"google.golang.org/api/iterator"
32	"google.golang.org/api/option"
33	)
34
35	// Dataset provides extensions to the bigquery Dataset and Dataset
36	// objects to streamline common actions.
37	// It encapsulates the Client and Dataset to simplify methods.
38	type Dataset struct {
39	*bigquery.Dataset // Exposes Dataset API directly.
40	BqClient *bigquery.Client
41	}
42
43	// NewDataset creates a Dataset for a project.
44	// httpClient is used to inject mocks for the bigquery client.
45	// if httpClient is nil, a suitable default client is used.
46	// Additional bigquery ClientOptions may be optionally passed as final
47	// clientOpts argument. This is useful for testing credentials.
48	func NewDataset(project, dataset string, clientOpts ...option.ClientOption) (Dataset, error) {	6✔
49	ctx := context.Background()	6✔
50	var bqClient *bigquery.Client	6✔
51	var err error	6✔
52	bqClient, err = bigquery.NewClient(ctx, project, clientOpts...)	6✔
53		6✔
54	if err != nil {	6✔
55	return Dataset{}, err	×
56	}	×
57
58	return Dataset{bqClient.Dataset(dataset), bqClient}, nil	6✔
59	}
60
61	// ResultQuery constructs a query with common QueryConfig settings for
62	// writing results to a table.
63	// Generally, may need to change WriteDisposition.
64	func (dsExt Dataset) ResultQuery(query string, dryRun bool) bigquery.Query {	4✔
65	q := dsExt.BqClient.Query(query)	4✔
66	q.QueryConfig.DryRun = dryRun	4✔
67	if strings.HasPrefix(query, "#legacySQL") {	4✔
68	q.QueryConfig.UseLegacySQL = true	×
69	}	×
70	// Default for unqualified table names in the query.
71	q.QueryConfig.DefaultProjectID = dsExt.ProjectID	4✔
72	q.QueryConfig.DefaultDatasetID = dsExt.DatasetID	4✔
73	return q	4✔
74	}
75
76	///////////////////////////////////////////////////////////////////
77	// Code to execute a single query and parse single row result.
78	///////////////////////////////////////////////////////////////////
79
80	// QueryAndParse executes a query that should return a single row, with
81	// all struct fields that match query columns filled in.
82	// The caller must pass in the address of an appropriate struct.
83	// TODO - extend this to also handle multirow results, by passing
84	// slice of structs.
85	func (dsExt *Dataset) QueryAndParse(q string, structPtr interface{}) error {	×
86	typeInfo := reflect.ValueOf(structPtr)	×
87		×
88	if typeInfo.Type().Kind() != reflect.Ptr {	×
89	return errors.New("Argument should be ptr to struct")	×
90	}	×
91	if reflect.Indirect(typeInfo).Kind() != reflect.Struct {	×
92	return errors.New("Argument should be ptr to struct")	×
93	}	×
94
95	query := dsExt.ResultQuery(q, false)	×
96	it, err := query.Read(context.Background())	×
97	if err != nil {	×
98	return err	×
99	}	×
100
101	// We expect a single result row, so proceed accordingly.
102	err = it.Next(structPtr)	×
103	if err != nil {	×
104	return err	×
105	}	×
106	var row map[string]bigquery.Value	×
107	// If there are more rows, then something is wrong.	×
108	err = it.Next(&row)	×
109	if err != iterator.Done {	×
110	return errors.New("multiple row data")	×
111	}	×
112	return nil	×
113	}
114
115	// PartitionInfo provides basic information about a partition.
116	type PartitionInfo struct {
117	PartitionID string
118	CreationTime time.Time
119	LastModified time.Time
120	}
121
122	// GetPartitionInfo provides basic information about a partition.
123	func (dsExt Dataset) GetPartitionInfo(table string, partition string) (PartitionInfo, error) {	×
124	// This uses legacy, because PARTITION_SUMMARY is not supported in standard.	×
125	queryString := fmt.Sprintf(	×
126	`#legacySQL	×
127	SELECT	×
128	partition_id as PartitionID,	×
129	msec_to_timestamp(creation_time) AS CreationTime,	×
130	msec_to_timestamp(last_modified_time) AS LastModified	×
131	FROM	×
132	[%s$__PARTITIONS_SUMMARY__]	×
133	where partition_id = "%s" `, table, partition)	×
134	pi := PartitionInfo{}	×
135		×
136	err := dsExt.QueryAndParse(queryString, &pi)	×
137	if err != nil {	×
138	log.Println(err, ":", queryString)	×
139	return PartitionInfo{}, err	×
140	}	×
141	return pi, nil	×
142	}
143
144	// DestQuery constructs a query with common Config settings for
145	// writing results to a table.
146	// If dest is nil, then this will create a DryRun query.
147	// TODO - should disposition be an opts... field instead?
148	func (dsExt Dataset) DestQuery(query string, dest bigquery.Table, disposition bigquery.TableWriteDisposition) *bigquery.Query {	4✔
149	q := dsExt.BqClient.Query(query)	4✔
150	if dest != nil {	6✔
151	q.QueryConfig.Dst = dest	2✔
152	} else {	4✔
153	q.QueryConfig.DryRun = true	2✔
154	}	2✔
155	q.QueryConfig.WriteDisposition = disposition	4✔
156	q.QueryConfig.AllowLargeResults = true	4✔
157	// Default for unqualified table names in the query.	4✔
158	q.QueryConfig.DefaultProjectID = dsExt.ProjectID	4✔
159	q.QueryConfig.DefaultDatasetID = dsExt.DatasetID	4✔
160	q.QueryConfig.DisableFlattenedResults = true	4✔
161	return q	4✔
162	}
163
164	// ExecDestQuery executes a destination or dryrun query, and returns status or error.
165	func (dsExt Dataset) ExecDestQuery(q bigquery.Query) (*bigquery.JobStatus, error) {	×
166	if q.QueryConfig.Dst == nil && q.QueryConfig.DryRun == false {	×
167	return nil, errors.New("query must be a destination or dry run")	×
168	}	×
169	job, err := q.Run(context.Background())	×
170	if err != nil {	×
171	return nil, err	×
172	}	×
173	log.Println("JobID:", job.ID())	×
174	status, err := job.Wait(context.Background())	×
175	if err != nil {	×
176	return status, err	×
177	}	×
178	return status, nil	×
179	}

m-lab / go / 1425

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous