Skip to content

Commit dc66aca

Browse files
authored
feat: add support for specifying a quota project (GoogleCloudPlatform#1044)
This commit adds a CLI flag quota_project such that users may provide a project against which the SQL Admin API quota is calculated. Typically, this flag will be used for internal development purposes. Note: there is an API option WithQuotaProject that does the same thing. However, that option is not compatible with WithHTTPClient and so we manually add the quota project instead.
1 parent 2b3df42 commit dc66aca

1 file changed

Lines changed: 44 additions & 1 deletion

File tree

cmd/cloud_sql_proxy/cloud_sql_proxy.go

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,8 +130,10 @@ unavailable.`,
130130
// Setting to choose what API to connect to
131131
host = flag.String("host", "",
132132
`When set, the proxy uses this host as the base API path. Example:
133-
https://sqladmin.googleapis.com`,
133+
https://sqladmin.googleapis.com`,
134134
)
135+
quotaProject = flag.String("quota_project", "",
136+
`Specifies the project to use for Cloud SQL Admin API quota tracking.`)
135137

136138
// Settings for healthcheck
137139
useHTTPHealthCheck = flag.Bool("use_http_health_check", false, "When set, creates an HTTP server that checks and communicates the health of the proxy client.")
@@ -374,6 +376,42 @@ func authenticatedClient(ctx context.Context) (*http.Client, oauth2.TokenSource,
374376
return oauth2.NewClient(ctx, src), src, nil
375377
}
376378

379+
// quotaProjectTransport is an http.RoundTripper that adds an X-Goog-User-Project
380+
// header to all requests for quota and billing purposes.
381+
//
382+
// For details, see:
383+
// https://cloud.google.com/apis/docs/system-parameters#definitions
384+
type quotaProjectTransport struct {
385+
base http.RoundTripper
386+
project string
387+
}
388+
389+
var _ http.RoundTripper = quotaProjectTransport{}
390+
391+
// RoundTrip adds a X-Goog-User-Project header to each request.
392+
func (t quotaProjectTransport) RoundTrip(req *http.Request) (*http.Response, error) {
393+
if req.Header == nil {
394+
req.Header = make(http.Header)
395+
}
396+
req.Header.Add("X-Goog-User-Project", t.project)
397+
return t.base.RoundTrip(req)
398+
}
399+
400+
// configureQuotaProject configures an HTTP client to use the provided project
401+
// for quota calculations for all requests.
402+
func configureQuotaProject(c *http.Client, project string) {
403+
// Copy the given client's tripper. Note that tripper can be nil, which is equivalent to
404+
// http.DefaultTransport. (See https://golang.org/pkg/net/http/#Client)
405+
base := c.Transport
406+
if base == nil {
407+
base = http.DefaultTransport
408+
}
409+
c.Transport = quotaProjectTransport{
410+
base: base,
411+
project: project,
412+
}
413+
}
414+
377415
func stringList(s string) []string {
378416
spl := strings.Split(s, ",")
379417
if len(spl) == 1 && spl[0] == "" {
@@ -523,6 +561,11 @@ func runProxy() int {
523561
return 1
524562
}
525563

564+
if *quotaProject != "" {
565+
logging.Infof("Using the project %q for SQL Admin API quota", *quotaProject)
566+
configureQuotaProject(client, *quotaProject)
567+
}
568+
526569
ins, err := listInstances(ctx, client, projList)
527570
if err != nil {
528571
logging.Errorf(err.Error())

0 commit comments

Comments
 (0)