You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

165 lines
4.9 KiB

  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. package bigquery
  15. import (
  16. "fmt"
  17. "io"
  18. "net/http"
  19. "time"
  20. gax "github.com/googleapis/gax-go"
  21. "cloud.google.com/go/internal"
  22. "cloud.google.com/go/internal/version"
  23. "google.golang.org/api/googleapi"
  24. "google.golang.org/api/option"
  25. htransport "google.golang.org/api/transport/http"
  26. "golang.org/x/net/context"
  27. bq "google.golang.org/api/bigquery/v2"
  28. )
  29. const (
  30. prodAddr = "https://www.googleapis.com/bigquery/v2/"
  31. Scope = "https://www.googleapis.com/auth/bigquery"
  32. userAgent = "gcloud-golang-bigquery/20160429"
  33. )
  34. var xGoogHeader = fmt.Sprintf("gl-go/%s gccl/%s", version.Go(), version.Repo)
  35. func setClientHeader(headers http.Header) {
  36. headers.Set("x-goog-api-client", xGoogHeader)
  37. }
  38. // Client may be used to perform BigQuery operations.
  39. type Client struct {
  40. // Location, if set, will be used as the default location for all subsequent
  41. // dataset creation and job operations. A location specified directly in one of
  42. // those operations will override this value.
  43. Location string
  44. projectID string
  45. bqs *bq.Service
  46. }
  47. // NewClient constructs a new Client which can perform BigQuery operations.
  48. // Operations performed via the client are billed to the specified GCP project.
  49. func NewClient(ctx context.Context, projectID string, opts ...option.ClientOption) (*Client, error) {
  50. o := []option.ClientOption{
  51. option.WithEndpoint(prodAddr),
  52. option.WithScopes(Scope),
  53. option.WithUserAgent(userAgent),
  54. }
  55. o = append(o, opts...)
  56. httpClient, endpoint, err := htransport.NewClient(ctx, o...)
  57. if err != nil {
  58. return nil, fmt.Errorf("bigquery: dialing: %v", err)
  59. }
  60. bqs, err := bq.New(httpClient)
  61. if err != nil {
  62. return nil, fmt.Errorf("bigquery: constructing client: %v", err)
  63. }
  64. bqs.BasePath = endpoint
  65. c := &Client{
  66. projectID: projectID,
  67. bqs: bqs,
  68. }
  69. return c, nil
  70. }
  71. // Close closes any resources held by the client.
  72. // Close should be called when the client is no longer needed.
  73. // It need not be called at program exit.
  74. func (c *Client) Close() error {
  75. return nil
  76. }
  77. // Calls the Jobs.Insert RPC and returns a Job.
  78. func (c *Client) insertJob(ctx context.Context, job *bq.Job, media io.Reader) (*Job, error) {
  79. call := c.bqs.Jobs.Insert(c.projectID, job).Context(ctx)
  80. setClientHeader(call.Header())
  81. if media != nil {
  82. call.Media(media)
  83. }
  84. var res *bq.Job
  85. var err error
  86. invoke := func() error {
  87. res, err = call.Do()
  88. return err
  89. }
  90. // A job with a client-generated ID can be retried; the presence of the
  91. // ID makes the insert operation idempotent.
  92. // We don't retry if there is media, because it is an io.Reader. We'd
  93. // have to read the contents and keep it in memory, and that could be expensive.
  94. // TODO(jba): Look into retrying if media != nil.
  95. if job.JobReference != nil && media == nil {
  96. err = runWithRetry(ctx, invoke)
  97. } else {
  98. err = invoke()
  99. }
  100. if err != nil {
  101. return nil, err
  102. }
  103. return bqToJob(res, c)
  104. }
  105. // Convert a number of milliseconds since the Unix epoch to a time.Time.
  106. // Treat an input of zero specially: convert it to the zero time,
  107. // rather than the start of the epoch.
  108. func unixMillisToTime(m int64) time.Time {
  109. if m == 0 {
  110. return time.Time{}
  111. }
  112. return time.Unix(0, m*1e6)
  113. }
  114. // runWithRetry calls the function until it returns nil or a non-retryable error, or
  115. // the context is done.
  116. // See the similar function in ../storage/invoke.go. The main difference is the
  117. // reason for retrying.
  118. func runWithRetry(ctx context.Context, call func() error) error {
  119. // These parameters match the suggestions in https://cloud.google.com/bigquery/sla.
  120. backoff := gax.Backoff{
  121. Initial: 1 * time.Second,
  122. Max: 32 * time.Second,
  123. Multiplier: 2,
  124. }
  125. return internal.Retry(ctx, backoff, func() (stop bool, err error) {
  126. err = call()
  127. if err == nil {
  128. return true, nil
  129. }
  130. return !retryableError(err), err
  131. })
  132. }
  133. // This is the correct definition of retryable according to the BigQuery team. It
  134. // also considers 502 ("Bad Gateway") and 503 ("Service Unavailable") errors
  135. // retryable; these are returned by systems between the client and the BigQuery
  136. // service.
  137. func retryableError(err error) bool {
  138. e, ok := err.(*googleapi.Error)
  139. if !ok {
  140. return false
  141. }
  142. var reason string
  143. if len(e.Errors) > 0 {
  144. reason = e.Errors[0].Reason
  145. }
  146. return e.Code == http.StatusServiceUnavailable || e.Code == http.StatusBadGateway || reason == "backendError" || reason == "rateLimitExceeded"
  147. }