From ab3dc1e26f5240ea3456de364c7c5cb8f40f9583 Mon Sep 17 00:00:00 2001
From: Bu Sun Kim <8822365+busunkim96@users.noreply.github.com>
Date: Fri, 6 Dec 2019 15:56:24 -0500
Subject: [PATCH] feat: send quota project id in x-goog-user-project header for
 OAuth2 credentials (#400)

When the 3LO credentials are used, the quota project ("quota_project_id") is sent on every outgoing request in the x-goog-user-project HTTP header/grpc metadata. The quota project is used for billing and quota purposes.
---
 docs/reference/google.auth.crypt.base.rst |  7 ++++++
 docs/reference/google.auth.crypt.rsa.rst  |  7 ++++++
 google/oauth2/credentials.py              | 19 ++++++++++++++++
 tests/oauth2/test_credentials.py          | 27 +++++++++++++++++++++++
 4 files changed, 60 insertions(+)
 create mode 100644 docs/reference/google.auth.crypt.base.rst
 create mode 100644 docs/reference/google.auth.crypt.rsa.rst

diff --git a/docs/reference/google.auth.crypt.base.rst b/docs/reference/google.auth.crypt.base.rst
new file mode 100644
index 000000000..a8996501a
--- /dev/null
+++ b/docs/reference/google.auth.crypt.base.rst
@@ -0,0 +1,7 @@
+google.auth.crypt.base module
+=============================
+
+.. automodule:: google.auth.crypt.base
+   :members:
+   :inherited-members:
+   :show-inheritance:
diff --git a/docs/reference/google.auth.crypt.rsa.rst b/docs/reference/google.auth.crypt.rsa.rst
new file mode 100644
index 000000000..7060b03c8
--- /dev/null
+++ b/docs/reference/google.auth.crypt.rsa.rst
@@ -0,0 +1,7 @@
+google.auth.crypt.rsa module
+============================
+
+.. automodule:: google.auth.crypt.rsa
+   :members:
+   :inherited-members:
+   :show-inheritance:
diff --git a/google/oauth2/credentials.py b/google/oauth2/credentials.py
index 3a32c0631..422c8ab10 100644
--- a/google/oauth2/credentials.py
+++ b/google/oauth2/credentials.py
@@ -58,6 +58,7 @@ def __init__(
         client_id=None,
         client_secret=None,
         scopes=None,
+        quota_project_id=None,
     ):
         """
         Args:
@@ -81,6 +82,9 @@ def __init__(
                 token if refresh information is provided (e.g. The refresh
                 token scopes are a superset of this or contain a wild card
                 scope like 'https://www.googleapis.com/auth/any-api').
+            quota_project_id (Optional[str]): The project ID used for quota and billing.
+                This project may be different from the project used to
+                create the credentials.
         """
         super(Credentials, self).__init__()
         self.token = token
@@ -90,6 +94,7 @@ def __init__(
         self._token_uri = token_uri
         self._client_id = client_id
         self._client_secret = client_secret
+        self._quota_project_id = quota_project_id
 
     @property
     def refresh_token(self):
@@ -123,6 +128,11 @@ def client_secret(self):
         """Optional[str]: The OAuth 2.0 client secret."""
         return self._client_secret
 
+    @property
+    def quota_project_id(self):
+        """Optional[str]: The project to use for quota and billing purposes."""
+        return self._quota_project_id
+
     @property
     def requires_scopes(self):
         """False: OAuth 2.0 credentials have their scopes set when
@@ -169,6 +179,12 @@ def refresh(self, request):
                     )
                 )
 
+    @_helpers.copy_docstring(credentials.Credentials)
+    def apply(self, headers, token=None):
+        super(Credentials, self).apply(headers, token=token)
+        if self.quota_project_id is not None:
+            headers["x-goog-user-project"] = self.quota_project_id
+
     @classmethod
     def from_authorized_user_info(cls, info, scopes=None):
         """Creates a Credentials instance from parsed authorized user info.
@@ -202,6 +218,9 @@ def from_authorized_user_info(cls, info, scopes=None):
             scopes=scopes,
             client_id=info["client_id"],
             client_secret=info["client_secret"],
+            quota_project_id=info.get(
+                "quota_project_id"
+            ),  # quota project may not exist
         )
 
     @classmethod
diff --git a/tests/oauth2/test_credentials.py b/tests/oauth2/test_credentials.py
index bb70f1516..59031d7de 100644
--- a/tests/oauth2/test_credentials.py
+++ b/tests/oauth2/test_credentials.py
@@ -294,6 +294,33 @@ def test_credentials_with_scopes_refresh_failure_raises_refresh_error(
         # expired.)
         assert creds.valid
 
+    def test_apply_with_quota_project_id(self):
+        creds = credentials.Credentials(
+            token="token",
+            refresh_token=self.REFRESH_TOKEN,
+            token_uri=self.TOKEN_URI,
+            client_id=self.CLIENT_ID,
+            client_secret=self.CLIENT_SECRET,
+            quota_project_id="quota-project-123",
+        )
+
+        headers = {}
+        creds.apply(headers)
+        assert headers["x-goog-user-project"] == "quota-project-123"
+
+    def test_apply_with_no_quota_project_id(self):
+        creds = credentials.Credentials(
+            token="token",
+            refresh_token=self.REFRESH_TOKEN,
+            token_uri=self.TOKEN_URI,
+            client_id=self.CLIENT_ID,
+            client_secret=self.CLIENT_SECRET,
+        )
+
+        headers = {}
+        creds.apply(headers)
+        assert "x-goog-user-project" not in headers
+
     def test_from_authorized_user_info(self):
         info = AUTH_USER_INFO.copy()