// Copyright 2015 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 bigquery

import (

"context"

"errors"

"fmt"

"strings"

"time"

"cloud.google.com/go/internal/optional"

"cloud.google.com/go/internal/trace"

bq "google.golang.org/api/bigquery/v2"

"google.golang.org/api/iterator"

)

// Dataset is a reference to a BigQuery dataset.

type Dataset struct {

ProjectID string

DatasetID string

c *Client

}

// DatasetMetadata contains information about a BigQuery dataset.

type DatasetMetadata struct {

// These fields can be set when creating a dataset.

Name string // The user-friendly name for this dataset.

Description string // The user-friendly description of this dataset.

Location string // The geo location of the dataset.

DefaultTableExpiration time.Duration // The default expiration time for new tables.

Labels map[string]string // User-provided labels.

Access []*AccessEntry // Access permissions.

DefaultEncryptionConfig *EncryptionConfig

// DefaultPartitionExpiration is the default expiration time for

// all newly created partitioned tables in the dataset.

DefaultPartitionExpiration time.Duration

// Defines the default collation specification of future tables

// created in the dataset. If a table is created in this dataset without

// table-level default collation, then the table inherits the dataset default

// collation, which is applied to the string fields that do not have explicit

// collation specified. A change to this field affects only tables created

// afterwards, and does not alter the existing tables.

// More information: https://cloud.google.com/bigquery/docs/reference/standard-sql/collation-concepts

DefaultCollation string

// For externally defined datasets, contains information about the configuration.

ExternalDatasetReference *ExternalDatasetReference

// MaxTimeTravel represents the number of hours for the max time travel for all tables

// in the dataset. Durations are rounded towards zero for the nearest hourly value.

MaxTimeTravel time.Duration

// Storage billing model to be used for all tables in the dataset.

// Can be set to PHYSICAL. Default is LOGICAL.

// Once you create a dataset with storage billing model set to physical bytes, you can't change it back to using logical bytes again.

// More details: https://cloud.google.com/bigquery/docs/datasets-intro#dataset_storage_billing_models

StorageBillingModel string

// These fields are read-only.

CreationTime time.Time

LastModifiedTime time.Time // When the dataset or any of its tables were modified.

FullID string // The full dataset ID in the form projectID:datasetID.

// The tags associated with this dataset. Tag keys are

// globally unique, and managed via the resource manager API.

// More information: https://cloud.google.com/resource-manager/docs/tags/tags-overview

Tags []*DatasetTag

// TRUE if the dataset and its table names are case-insensitive, otherwise

// FALSE. By default, this is FALSE, which means the dataset and its table

// names are case-sensitive. This field does not affect routine references.

IsCaseInsensitive bool

// ETag is the ETag obtained when reading metadata. Pass it to Dataset.Update to

// ensure that the metadata hasn't changed since it was read.

ETag string

}

// DatasetTag is a representation of a single tag key/value.

type DatasetTag struct {

// TagKey is the namespaced friendly name of the tag key, e.g.

// "12345/environment" where 12345 is org id.

TagKey string

// TagValue is the friendly short name of the tag value, e.g.

// "production".

TagValue string

}

const (

// LogicalStorageBillingModel indicates billing for logical bytes.

LogicalStorageBillingModel = ""

// PhysicalStorageBillingModel indicates billing for physical bytes.

PhysicalStorageBillingModel = "PHYSICAL"

)

func bqToDatasetTag(in *bq.DatasetTags) *DatasetTag {

if in == nil {

return nil

}

return &DatasetTag{

TagKey: in.TagKey,

TagValue: in.TagValue,

}

}

// DatasetMetadataToUpdate is used when updating a dataset's metadata.

// Only non-nil fields will be updated.

type DatasetMetadataToUpdate struct {

Description optional.String // The user-friendly description of this table.

Name optional.String // The user-friendly name for this dataset.

// DefaultTableExpiration is the default expiration time for new tables.

// If set to time.Duration(0), new tables never expire.

DefaultTableExpiration optional.Duration

// DefaultTableExpiration is the default expiration time for

// all newly created partitioned tables.

// If set to time.Duration(0), new table partitions never expire.

DefaultPartitionExpiration optional.Duration

// DefaultEncryptionConfig defines CMEK settings for new resources created

// in the dataset.

DefaultEncryptionConfig *EncryptionConfig

// Defines the default collation specification of future tables

// created in the dataset.

DefaultCollation optional.String

// For externally defined datasets, contains information about the configuration.

ExternalDatasetReference *ExternalDatasetReference

// MaxTimeTravel represents the number of hours for the max time travel for all tables

// in the dataset. Durations are rounded towards zero for the nearest hourly value.

MaxTimeTravel optional.Duration

// Storage billing model to be used for all tables in the dataset.

// Can be set to PHYSICAL. Default is LOGICAL.

// Once you change a dataset's storage billing model to use physical bytes, you can't change it back to using logical bytes again.

// More details: https://cloud.google.com/bigquery/docs/datasets-intro#dataset_storage_billing_models

StorageBillingModel optional.String

// The entire access list. It is not possible to replace individual entries.

Access []*AccessEntry

// TRUE if the dataset and its table names are case-insensitive, otherwise

// FALSE. By default, this is FALSE, which means the dataset and its table

// names are case-sensitive. This field does not affect routine references.

IsCaseInsensitive optional.Bool

labelUpdater

}

// Dataset creates a handle to a BigQuery dataset in the client's project.

func (c *Client) Dataset(id string) *Dataset {

return c.DatasetInProject(c.projectID, id)

}

// DatasetInProject creates a handle to a BigQuery dataset in the specified project.

func (c *Client) DatasetInProject(projectID, datasetID string) *Dataset {

return &Dataset{

ProjectID: projectID,

DatasetID: datasetID,

c: c,

}

}

// Identifier returns the ID of the dataset in the requested format.

//

// For Standard SQL format, the identifier will be quoted if the

// ProjectID contains dash (-) characters.

func (d *Dataset) Identifier(f IdentifierFormat) (string, error) {

switch f {

case LegacySQLID:

return fmt.Sprintf("%s:%s", d.ProjectID, d.DatasetID), nil

case StandardSQLID:

// Quote project identifiers if they have a dash character.

if strings.Contains(d.ProjectID, "-") {

return fmt.Sprintf("`%s`.%s", d.ProjectID, d.DatasetID), nil

}

return fmt.Sprintf("%s.%s", d.ProjectID, d.DatasetID), nil

default:

return "", ErrUnknownIdentifierFormat

}

}

// Create creates a dataset in the BigQuery service.

//

// An error will be returned if the dataset already exists.

// Pass in a DatasetMetadata value to configure the dataset.

func (d *Dataset) Create(ctx context.Context, md *DatasetMetadata) (err error) {

return d.CreateWithOptions(ctx, md)

}

// CreateWithOptions creates a dataset in the BigQuery service, and

// provides additional options to control the behavior of the call.

//

// An error will be returned if the dataset already exists.

// Pass in a DatasetMetadata value to configure the dataset.

func (d *Dataset) CreateWithOptions(ctx context.Context, md *DatasetMetadata, opts ...DatasetOption) (err error) {

ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Create")

defer func() { trace.EndSpan(ctx, err) }()

cOpt := &dsCallOption{}

for _, o := range opts {

o(cOpt)

}

ds, err := md.toBQ()

if err != nil {

return err

}

ds.DatasetReference = &bq.DatasetReference{DatasetId: d.DatasetID}

// Use Client.Location as a default.

if ds.Location == "" {

ds.Location = d.c.Location

}

ctx = setProjectItemTraceMetadata(ctx, d.ProjectID, "datasets")

call := d.c.bqs.Datasets.Insert(d.ProjectID, ds).Context(ctx)

setClientHeader(call.Header())

if cOpt.accessPolicyVersion != nil {

call.AccessPolicyVersion(int64(optional.ToInt(cOpt.accessPolicyVersion)))

}

_, err = call.Do()

return err

}

func (dm *DatasetMetadata) toBQ() (*bq.Dataset, error) {

ds := &bq.Dataset{}

if dm == nil {

return ds, nil

}

ds.FriendlyName = dm.Name

ds.Description = dm.Description

ds.Location = dm.Location

ds.DefaultTableExpirationMs = int64(dm.DefaultTableExpiration / time.Millisecond)

ds.DefaultPartitionExpirationMs = int64(dm.DefaultPartitionExpiration / time.Millisecond)

ds.DefaultCollation = dm.DefaultCollation

ds.MaxTimeTravelHours = int64(dm.MaxTimeTravel / time.Hour)

ds.StorageBillingModel = string(dm.StorageBillingModel)

ds.IsCaseInsensitive = dm.IsCaseInsensitive

ds.Labels = dm.Labels

// These fields are read-only, but we'll populate them should there be a change in the service behavior.

if !dm.CreationTime.IsZero() {

ds.CreationTime = dm.CreationTime.UnixMilli()

}

if !dm.LastModifiedTime.IsZero() {

ds.LastModifiedTime = dm.LastModifiedTime.UnixMilli()

}

ds.Etag = dm.ETag

ds.Id = dm.FullID

// access list

var err error

ds.Access, err = accessListToBQ(dm.Access)

if err != nil {

return nil, err

}

if dm.DefaultEncryptionConfig != nil {

ds.DefaultEncryptionConfiguration = dm.DefaultEncryptionConfig.toBQ()

}

if dm.ExternalDatasetReference != nil {

ds.ExternalDatasetReference = dm.ExternalDatasetReference.toBQ()

}

return ds, nil

}

func accessListToBQ(a []*AccessEntry) ([]*bq.DatasetAccess, error) {

var q []*bq.DatasetAccess

for _, e := range a {

a, err := e.toBQ()

if err != nil {

return nil, err

}

q = append(q, a)

}

return q, nil

}

// Delete deletes the dataset. Delete will fail if the dataset is not empty.

func (d *Dataset) Delete(ctx context.Context) (err error) {

return d.deleteInternal(ctx, false)

}

// DeleteWithContents deletes the dataset, as well as contained resources.

func (d *Dataset) DeleteWithContents(ctx context.Context) (err error) {

return d.deleteInternal(ctx, true)

}

func (d *Dataset) deleteInternal(ctx context.Context, deleteContents bool) (err error) {

ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Delete")

defer func() { trace.EndSpan(ctx, err) }()

ctx = setDatasetTraceMetadata(ctx, d.ProjectID, d.DatasetID)

call := d.c.bqs.Datasets.Delete(d.ProjectID, d.DatasetID).Context(ctx).DeleteContents(deleteContents)

setClientHeader(call.Header())

return runWithRetry(ctx, func() (err error) {

sCtx := trace.StartSpan(ctx, "bigquery.datasets.delete")

err = call.Do()

trace.EndSpan(sCtx, err)

return err

})

}

// Metadata fetches the metadata for the dataset.

func (d *Dataset) Metadata(ctx context.Context) (md *DatasetMetadata, err error) {

return d.MetadataWithOptions(ctx)

}

// MetadataWithOptions fetches metadata for the dataset, and provides additional options for

// controlling the request.

func (d *Dataset) MetadataWithOptions(ctx context.Context, opts ...DatasetOption) (md *DatasetMetadata, err error) {

ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Metadata")

defer func() { trace.EndSpan(ctx, err) }()

cOpt := &dsCallOption{}

for _, o := range opts {

o(cOpt)

}

ctx = setDatasetTraceMetadata(ctx, d.ProjectID, d.DatasetID)

call := d.c.bqs.Datasets.Get(d.ProjectID, d.DatasetID).Context(ctx)

setClientHeader(call.Header())

if cOpt.accessPolicyVersion != nil {

call.AccessPolicyVersion(int64(optional.ToInt(cOpt.accessPolicyVersion)))

}

if cOpt.datasetView != nil {

call.DatasetView(optional.ToString(cOpt.datasetView))

}

var ds *bq.Dataset

if err := runWithRetry(ctx, func() (err error) {

sCtx := trace.StartSpan(ctx, "bigquery.datasets.get")

ds, err = call.Do()

trace.EndSpan(sCtx, err)

return err

}); err != nil {

return nil, err

}

return bqToDatasetMetadata(ds, d.c)

}

// dsCallOption provides a general option holder for dataset RPCs

type dsCallOption struct {

accessPolicyVersion optional.Int

datasetView optional.String

updateMode optional.String

}

// DatasetOption provides an option type for customizing requests against the Dataset

// service.

type DatasetOption func(*dsCallOption)

// WithAccessPolicyVersion is an option that enabled setting of the Access Policy Version for a request

// where appropriate. Valid values are 0, 1, and 3.

//

// Requests specifying an invalid value will be rejected.

// Requests for conditional access policy binding in datasets must specify version 3.

//

// Dataset with no conditional role bindings in access policy may specify any valid value

// or leave the field unset.

//

// This field will be mapped to [IAM Policy version] (https://cloud.google.com/iam/docs/policies#versions)

// and will be used to fetch policy from IAM. If unset or if 0 or 1 value is used for

// dataset with conditional bindings, access entry with condition will have role string

// appended by 'withcond' string followed by a hash value.

//

// Please refer https://cloud.google.com/iam/docs/troubleshooting-withcond for more details.

func WithAccessPolicyVersion(apv int) DatasetOption {

return func(o *dsCallOption) {

o.accessPolicyVersion = apv

}

}

// DatasetView specifies which details about a dataset are desired.

type DatasetView string

const (

// DatasetMetadataView populates metadata information for the dataset,

// such as friendlyName, description, labels, etc.

DatasetMetadataView DatasetView = "METADATA"

// DatasetACLView populates information for the dataset, which defines

// dataset access for one or more entities.

DatasetACLView DatasetView = "ACL"

// DatasetFullView populates both dataset metadata and ACL information.

DatasetFullView DatasetView = "FULL"

// UnspecifiedDatasetView is the default value, which will be treated as DatasetFullView

UnspecifiedDatasetView DatasetView = "DATASET_VIEW_UNSPECIFIED"

)

// WithDatasetView specifies the view that determines which dataset information

// is returned. By default, metadata and ACL information are returned.

func WithDatasetView(view DatasetView) DatasetOption {

return func(o *dsCallOption) {

o.datasetView = string(view)

}

}

// DatasetUpdateMode specifies which fields of a dataset are going to be affected

// by update/patch operations.

type DatasetUpdateMode string

const (

// DatasetMetadataUpdateMode targets metadata information for the dataset,

// such as friendlyName, description, labels, etc.

DatasetMetadataUpdateMode DatasetUpdateMode = "UPDATE_METADATA"

// DatasetACLUpdateMode targets ACL information for the dataset,

// which defines dataset access for one or more entities.

DatasetACLUpdateMode DatasetUpdateMode = "UPDATE_ACL"

// DatasetFullUpdateMode targets both dataset metadata and ACL

// information on update operations.

DatasetFullUpdateMode DatasetUpdateMode = "UPDATE_FULL"

// UnspecifiedDatasetUpdateMode is the default value, which will be treated as DatasetFullUpdateMode

UnspecifiedDatasetUpdateMode DatasetUpdateMode = "UPDATE_MODE_UNSPECIFIED"

)

// WithUpdateMode specifies the fields of dataset that the update/patch

// operation is targeting. By default, both metadata and ACL fields are updated.

func WithUpdateMode(mode DatasetUpdateMode) DatasetOption {

return func(o *dsCallOption) {

o.updateMode = string(mode)

}

}

func bqToDatasetMetadata(d *bq.Dataset, c *Client) (*DatasetMetadata, error) {

dm := &DatasetMetadata{

CreationTime: unixMillisToTime(d.CreationTime),

LastModifiedTime: unixMillisToTime(d.LastModifiedTime),

DefaultTableExpiration: time.Duration(d.DefaultTableExpirationMs) * time.Millisecond,

DefaultPartitionExpiration: time.Duration(d.DefaultPartitionExpirationMs) * time.Millisecond,

DefaultCollation: d.DefaultCollation,

ExternalDatasetReference: bqToExternalDatasetReference(d.ExternalDatasetReference),

MaxTimeTravel: time.Duration(d.MaxTimeTravelHours) * time.Hour,

StorageBillingModel: d.StorageBillingModel,

DefaultEncryptionConfig: bqToEncryptionConfig(d.DefaultEncryptionConfiguration),

Description: d.Description,

Name: d.FriendlyName,

FullID: d.Id,

Location: d.Location,

Labels: d.Labels,

IsCaseInsensitive: d.IsCaseInsensitive,

ETag: d.Etag,

}

for _, a := range d.Access {

e, err := bqToAccessEntry(a, c)

if err != nil {

return nil, err

}

dm.Access = append(dm.Access, e)

}

for _, bqTag := range d.Tags {

tag := bqToDatasetTag(bqTag)

if tag != nil {

dm.Tags = append(dm.Tags, tag)

}

}

return dm, nil

}

// Update modifies specific Dataset metadata fields.

// To perform a read-modify-write that protects against intervening reads,

// set the etag argument to the DatasetMetadata.ETag field from the read.

// Pass the empty string for etag for a "blind write" that will always succeed.

func (d *Dataset) Update(ctx context.Context, dm DatasetMetadataToUpdate, etag string) (md *DatasetMetadata, err error) {

return d.UpdateWithOptions(ctx, dm, etag)

}

// UpdateWithOptions modifies specific Dataset metadata fields and

// provides an interface for specifying additional options to the request.

//

// To perform a read-modify-write that protects against intervening reads,

// set the etag argument to the DatasetMetadata.ETag field from the read.

// Pass the empty string for etag for a "blind write" that will always succeed.

func (d *Dataset) UpdateWithOptions(ctx context.Context, dm DatasetMetadataToUpdate, etag string, opts ...DatasetOption) (md *DatasetMetadata, err error) {

ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Update")

defer func() { trace.EndSpan(ctx, err) }()

cOpt := &dsCallOption{}

for _, o := range opts {

o(cOpt)

}

ds, err := dm.toBQ()

if err != nil {

return nil, err

}

ctx = setDatasetTraceMetadata(ctx, d.ProjectID, d.DatasetID)

call := d.c.bqs.Datasets.Patch(d.ProjectID, d.DatasetID, ds).Context(ctx)

setClientHeader(call.Header())

if etag != "" {

call.Header().Set("If-Match", etag)

}

if cOpt.accessPolicyVersion != nil {

call.AccessPolicyVersion(int64(optional.ToInt(cOpt.accessPolicyVersion)))

}

if cOpt.updateMode != nil {

call.UpdateMode(optional.ToString(cOpt.updateMode))

}

var ds2 *bq.Dataset

if err := runWithRetry(ctx, func() (err error) {

sCtx := trace.StartSpan(ctx, "bigquery.datasets.patch")

ds2, err = call.Do()

trace.EndSpan(sCtx, err)

return err

}); err != nil {

return nil, err

}

return bqToDatasetMetadata(ds2, d.c)

}

func (dm *DatasetMetadataToUpdate) toBQ() (*bq.Dataset, error) {

ds := &bq.Dataset{}

forceSend := func(field string) {

ds.ForceSendFields = append(ds.ForceSendFields, field)

}

if dm.Description != nil {

ds.Description = optional.ToString(dm.Description)

forceSend("Description")

}

if dm.Name != nil {

ds.FriendlyName = optional.ToString(dm.Name)

forceSend("FriendlyName")

}

if dm.DefaultTableExpiration != nil {

dur := optional.ToDuration(dm.DefaultTableExpiration)

if dur == 0 {

// Send a null to delete the field.

ds.NullFields = append(ds.NullFields, "DefaultTableExpirationMs")

} else {

ds.DefaultTableExpirationMs = int64(dur / time.Millisecond)

}

}

if dm.DefaultPartitionExpiration != nil {

dur := optional.ToDuration(dm.DefaultPartitionExpiration)

if dur == 0 {

// Send a null to delete the field.

ds.NullFields = append(ds.NullFields, "DefaultPartitionExpirationMs")

} else {

ds.DefaultPartitionExpirationMs = int64(dur / time.Millisecond)

}

}

if dm.DefaultCollation != nil {

ds.DefaultCollation = optional.ToString(dm.DefaultCollation)

forceSend("DefaultCollation")

}

if dm.ExternalDatasetReference != nil {

ds.ExternalDatasetReference = dm.ExternalDatasetReference.toBQ()

forceSend("ExternalDatasetReference")

}

if dm.MaxTimeTravel != nil {

dur := optional.ToDuration(dm.MaxTimeTravel)

if dur == 0 {

// Send a null to delete the field.

ds.NullFields = append(ds.NullFields, "MaxTimeTravelHours")

} else {

ds.MaxTimeTravelHours = int64(dur / time.Hour)

}

}

if dm.StorageBillingModel != nil {

ds.StorageBillingModel = optional.ToString(dm.StorageBillingModel)

forceSend("StorageBillingModel")

}

if dm.DefaultEncryptionConfig != nil {

ds.DefaultEncryptionConfiguration = dm.DefaultEncryptionConfig.toBQ()

ds.DefaultEncryptionConfiguration.ForceSendFields = []string{"KmsKeyName"}

}

if dm.Access != nil {

var err error

ds.Access, err = accessListToBQ(dm.Access)

if err != nil {

return nil, err

}

if len(ds.Access) == 0 {

ds.NullFields = append(ds.NullFields, "Access")

}

}

if dm.IsCaseInsensitive != nil {

ds.IsCaseInsensitive = optional.ToBool(dm.IsCaseInsensitive)

forceSend("IsCaseInsensitive")

}

labels, forces, nulls := dm.update()

ds.Labels = labels

ds.ForceSendFields = append(ds.ForceSendFields, forces...)

ds.NullFields = append(ds.NullFields, nulls...)

return ds, nil

}

// Table creates a handle to a BigQuery table in the dataset.

// To determine if a table exists, call Table.Metadata.

// If the table does not already exist, use Table.Create to create it.

func (d *Dataset) Table(tableID string) *Table {

return &Table{ProjectID: d.ProjectID, DatasetID: d.DatasetID, TableID: tableID, c: d.c}

}

// Tables returns an iterator over the tables in the Dataset.

func (d *Dataset) Tables(ctx context.Context) *TableIterator {

it := &TableIterator{

ctx: ctx,

dataset: d,

}

it.pageInfo, it.nextFunc = iterator.NewPageInfo(

it.fetch,

func() int { return len(it.tables) },

func() interface{} { b := it.tables; it.tables = nil; return b })

return it

}

// A TableIterator is an iterator over Tables.

type TableIterator struct {

ctx context.Context

dataset *Dataset

tables []*Table

pageInfo *iterator.PageInfo

nextFunc func() error

}

// Next returns the next result. Its second return value is Done if there are

// no more results. Once Next returns Done, all subsequent calls will return

// Done.

func (it *TableIterator) Next() (*Table, error) {

if err := it.nextFunc(); err != nil {

return nil, err

}

t := it.tables[0]

it.tables = it.tables[1:]

return t, nil

}

// PageInfo supports pagination. See the google.golang.org/api/iterator package for details.

func (it *TableIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }

// listTables exists to aid testing.

var listTables = func(it *TableIterator, pageSize int, pageToken string) (*bq.TableList, error) {

it.ctx = setDatasetItemTraceMetadata(it.ctx, it.dataset.ProjectID, it.dataset.DatasetID, "tables")

call := it.dataset.c.bqs.Tables.List(it.dataset.ProjectID, it.dataset.DatasetID).

PageToken(pageToken).

Context(it.ctx)

setClientHeader(call.Header())

if pageSize > 0 {

call.MaxResults(int64(pageSize))

}

var res *bq.TableList

err := runWithRetry(it.ctx, func() (err error) {

sCtx := trace.StartSpan(it.ctx, "bigquery.tables.list")

res, err = call.Do()

trace.EndSpan(sCtx, err)

return err

})

return res, err

}

func (it *TableIterator) fetch(pageSize int, pageToken string) (string, error) {

res, err := listTables(it, pageSize, pageToken)

if err != nil {

return "", err

}

for _, t := range res.Tables {

it.tables = append(it.tables, bqToTable(t.TableReference, it.dataset.c))

}

return res.NextPageToken, nil

}

func bqToTable(tr *bq.TableReference, c *Client) *Table {

if tr == nil {

return nil

}

return &Table{

ProjectID: tr.ProjectId,

DatasetID: tr.DatasetId,

TableID: tr.TableId,

c: c,

}

}

// Model creates a handle to a BigQuery model in the dataset.

// To determine if a model exists, call Model.Metadata.

// If the model does not already exist, you can create it via execution

// of a CREATE MODEL query.

func (d *Dataset) Model(modelID string) *Model {

return &Model{ProjectID: d.ProjectID, DatasetID: d.DatasetID, ModelID: modelID, c: d.c}

}

// Models returns an iterator over the models in the Dataset.

func (d *Dataset) Models(ctx context.Context) *ModelIterator {

it := &ModelIterator{

ctx: ctx,

dataset: d,

}

it.pageInfo, it.nextFunc = iterator.NewPageInfo(

it.fetch,

func() int { return len(it.models) },

func() interface{} { b := it.models; it.models = nil; return b })

return it

}

// A ModelIterator is an iterator over Models.

type ModelIterator struct {

ctx context.Context

dataset *Dataset

models []*Model

pageInfo *iterator.PageInfo

nextFunc func() error

}

// Next returns the next result. Its second return value is Done if there are

// no more results. Once Next returns Done, all subsequent calls will return

// Done.

func (it *ModelIterator) Next() (*Model, error) {

if err := it.nextFunc(); err != nil {

return nil, err

}

t := it.models[0]

it.models = it.models[1:]

return t, nil

}

// PageInfo supports pagination. See the google.golang.org/api/iterator package for details.

func (it *ModelIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }

// listTables exists to aid testing.

var listModels = func(it *ModelIterator, pageSize int, pageToken string) (*bq.ListModelsResponse, error) {

it.ctx = setDatasetItemTraceMetadata(it.ctx, it.dataset.ProjectID, it.dataset.DatasetID, "models")

call := it.dataset.c.bqs.Models.List(it.dataset.ProjectID, it.dataset.DatasetID).

PageToken(pageToken).

Context(it.ctx)

setClientHeader(call.Header())

if pageSize > 0 {

call.MaxResults(int64(pageSize))

}

var res *bq.ListModelsResponse

err := runWithRetry(it.ctx, func() (err error) {

sCtx := trace.StartSpan(it.ctx, "bigquery.models.list")

res, err = call.Do()

trace.EndSpan(sCtx, err)

return err

})

return res, err

}

func (it *ModelIterator) fetch(pageSize int, pageToken string) (string, error) {

res, err := listModels(it, pageSize, pageToken)

if err != nil {

return "", err

}

for _, t := range res.Models {

it.models = append(it.models, bqToModel(t.ModelReference, it.dataset.c))

}

return res.NextPageToken, nil

}

func bqToModel(mr *bq.ModelReference, c *Client) *Model {

if mr == nil {

return nil

}

return &Model{

ProjectID: mr.ProjectId,

DatasetID: mr.DatasetId,

ModelID: mr.ModelId,

c: c,

}

}

// Routine creates a handle to a BigQuery routine in the dataset.

// To determine if a routine exists, call Routine.Metadata.

func (d *Dataset) Routine(routineID string) *Routine {

return &Routine{

ProjectID: d.ProjectID,

DatasetID: d.DatasetID,

RoutineID: routineID,

c: d.c}

}

// Routines returns an iterator over the routines in the Dataset.

func (d *Dataset) Routines(ctx context.Context) *RoutineIterator {

it := &RoutineIterator{

ctx: ctx,

dataset: d,

}

it.pageInfo, it.nextFunc = iterator.NewPageInfo(

it.fetch,

func() int { return len(it.routines) },

func() interface{} { b := it.routines; it.routines = nil; return b })

return it

}

// A RoutineIterator is an iterator over Routines.

type RoutineIterator struct {

ctx context.Context

dataset *Dataset

routines []*Routine

pageInfo *iterator.PageInfo

nextFunc func() error

}

// Next returns the next result. Its second return value is Done if there are

// no more results. Once Next returns Done, all subsequent calls will return

// Done.

func (it *RoutineIterator) Next() (*Routine, error) {

if err := it.nextFunc(); err != nil {

return nil, err

}

t := it.routines[0]

it.routines = it.routines[1:]

return t, nil

}

// PageInfo supports pagination. See the google.golang.org/api/iterator package for details.

func (it *RoutineIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }

// listRoutines exists to aid testing.

var listRoutines = func(it *RoutineIterator, pageSize int, pageToken string) (*bq.ListRoutinesResponse, error) {

it.ctx = setDatasetItemTraceMetadata(it.ctx, it.dataset.ProjectID, it.dataset.DatasetID, "routines")

call := it.dataset.c.bqs.Routines.List(it.dataset.ProjectID, it.dataset.DatasetID).

PageToken(pageToken).

Context(it.ctx)

setClientHeader(call.Header())

if pageSize > 0 {

call.MaxResults(int64(pageSize))

}

var res *bq.ListRoutinesResponse

err := runWithRetry(it.ctx, func() (err error) {

sCtx := trace.StartSpan(it.ctx, "bigquery.routines.list")

res, err = call.Do()

trace.EndSpan(sCtx, err)

return err

})

return res, err

}

func (it *RoutineIterator) fetch(pageSize int, pageToken string) (string, error) {

res, err := listRoutines(it, pageSize, pageToken)

if err != nil {

return "", err

}

for _, t := range res.Routines {

it.routines = append(it.routines, bqToRoutine(t.RoutineReference, it.dataset.c))

}

return res.NextPageToken, nil

}

func bqToRoutine(mr *bq.RoutineReference, c *Client) *Routine {

if mr == nil {

return nil

}

return &Routine{

ProjectID: mr.ProjectId,

DatasetID: mr.DatasetId,

RoutineID: mr.RoutineId,

c: c,

}

}

// Datasets returns an iterator over the datasets in a project.

// The Client's project is used by default, but that can be

// changed by setting ProjectID on the returned iterator before calling Next.

func (c *Client) Datasets(ctx context.Context) *DatasetIterator {

return c.DatasetsInProject(ctx, c.projectID)

}

// DatasetsInProject returns an iterator over the datasets in the provided project.

//

// Deprecated: call Client.Datasets, then set ProjectID on the returned iterator.

func (c *Client) DatasetsInProject(ctx context.Context, projectID string) *DatasetIterator {

it := &DatasetIterator{

ctx: ctx,

c: c,

ProjectID: projectID,

}

it.pageInfo, it.nextFunc = iterator.NewPageInfo(

it.fetch,

func() int { return len(it.items) },

func() interface{} { b := it.items; it.items = nil; return b })

return it

}

// DatasetIterator iterates over the datasets in a project.

type DatasetIterator struct {

// ListHidden causes hidden datasets to be listed when set to true.

// Set before the first call to Next.

ListHidden bool

// Filter restricts the datasets returned by label. The filter syntax is described in

// https://cloud.google.com/bigquery/docs/labeling-datasets#filtering_datasets_using_labels

// Set before the first call to Next.

Filter string

// The project ID of the listed datasets.

// Set before the first call to Next.

ProjectID string

ctx context.Context

c *Client

pageInfo *iterator.PageInfo

nextFunc func() error

items []*Dataset

}

// PageInfo supports pagination. See the google.golang.org/api/iterator package for details.

func (it *DatasetIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }

// Next returns the next Dataset. Its second return value is iterator.Done if

// there are no more results. Once Next returns Done, all subsequent calls will

// return Done.

func (it *DatasetIterator) Next() (*Dataset, error) {

if err := it.nextFunc(); err != nil {

return nil, err

}

item := it.items[0]

it.items = it.items[1:]

return item, nil

}

// for testing

var listDatasets = func(it *DatasetIterator, pageSize int, pageToken string) (*bq.DatasetList, error) {

it.ctx = setProjectItemTraceMetadata(it.ctx, it.ProjectID, "datasets")

call := it.c.bqs.Datasets.List(it.ProjectID).

Context(it.ctx).

PageToken(pageToken).

All(it.ListHidden)

setClientHeader(call.Header())

if pageSize > 0 {

call.MaxResults(int64(pageSize))

}

if it.Filter != "" {

call.Filter(it.Filter)

}

var res *bq.DatasetList

err := runWithRetry(it.ctx, func() (err error) {

sCtx := trace.StartSpan(it.ctx, "bigquery.datasets.list")

res, err = call.Do()

trace.EndSpan(sCtx, err)

return err

})

return res, err

}

func (it *DatasetIterator) fetch(pageSize int, pageToken string) (string, error) {

res, err := listDatasets(it, pageSize, pageToken)

if err != nil {

return "", err

}

for _, d := range res.Datasets {

it.items = append(it.items, &Dataset{

ProjectID: d.DatasetReference.ProjectId,

DatasetID: d.DatasetReference.DatasetId,

c: it.c,

})

}

return res.NextPageToken, nil

}

// An AccessEntry describes the permissions that an entity has on a dataset.

type AccessEntry struct {

Role AccessRole // The role of the entity

EntityType EntityType // The type of entity

Entity string // The entity (individual or group) granted access

View *Table // The view granted access (EntityType must be ViewEntity)

Routine *Routine // The routine granted access (only UDF currently supported)

Dataset *DatasetAccessEntry // The resources within a dataset granted access.

Condition *Expr // Condition for the access binding.

}

// Expr represents the conditional information related to dataset access policies.

type Expr struct {

// Textual representation of an expression in Common Expression Language syntax.

Expression string

// Optional. Title for the expression, i.e. a short string describing

// its purpose. This can be used e.g. in UIs which allow to enter the

// expression.

Title string

// Optional. Description of the expression. This is a longer text which

// describes the expression, e.g. when hovered over it in a UI.

Description string