dataset.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  	"context"
    19  	"errors"
    20  	"fmt"
    21  	"strings"
    22  	"time"
    23  
    24  	"cloud.google.com/go/internal/optional"
    25  	"cloud.google.com/go/internal/trace"
    26  	bq "google.golang.org/api/bigquery/v2"
    27  	"google.golang.org/api/iterator"
    28  )
    29  
    30  // Dataset is a reference to a BigQuery dataset.
    31  type Dataset struct {
    32  	ProjectID string
    33  	DatasetID string
    34  	c         *Client
    35  }
    36  
    37  // DatasetMetadata contains information about a BigQuery dataset.
    38  type DatasetMetadata struct {
    39  	// These fields can be set when creating a dataset.
    40  	Name                    string            // The user-friendly name for this dataset.
    41  	Description             string            // The user-friendly description of this dataset.
    42  	Location                string            // The geo location of the dataset.
    43  	DefaultTableExpiration  time.Duration     // The default expiration time for new tables.
    44  	Labels                  map[string]string // User-provided labels.
    45  	Access                  []*AccessEntry    // Access permissions.
    46  	DefaultEncryptionConfig *EncryptionConfig
    47  
    48  	// DefaultPartitionExpiration is the default expiration time for
    49  	// all newly created partitioned tables in the dataset.
    50  	DefaultPartitionExpiration time.Duration
    51  
    52  	// Defines the default collation specification of future tables
    53  	// created in the dataset. If a table is created in this dataset without
    54  	// table-level default collation, then the table inherits the dataset default
    55  	// collation, which is applied to the string fields that do not have explicit
    56  	// collation specified. A change to this field affects only tables created
    57  	// afterwards, and does not alter the existing tables.
    58  	// More information: https://cloud.google.com/bigquery/docs/reference/standard-sql/collation-concepts
    59  	DefaultCollation string
    60  
    61  	// For externally defined datasets, contains information about the configuration.
    62  	ExternalDatasetReference *ExternalDatasetReference
    63  
    64  	// MaxTimeTravel represents the number of hours for the max time travel for all tables
    65  	// in the dataset.  Durations are rounded towards zero for the nearest hourly value.
    66  	MaxTimeTravel time.Duration
    67  
    68  	// Storage billing model to be used for all tables in the dataset.
    69  	// Can be set to PHYSICAL. Default is LOGICAL.
    70  	// Once you create a dataset with storage billing model set to physical bytes, you can't change it back to using logical bytes again.
    71  	// More details: https://cloud.google.com/bigquery/docs/datasets-intro#dataset_storage_billing_models
    72  	StorageBillingModel string
    73  
    74  	// These fields are read-only.
    75  	CreationTime     time.Time
    76  	LastModifiedTime time.Time // When the dataset or any of its tables were modified.
    77  	FullID           string    // The full dataset ID in the form projectID:datasetID.
    78  
    79  	// The tags associated with this dataset. Tag keys are
    80  	// globally unique, and managed via the resource manager API.
    81  	// More information: https://cloud.google.com/resource-manager/docs/tags/tags-overview
    82  	Tags []*DatasetTag
    83  
    84  	// ETag is the ETag obtained when reading metadata. Pass it to Dataset.Update to
    85  	// ensure that the metadata hasn't changed since it was read.
    86  	ETag string
    87  }
    88  
    89  // DatasetTag is a representation of a single tag key/value.
    90  type DatasetTag struct {
    91  	// TagKey is the namespaced friendly name of the tag key, e.g.
    92  	// "12345/environment" where 12345 is org id.
    93  	TagKey string
    94  
    95  	// TagValue is the friendly short name of the tag value, e.g.
    96  	// "production".
    97  	TagValue string
    98  }
    99  
   100  const (
   101  	// LogicalStorageBillingModel indicates billing for logical bytes.
   102  	LogicalStorageBillingModel = ""
   103  
   104  	// PhysicalStorageBillingModel indicates billing for physical bytes.
   105  	PhysicalStorageBillingModel = "PHYSICAL"
   106  )
   107  
   108  func bqToDatasetTag(in *bq.DatasetTags) *DatasetTag {
   109  	if in == nil {
   110  		return nil
   111  	}
   112  	return &DatasetTag{
   113  		TagKey:   in.TagKey,
   114  		TagValue: in.TagValue,
   115  	}
   116  }
   117  
   118  // DatasetMetadataToUpdate is used when updating a dataset's metadata.
   119  // Only non-nil fields will be updated.
   120  type DatasetMetadataToUpdate struct {
   121  	Description optional.String // The user-friendly description of this table.
   122  	Name        optional.String // The user-friendly name for this dataset.
   123  
   124  	// DefaultTableExpiration is the default expiration time for new tables.
   125  	// If set to time.Duration(0), new tables never expire.
   126  	DefaultTableExpiration optional.Duration
   127  
   128  	// DefaultTableExpiration is the default expiration time for
   129  	// all newly created partitioned tables.
   130  	// If set to time.Duration(0), new table partitions never expire.
   131  	DefaultPartitionExpiration optional.Duration
   132  
   133  	// DefaultEncryptionConfig defines CMEK settings for new resources created
   134  	// in the dataset.
   135  	DefaultEncryptionConfig *EncryptionConfig
   136  
   137  	// Defines the default collation specification of future tables
   138  	// created in the dataset.
   139  	DefaultCollation optional.String
   140  
   141  	// For externally defined datasets, contains information about the configuration.
   142  	ExternalDatasetReference *ExternalDatasetReference
   143  
   144  	// MaxTimeTravel represents the number of hours for the max time travel for all tables
   145  	// in the dataset.  Durations are rounded towards zero for the nearest hourly value.
   146  	MaxTimeTravel optional.Duration
   147  
   148  	// Storage billing model to be used for all tables in the dataset.
   149  	// Can be set to PHYSICAL. Default is LOGICAL.
   150  	// Once you change a dataset's storage billing model to use physical bytes, you can't change it back to using logical bytes again.
   151  	// More details: https://cloud.google.com/bigquery/docs/datasets-intro#dataset_storage_billing_models
   152  	StorageBillingModel optional.String
   153  
   154  	// The entire access list. It is not possible to replace individual entries.
   155  	Access []*AccessEntry
   156  
   157  	labelUpdater
   158  }
   159  
   160  // Dataset creates a handle to a BigQuery dataset in the client's project.
   161  func (c *Client) Dataset(id string) *Dataset {
   162  	return c.DatasetInProject(c.projectID, id)
   163  }
   164  
   165  // DatasetInProject creates a handle to a BigQuery dataset in the specified project.
   166  func (c *Client) DatasetInProject(projectID, datasetID string) *Dataset {
   167  	return &Dataset{
   168  		ProjectID: projectID,
   169  		DatasetID: datasetID,
   170  		c:         c,
   171  	}
   172  }
   173  
   174  // Identifier returns the ID of the dataset in the requested format.
   175  //
   176  // For Standard SQL format, the identifier will be quoted if the
   177  // ProjectID contains dash (-) characters.
   178  func (d *Dataset) Identifier(f IdentifierFormat) (string, error) {
   179  	switch f {
   180  	case LegacySQLID:
   181  		return fmt.Sprintf("%s:%s", d.ProjectID, d.DatasetID), nil
   182  	case StandardSQLID:
   183  		// Quote project identifiers if they have a dash character.
   184  		if strings.Contains(d.ProjectID, "-") {
   185  			return fmt.Sprintf("`%s`.%s", d.ProjectID, d.DatasetID), nil
   186  		}
   187  		return fmt.Sprintf("%s.%s", d.ProjectID, d.DatasetID), nil
   188  	default:
   189  		return "", ErrUnknownIdentifierFormat
   190  	}
   191  }
   192  
   193  // Create creates a dataset in the BigQuery service. An error will be returned if the
   194  // dataset already exists. Pass in a DatasetMetadata value to configure the dataset.
   195  func (d *Dataset) Create(ctx context.Context, md *DatasetMetadata) (err error) {
   196  	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Create")
   197  	defer func() { trace.EndSpan(ctx, err) }()
   198  
   199  	ds, err := md.toBQ()
   200  	if err != nil {
   201  		return err
   202  	}
   203  	ds.DatasetReference = &bq.DatasetReference{DatasetId: d.DatasetID}
   204  	// Use Client.Location as a default.
   205  	if ds.Location == "" {
   206  		ds.Location = d.c.Location
   207  	}
   208  	call := d.c.bqs.Datasets.Insert(d.ProjectID, ds).Context(ctx)
   209  	setClientHeader(call.Header())
   210  	_, err = call.Do()
   211  	return err
   212  }
   213  
   214  func (dm *DatasetMetadata) toBQ() (*bq.Dataset, error) {
   215  	ds := &bq.Dataset{}
   216  	if dm == nil {
   217  		return ds, nil
   218  	}
   219  	ds.FriendlyName = dm.Name
   220  	ds.Description = dm.Description
   221  	ds.Location = dm.Location
   222  	ds.DefaultTableExpirationMs = int64(dm.DefaultTableExpiration / time.Millisecond)
   223  	ds.DefaultPartitionExpirationMs = int64(dm.DefaultPartitionExpiration / time.Millisecond)
   224  	ds.DefaultCollation = dm.DefaultCollation
   225  	ds.MaxTimeTravelHours = int64(dm.MaxTimeTravel / time.Hour)
   226  	ds.StorageBillingModel = string(dm.StorageBillingModel)
   227  	ds.Labels = dm.Labels
   228  	var err error
   229  	ds.Access, err = accessListToBQ(dm.Access)
   230  	if err != nil {
   231  		return nil, err
   232  	}
   233  	if !dm.CreationTime.IsZero() {
   234  		return nil, errors.New("bigquery: Dataset.CreationTime is not writable")
   235  	}
   236  	if !dm.LastModifiedTime.IsZero() {
   237  		return nil, errors.New("bigquery: Dataset.LastModifiedTime is not writable")
   238  	}
   239  	if dm.FullID != "" {
   240  		return nil, errors.New("bigquery: Dataset.FullID is not writable")
   241  	}
   242  	if dm.ETag != "" {
   243  		return nil, errors.New("bigquery: Dataset.ETag is not writable")
   244  	}
   245  	if dm.DefaultEncryptionConfig != nil {
   246  		ds.DefaultEncryptionConfiguration = dm.DefaultEncryptionConfig.toBQ()
   247  	}
   248  	if dm.ExternalDatasetReference != nil {
   249  		ds.ExternalDatasetReference = dm.ExternalDatasetReference.toBQ()
   250  	}
   251  	return ds, nil
   252  }
   253  
   254  func accessListToBQ(a []*AccessEntry) ([]*bq.DatasetAccess, error) {
   255  	var q []*bq.DatasetAccess
   256  	for _, e := range a {
   257  		a, err := e.toBQ()
   258  		if err != nil {
   259  			return nil, err
   260  		}
   261  		q = append(q, a)
   262  	}
   263  	return q, nil
   264  }
   265  
   266  // Delete deletes the dataset.  Delete will fail if the dataset is not empty.
   267  func (d *Dataset) Delete(ctx context.Context) (err error) {
   268  	return d.deleteInternal(ctx, false)
   269  }
   270  
   271  // DeleteWithContents deletes the dataset, as well as contained resources.
   272  func (d *Dataset) DeleteWithContents(ctx context.Context) (err error) {
   273  	return d.deleteInternal(ctx, true)
   274  }
   275  
   276  func (d *Dataset) deleteInternal(ctx context.Context, deleteContents bool) (err error) {
   277  	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Delete")
   278  	defer func() { trace.EndSpan(ctx, err) }()
   279  
   280  	call := d.c.bqs.Datasets.Delete(d.ProjectID, d.DatasetID).Context(ctx).DeleteContents(deleteContents)
   281  	setClientHeader(call.Header())
   282  	return runWithRetry(ctx, func() (err error) {
   283  		sCtx := trace.StartSpan(ctx, "bigquery.datasets.delete")
   284  		err = call.Do()
   285  		trace.EndSpan(sCtx, err)
   286  		return err
   287  	})
   288  }
   289  
   290  // Metadata fetches the metadata for the dataset.
   291  func (d *Dataset) Metadata(ctx context.Context) (md *DatasetMetadata, err error) {
   292  	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Metadata")
   293  	defer func() { trace.EndSpan(ctx, err) }()
   294  
   295  	call := d.c.bqs.Datasets.Get(d.ProjectID, d.DatasetID).Context(ctx)
   296  	setClientHeader(call.Header())
   297  	var ds *bq.Dataset
   298  	if err := runWithRetry(ctx, func() (err error) {
   299  		sCtx := trace.StartSpan(ctx, "bigquery.datasets.get")
   300  		ds, err = call.Do()
   301  		trace.EndSpan(sCtx, err)
   302  		return err
   303  	}); err != nil {
   304  		return nil, err
   305  	}
   306  	return bqToDatasetMetadata(ds, d.c)
   307  }
   308  
   309  func bqToDatasetMetadata(d *bq.Dataset, c *Client) (*DatasetMetadata, error) {
   310  	dm := &DatasetMetadata{
   311  		CreationTime:               unixMillisToTime(d.CreationTime),
   312  		LastModifiedTime:           unixMillisToTime(d.LastModifiedTime),
   313  		DefaultTableExpiration:     time.Duration(d.DefaultTableExpirationMs) * time.Millisecond,
   314  		DefaultPartitionExpiration: time.Duration(d.DefaultPartitionExpirationMs) * time.Millisecond,
   315  		DefaultCollation:           d.DefaultCollation,
   316  		ExternalDatasetReference:   bqToExternalDatasetReference(d.ExternalDatasetReference),
   317  		MaxTimeTravel:              time.Duration(d.MaxTimeTravelHours) * time.Hour,
   318  		StorageBillingModel:        d.StorageBillingModel,
   319  		DefaultEncryptionConfig:    bqToEncryptionConfig(d.DefaultEncryptionConfiguration),
   320  		Description:                d.Description,
   321  		Name:                       d.FriendlyName,
   322  		FullID:                     d.Id,
   323  		Location:                   d.Location,
   324  		Labels:                     d.Labels,
   325  		ETag:                       d.Etag,
   326  	}
   327  	for _, a := range d.Access {
   328  		e, err := bqToAccessEntry(a, c)
   329  		if err != nil {
   330  			return nil, err
   331  		}
   332  		dm.Access = append(dm.Access, e)
   333  	}
   334  	for _, bqTag := range d.Tags {
   335  		tag := bqToDatasetTag(bqTag)
   336  		if tag != nil {
   337  			dm.Tags = append(dm.Tags, tag)
   338  		}
   339  	}
   340  	return dm, nil
   341  }
   342  
   343  // Update modifies specific Dataset metadata fields.
   344  // To perform a read-modify-write that protects against intervening reads,
   345  // set the etag argument to the DatasetMetadata.ETag field from the read.
   346  // Pass the empty string for etag for a "blind write" that will always succeed.
   347  func (d *Dataset) Update(ctx context.Context, dm DatasetMetadataToUpdate, etag string) (md *DatasetMetadata, err error) {
   348  	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.Dataset.Update")
   349  	defer func() { trace.EndSpan(ctx, err) }()
   350  
   351  	ds, err := dm.toBQ()
   352  	if err != nil {
   353  		return nil, err
   354  	}
   355  	call := d.c.bqs.Datasets.Patch(d.ProjectID, d.DatasetID, ds).Context(ctx)
   356  	setClientHeader(call.Header())
   357  	if etag != "" {
   358  		call.Header().Set("If-Match", etag)
   359  	}
   360  	var ds2 *bq.Dataset
   361  	if err := runWithRetry(ctx, func() (err error) {
   362  		sCtx := trace.StartSpan(ctx, "bigquery.datasets.patch")
   363  		ds2, err = call.Do()
   364  		trace.EndSpan(sCtx, err)
   365  		return err
   366  	}); err != nil {
   367  		return nil, err
   368  	}
   369  	return bqToDatasetMetadata(ds2, d.c)
   370  }
   371  
   372  func (dm *DatasetMetadataToUpdate) toBQ() (*bq.Dataset, error) {
   373  	ds := &bq.Dataset{}
   374  	forceSend := func(field string) {
   375  		ds.ForceSendFields = append(ds.ForceSendFields, field)
   376  	}
   377  
   378  	if dm.Description != nil {
   379  		ds.Description = optional.ToString(dm.Description)
   380  		forceSend("Description")
   381  	}
   382  	if dm.Name != nil {
   383  		ds.FriendlyName = optional.ToString(dm.Name)
   384  		forceSend("FriendlyName")
   385  	}
   386  	if dm.DefaultTableExpiration != nil {
   387  		dur := optional.ToDuration(dm.DefaultTableExpiration)
   388  		if dur == 0 {
   389  			// Send a null to delete the field.
   390  			ds.NullFields = append(ds.NullFields, "DefaultTableExpirationMs")
   391  		} else {
   392  			ds.DefaultTableExpirationMs = int64(dur / time.Millisecond)
   393  		}
   394  	}
   395  	if dm.DefaultPartitionExpiration != nil {
   396  		dur := optional.ToDuration(dm.DefaultPartitionExpiration)
   397  		if dur == 0 {
   398  			// Send a null to delete the field.
   399  			ds.NullFields = append(ds.NullFields, "DefaultPartitionExpirationMs")
   400  		} else {
   401  			ds.DefaultPartitionExpirationMs = int64(dur / time.Millisecond)
   402  		}
   403  	}
   404  	if dm.DefaultCollation != nil {
   405  		ds.DefaultCollation = optional.ToString(dm.DefaultCollation)
   406  		forceSend("DefaultCollation")
   407  	}
   408  	if dm.ExternalDatasetReference != nil {
   409  		ds.ExternalDatasetReference = dm.ExternalDatasetReference.toBQ()
   410  		forceSend("ExternalDatasetReference")
   411  	}
   412  	if dm.MaxTimeTravel != nil {
   413  		dur := optional.ToDuration(dm.MaxTimeTravel)
   414  		if dur == 0 {
   415  			// Send a null to delete the field.
   416  			ds.NullFields = append(ds.NullFields, "MaxTimeTravelHours")
   417  		} else {
   418  			ds.MaxTimeTravelHours = int64(dur / time.Hour)
   419  		}
   420  	}
   421  	if dm.StorageBillingModel != nil {
   422  		ds.StorageBillingModel = optional.ToString(dm.StorageBillingModel)
   423  		forceSend("StorageBillingModel")
   424  	}
   425  	if dm.DefaultEncryptionConfig != nil {
   426  		ds.DefaultEncryptionConfiguration = dm.DefaultEncryptionConfig.toBQ()
   427  		ds.DefaultEncryptionConfiguration.ForceSendFields = []string{"KmsKeyName"}
   428  	}
   429  	if dm.Access != nil {
   430  		var err error
   431  		ds.Access, err = accessListToBQ(dm.Access)
   432  		if err != nil {
   433  			return nil, err
   434  		}
   435  		if len(ds.Access) == 0 {
   436  			ds.NullFields = append(ds.NullFields, "Access")
   437  		}
   438  	}
   439  	labels, forces, nulls := dm.update()
   440  	ds.Labels = labels
   441  	ds.ForceSendFields = append(ds.ForceSendFields, forces...)
   442  	ds.NullFields = append(ds.NullFields, nulls...)
   443  	return ds, nil
   444  }
   445  
   446  // Table creates a handle to a BigQuery table in the dataset.
   447  // To determine if a table exists, call Table.Metadata.
   448  // If the table does not already exist, use Table.Create to create it.
   449  func (d *Dataset) Table(tableID string) *Table {
   450  	return &Table{ProjectID: d.ProjectID, DatasetID: d.DatasetID, TableID: tableID, c: d.c}
   451  }
   452  
   453  // Tables returns an iterator over the tables in the Dataset.
   454  func (d *Dataset) Tables(ctx context.Context) *TableIterator {
   455  	it := &TableIterator{
   456  		ctx:     ctx,
   457  		dataset: d,
   458  	}
   459  	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
   460  		it.fetch,
   461  		func() int { return len(it.tables) },
   462  		func() interface{} { b := it.tables; it.tables = nil; return b })
   463  	return it
   464  }
   465  
   466  // A TableIterator is an iterator over Tables.
   467  type TableIterator struct {
   468  	ctx      context.Context
   469  	dataset  *Dataset
   470  	tables   []*Table
   471  	pageInfo *iterator.PageInfo
   472  	nextFunc func() error
   473  }
   474  
   475  // Next returns the next result. Its second return value is Done if there are
   476  // no more results. Once Next returns Done, all subsequent calls will return
   477  // Done.
   478  func (it *TableIterator) Next() (*Table, error) {
   479  	if err := it.nextFunc(); err != nil {
   480  		return nil, err
   481  	}
   482  	t := it.tables[0]
   483  	it.tables = it.tables[1:]
   484  	return t, nil
   485  }
   486  
   487  // PageInfo supports pagination. See the google.golang.org/api/iterator package for details.
   488  func (it *TableIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }
   489  
   490  // listTables exists to aid testing.
   491  var listTables = func(it *TableIterator, pageSize int, pageToken string) (*bq.TableList, error) {
   492  	call := it.dataset.c.bqs.Tables.List(it.dataset.ProjectID, it.dataset.DatasetID).
   493  		PageToken(pageToken).
   494  		Context(it.ctx)
   495  	setClientHeader(call.Header())
   496  	if pageSize > 0 {
   497  		call.MaxResults(int64(pageSize))
   498  	}
   499  	var res *bq.TableList
   500  	err := runWithRetry(it.ctx, func() (err error) {
   501  		sCtx := trace.StartSpan(it.ctx, "bigquery.tables.list")
   502  		res, err = call.Do()
   503  		trace.EndSpan(sCtx, err)
   504  		return err
   505  	})
   506  	return res, err
   507  }
   508  
   509  func (it *TableIterator) fetch(pageSize int, pageToken string) (string, error) {
   510  	res, err := listTables(it, pageSize, pageToken)
   511  	if err != nil {
   512  		return "", err
   513  	}
   514  	for _, t := range res.Tables {
   515  		it.tables = append(it.tables, bqToTable(t.TableReference, it.dataset.c))
   516  	}
   517  	return res.NextPageToken, nil
   518  }
   519  
   520  func bqToTable(tr *bq.TableReference, c *Client) *Table {
   521  	if tr == nil {
   522  		return nil
   523  	}
   524  	return &Table{
   525  		ProjectID: tr.ProjectId,
   526  		DatasetID: tr.DatasetId,
   527  		TableID:   tr.TableId,
   528  		c:         c,
   529  	}
   530  }
   531  
   532  // Model creates a handle to a BigQuery model in the dataset.
   533  // To determine if a model exists, call Model.Metadata.
   534  // If the model does not already exist, you can create it via execution
   535  // of a CREATE MODEL query.
   536  func (d *Dataset) Model(modelID string) *Model {
   537  	return &Model{ProjectID: d.ProjectID, DatasetID: d.DatasetID, ModelID: modelID, c: d.c}
   538  }
   539  
   540  // Models returns an iterator over the models in the Dataset.
   541  func (d *Dataset) Models(ctx context.Context) *ModelIterator {
   542  	it := &ModelIterator{
   543  		ctx:     ctx,
   544  		dataset: d,
   545  	}
   546  	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
   547  		it.fetch,
   548  		func() int { return len(it.models) },
   549  		func() interface{} { b := it.models; it.models = nil; return b })
   550  	return it
   551  }
   552  
   553  // A ModelIterator is an iterator over Models.
   554  type ModelIterator struct {
   555  	ctx      context.Context
   556  	dataset  *Dataset
   557  	models   []*Model
   558  	pageInfo *iterator.PageInfo
   559  	nextFunc func() error
   560  }
   561  
   562  // Next returns the next result. Its second return value is Done if there are
   563  // no more results. Once Next returns Done, all subsequent calls will return
   564  // Done.
   565  func (it *ModelIterator) Next() (*Model, error) {
   566  	if err := it.nextFunc(); err != nil {
   567  		return nil, err
   568  	}
   569  	t := it.models[0]
   570  	it.models = it.models[1:]
   571  	return t, nil
   572  }
   573  
   574  // PageInfo supports pagination. See the google.golang.org/api/iterator package for details.
   575  func (it *ModelIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }
   576  
   577  // listTables exists to aid testing.
   578  var listModels = func(it *ModelIterator, pageSize int, pageToken string) (*bq.ListModelsResponse, error) {
   579  	call := it.dataset.c.bqs.Models.List(it.dataset.ProjectID, it.dataset.DatasetID).
   580  		PageToken(pageToken).
   581  		Context(it.ctx)
   582  	setClientHeader(call.Header())
   583  	if pageSize > 0 {
   584  		call.MaxResults(int64(pageSize))
   585  	}
   586  	var res *bq.ListModelsResponse
   587  	err := runWithRetry(it.ctx, func() (err error) {
   588  		sCtx := trace.StartSpan(it.ctx, "bigquery.models.list")
   589  		res, err = call.Do()
   590  		trace.EndSpan(sCtx, err)
   591  		return err
   592  	})
   593  	return res, err
   594  }
   595  
   596  func (it *ModelIterator) fetch(pageSize int, pageToken string) (string, error) {
   597  	res, err := listModels(it, pageSize, pageToken)
   598  	if err != nil {
   599  		return "", err
   600  	}
   601  	for _, t := range res.Models {
   602  		it.models = append(it.models, bqToModel(t.ModelReference, it.dataset.c))
   603  	}
   604  	return res.NextPageToken, nil
   605  }
   606  
   607  func bqToModel(mr *bq.ModelReference, c *Client) *Model {
   608  	if mr == nil {
   609  		return nil
   610  	}
   611  	return &Model{
   612  		ProjectID: mr.ProjectId,
   613  		DatasetID: mr.DatasetId,
   614  		ModelID:   mr.ModelId,
   615  		c:         c,
   616  	}
   617  }
   618  
   619  // Routine creates a handle to a BigQuery routine in the dataset.
   620  // To determine if a routine exists, call Routine.Metadata.
   621  func (d *Dataset) Routine(routineID string) *Routine {
   622  	return &Routine{
   623  		ProjectID: d.ProjectID,
   624  		DatasetID: d.DatasetID,
   625  		RoutineID: routineID,
   626  		c:         d.c}
   627  }
   628  
   629  // Routines returns an iterator over the routines in the Dataset.
   630  func (d *Dataset) Routines(ctx context.Context) *RoutineIterator {
   631  	it := &RoutineIterator{
   632  		ctx:     ctx,
   633  		dataset: d,
   634  	}
   635  	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
   636  		it.fetch,
   637  		func() int { return len(it.routines) },
   638  		func() interface{} { b := it.routines; it.routines = nil; return b })
   639  	return it
   640  }
   641  
   642  // A RoutineIterator is an iterator over Routines.
   643  type RoutineIterator struct {
   644  	ctx      context.Context
   645  	dataset  *Dataset
   646  	routines []*Routine
   647  	pageInfo *iterator.PageInfo
   648  	nextFunc func() error
   649  }
   650  
   651  // Next returns the next result. Its second return value is Done if there are
   652  // no more results. Once Next returns Done, all subsequent calls will return
   653  // Done.
   654  func (it *RoutineIterator) Next() (*Routine, error) {
   655  	if err := it.nextFunc(); err != nil {
   656  		return nil, err
   657  	}
   658  	t := it.routines[0]
   659  	it.routines = it.routines[1:]
   660  	return t, nil
   661  }
   662  
   663  // PageInfo supports pagination. See the google.golang.org/api/iterator package for details.
   664  func (it *RoutineIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }
   665  
   666  // listRoutines exists to aid testing.
   667  var listRoutines = func(it *RoutineIterator, pageSize int, pageToken string) (*bq.ListRoutinesResponse, error) {
   668  	call := it.dataset.c.bqs.Routines.List(it.dataset.ProjectID, it.dataset.DatasetID).
   669  		PageToken(pageToken).
   670  		Context(it.ctx)
   671  	setClientHeader(call.Header())
   672  	if pageSize > 0 {
   673  		call.MaxResults(int64(pageSize))
   674  	}
   675  	var res *bq.ListRoutinesResponse
   676  	err := runWithRetry(it.ctx, func() (err error) {
   677  		sCtx := trace.StartSpan(it.ctx, "bigquery.routines.list")
   678  		res, err = call.Do()
   679  		trace.EndSpan(sCtx, err)
   680  		return err
   681  	})
   682  	return res, err
   683  }
   684  
   685  func (it *RoutineIterator) fetch(pageSize int, pageToken string) (string, error) {
   686  	res, err := listRoutines(it, pageSize, pageToken)
   687  	if err != nil {
   688  		return "", err
   689  	}
   690  	for _, t := range res.Routines {
   691  		it.routines = append(it.routines, bqToRoutine(t.RoutineReference, it.dataset.c))
   692  	}
   693  	return res.NextPageToken, nil
   694  }
   695  
   696  func bqToRoutine(mr *bq.RoutineReference, c *Client) *Routine {
   697  	if mr == nil {
   698  		return nil
   699  	}
   700  	return &Routine{
   701  		ProjectID: mr.ProjectId,
   702  		DatasetID: mr.DatasetId,
   703  		RoutineID: mr.RoutineId,
   704  		c:         c,
   705  	}
   706  }
   707  
   708  // Datasets returns an iterator over the datasets in a project.
   709  // The Client's project is used by default, but that can be
   710  // changed by setting ProjectID on the returned iterator before calling Next.
   711  func (c *Client) Datasets(ctx context.Context) *DatasetIterator {
   712  	return c.DatasetsInProject(ctx, c.projectID)
   713  }
   714  
   715  // DatasetsInProject returns an iterator over the datasets in the provided project.
   716  //
   717  // Deprecated: call Client.Datasets, then set ProjectID on the returned iterator.
   718  func (c *Client) DatasetsInProject(ctx context.Context, projectID string) *DatasetIterator {
   719  	it := &DatasetIterator{
   720  		ctx:       ctx,
   721  		c:         c,
   722  		ProjectID: projectID,
   723  	}
   724  	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
   725  		it.fetch,
   726  		func() int { return len(it.items) },
   727  		func() interface{} { b := it.items; it.items = nil; return b })
   728  	return it
   729  }
   730  
   731  // DatasetIterator iterates over the datasets in a project.
   732  type DatasetIterator struct {
   733  	// ListHidden causes hidden datasets to be listed when set to true.
   734  	// Set before the first call to Next.
   735  	ListHidden bool
   736  
   737  	// Filter restricts the datasets returned by label. The filter syntax is described in
   738  	// https://cloud.google.com/bigquery/docs/labeling-datasets#filtering_datasets_using_labels
   739  	// Set before the first call to Next.
   740  	Filter string
   741  
   742  	// The project ID of the listed datasets.
   743  	// Set before the first call to Next.
   744  	ProjectID string
   745  
   746  	ctx      context.Context
   747  	c        *Client
   748  	pageInfo *iterator.PageInfo
   749  	nextFunc func() error
   750  	items    []*Dataset
   751  }
   752  
   753  // PageInfo supports pagination. See the google.golang.org/api/iterator package for details.
   754  func (it *DatasetIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }
   755  
   756  // Next returns the next Dataset. Its second return value is iterator.Done if
   757  // there are no more results. Once Next returns Done, all subsequent calls will
   758  // return Done.
   759  func (it *DatasetIterator) Next() (*Dataset, error) {
   760  	if err := it.nextFunc(); err != nil {
   761  		return nil, err
   762  	}
   763  	item := it.items[0]
   764  	it.items = it.items[1:]
   765  	return item, nil
   766  }
   767  
   768  // for testing
   769  var listDatasets = func(it *DatasetIterator, pageSize int, pageToken string) (*bq.DatasetList, error) {
   770  	call := it.c.bqs.Datasets.List(it.ProjectID).
   771  		Context(it.ctx).
   772  		PageToken(pageToken).
   773  		All(it.ListHidden)
   774  	setClientHeader(call.Header())
   775  	if pageSize > 0 {
   776  		call.MaxResults(int64(pageSize))
   777  	}
   778  	if it.Filter != "" {
   779  		call.Filter(it.Filter)
   780  	}
   781  	var res *bq.DatasetList
   782  	err := runWithRetry(it.ctx, func() (err error) {
   783  		sCtx := trace.StartSpan(it.ctx, "bigquery.datasets.list")
   784  		res, err = call.Do()
   785  		trace.EndSpan(sCtx, err)
   786  		return err
   787  	})
   788  	return res, err
   789  }
   790  
   791  func (it *DatasetIterator) fetch(pageSize int, pageToken string) (string, error) {
   792  	res, err := listDatasets(it, pageSize, pageToken)
   793  	if err != nil {
   794  		return "", err
   795  	}
   796  	for _, d := range res.Datasets {
   797  		it.items = append(it.items, &Dataset{
   798  			ProjectID: d.DatasetReference.ProjectId,
   799  			DatasetID: d.DatasetReference.DatasetId,
   800  			c:         it.c,
   801  		})
   802  	}
   803  	return res.NextPageToken, nil
   804  }
   805  
   806  // An AccessEntry describes the permissions that an entity has on a dataset.
   807  type AccessEntry struct {
   808  	Role       AccessRole          // The role of the entity
   809  	EntityType EntityType          // The type of entity
   810  	Entity     string              // The entity (individual or group) granted access
   811  	View       *Table              // The view granted access (EntityType must be ViewEntity)
   812  	Routine    *Routine            // The routine granted access (only UDF currently supported)
   813  	Dataset    *DatasetAccessEntry // The resources within a dataset granted access.
   814  }
   815  
   816  // AccessRole is the level of access to grant to a dataset.
   817  type AccessRole string
   818  
   819  const (
   820  	// OwnerRole is the OWNER AccessRole.
   821  	OwnerRole AccessRole = "OWNER"
   822  	// ReaderRole is the READER AccessRole.
   823  	ReaderRole AccessRole = "READER"
   824  	// WriterRole is the WRITER AccessRole.
   825  	WriterRole AccessRole = "WRITER"
   826  )
   827  
   828  // EntityType is the type of entity in an AccessEntry.
   829  type EntityType int
   830  
   831  const (
   832  	// DomainEntity is a domain (e.g. "example.com").
   833  	DomainEntity EntityType = iota + 1
   834  
   835  	// GroupEmailEntity is an email address of a Google Group.
   836  	GroupEmailEntity
   837  
   838  	// UserEmailEntity is an email address of an individual user.
   839  	UserEmailEntity
   840  
   841  	// SpecialGroupEntity is a special group: one of projectOwners, projectReaders, projectWriters or
   842  	// allAuthenticatedUsers.
   843  	SpecialGroupEntity
   844  
   845  	// ViewEntity is a BigQuery logical view.
   846  	ViewEntity
   847  
   848  	// IAMMemberEntity represents entities present in IAM but not represented using
   849  	// the other entity types.
   850  	IAMMemberEntity
   851  
   852  	// RoutineEntity is a BigQuery routine, referencing a User Defined Function (UDF).
   853  	RoutineEntity
   854  
   855  	// DatasetEntity is BigQuery dataset, present in the access list.
   856  	DatasetEntity
   857  )
   858  
   859  func (e *AccessEntry) toBQ() (*bq.DatasetAccess, error) {
   860  	q := &bq.DatasetAccess{Role: string(e.Role)}
   861  	switch e.EntityType {
   862  	case DomainEntity:
   863  		q.Domain = e.Entity
   864  	case GroupEmailEntity:
   865  		q.GroupByEmail = e.Entity
   866  	case UserEmailEntity:
   867  		q.UserByEmail = e.Entity
   868  	case SpecialGroupEntity:
   869  		q.SpecialGroup = e.Entity
   870  	case ViewEntity:
   871  		q.View = e.View.toBQ()
   872  	case IAMMemberEntity:
   873  		q.IamMember = e.Entity
   874  	case RoutineEntity:
   875  		q.Routine = e.Routine.toBQ()
   876  	case DatasetEntity:
   877  		q.Dataset = e.Dataset.toBQ()
   878  	default:
   879  		return nil, fmt.Errorf("bigquery: unknown entity type %d", e.EntityType)
   880  	}
   881  	return q, nil
   882  }
   883  
   884  func bqToAccessEntry(q *bq.DatasetAccess, c *Client) (*AccessEntry, error) {
   885  	e := &AccessEntry{Role: AccessRole(q.Role)}
   886  	switch {
   887  	case q.Domain != "":
   888  		e.Entity = q.Domain
   889  		e.EntityType = DomainEntity
   890  	case q.GroupByEmail != "":
   891  		e.Entity = q.GroupByEmail
   892  		e.EntityType = GroupEmailEntity
   893  	case q.UserByEmail != "":
   894  		e.Entity = q.UserByEmail
   895  		e.EntityType = UserEmailEntity
   896  	case q.SpecialGroup != "":
   897  		e.Entity = q.SpecialGroup
   898  		e.EntityType = SpecialGroupEntity
   899  	case q.View != nil:
   900  		e.View = c.DatasetInProject(q.View.ProjectId, q.View.DatasetId).Table(q.View.TableId)
   901  		e.EntityType = ViewEntity
   902  	case q.IamMember != "":
   903  		e.Entity = q.IamMember
   904  		e.EntityType = IAMMemberEntity
   905  	case q.Routine != nil:
   906  		e.Routine = c.DatasetInProject(q.Routine.ProjectId, q.Routine.DatasetId).Routine(q.Routine.RoutineId)
   907  		e.EntityType = RoutineEntity
   908  	case q.Dataset != nil:
   909  		e.Dataset = bqToDatasetAccessEntry(q.Dataset, c)
   910  		e.EntityType = DatasetEntity
   911  	default:
   912  		return nil, errors.New("bigquery: invalid access value")
   913  	}
   914  	return e, nil
   915  }
   916  
   917  // DatasetAccessEntry is an access entry that refers to resources within
   918  // another dataset.
   919  type DatasetAccessEntry struct {
   920  	// The dataset to which this entry applies.
   921  	Dataset *Dataset
   922  	// The list of target types within the dataset
   923  	// to which this entry applies.
   924  	//
   925  	// Current supported values:
   926  	//
   927  	// VIEWS - This entry applies to views in the dataset.
   928  	TargetTypes []string
   929  }
   930  
   931  func (dae *DatasetAccessEntry) toBQ() *bq.DatasetAccessEntry {
   932  	if dae == nil {
   933  		return nil
   934  	}
   935  	return &bq.DatasetAccessEntry{
   936  		Dataset: &bq.DatasetReference{
   937  			ProjectId: dae.Dataset.ProjectID,
   938  			DatasetId: dae.Dataset.DatasetID,
   939  		},
   940  		TargetTypes: dae.TargetTypes,
   941  	}
   942  }
   943  
   944  func bqToDatasetAccessEntry(entry *bq.DatasetAccessEntry, c *Client) *DatasetAccessEntry {
   945  	if entry == nil {
   946  		return nil
   947  	}
   948  	return &DatasetAccessEntry{
   949  		Dataset:     c.DatasetInProject(entry.Dataset.ProjectId, entry.Dataset.DatasetId),
   950  		TargetTypes: entry.TargetTypes,
   951  	}
   952  }
   953  
   954  // ExternalDatasetReference provides information about external dataset metadata.
   955  type ExternalDatasetReference struct {
   956  	//The connection id that is used to access the external_source.
   957  	// Format: projects/{project_id}/locations/{location_id}/connections/{connection_id}
   958  	Connection string
   959  
   960  	// External source that backs this dataset.
   961  	ExternalSource string
   962  }
   963  
   964  func bqToExternalDatasetReference(bq *bq.ExternalDatasetReference) *ExternalDatasetReference {
   965  	if bq == nil {
   966  		return nil
   967  	}
   968  	return &ExternalDatasetReference{
   969  		Connection:     bq.Connection,
   970  		ExternalSource: bq.ExternalSource,
   971  	}
   972  }
   973  
   974  func (edr *ExternalDatasetReference) toBQ() *bq.ExternalDatasetReference {
   975  	if edr == nil {
   976  		return nil
   977  	}
   978  	return &bq.ExternalDatasetReference{
   979  		Connection:     edr.Connection,
   980  		ExternalSource: edr.ExternalSource,
   981  	}
   982  }
   983
View as plain text