feat: send quota project id in x-goog-user-project for OAuth2 credent…

…ials (#412)

* 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.

* feat: add `__setstate__` and `__getstate__` to `oauth2.credentials` class

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:

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:

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,27 @@ def __init__(
         self._token_uri = token_uri
         self._client_id = client_id
         self._client_secret = client_secret
+        self._quota_project_id = quota_project_id
+
+    def __getstate__(self):
+        """A __getstate__ method must exist for the __setstate__ to be called
+        This is identical to the default implementation.
+        See https://docs.python.org/3.7/library/pickle.html#object.__setstate__
+        """
+        return self.__dict__
+
+    def __setstate__(self, d):
+        """Credentials pickled with older versions of the class do not have
+        all the attributes."""
+        self.token = d.get("token")
+        self.expiry = d.get("expiry")
+        self._refresh_token = d.get("_refresh_token")
+        self._id_token = d.get("_id_token")
+        self._scopes = d.get("_scopes")
+        self._token_uri = d.get("_token_uri")
+        self._client_id = d.get("_client_id")
+        self._client_secret = d.get("_client_secret")
+        self._quota_project_id = d.get("_quota_project_id")
 
     @property
     def refresh_token(self):
@@ -123,6 +148,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 +199,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 +238,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

tests/oauth2/test_credentials.py

@@ -15,6 +15,8 @@
 import datetime
 import json
 import os
+import pickle
+import sys
 
 import mock
 import pytest
@@ -294,6 +296,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()
 
@@ -355,3 +384,40 @@ def test_to_json(self):
         assert json_asdict.get("client_id") == creds.client_id
         assert json_asdict.get("scopes") == creds.scopes
         assert json_asdict.get("client_secret") is None
+
+    def test_pickle_and_unpickle(self):
+        creds = self.make_credentials()
+        unpickled = pickle.loads(pickle.dumps(creds))
+
+        # make sure attributes aren't lost during pickling
+        assert list(creds.__dict__).sort() == list(unpickled.__dict__).sort()
+
+        for attr in list(creds.__dict__):
+            assert getattr(creds, attr) == getattr(unpickled, attr)
+
+    def test_pickle_with_missing_attribute(self):
+        creds = self.make_credentials()
+
+        # remove an optional attribute before pickling
+        # this mimics a pickle created with a previous class definition with
+        # fewer attributes
+        del creds.__dict__["_quota_project_id"]
+
+        unpickled = pickle.loads(pickle.dumps(creds))
+
+        # Attribute should be initialized by `__setstate__`
+        assert unpickled.quota_project_id is None
+
+    # pickles are not compatible across versions
+    @pytest.mark.skipif(
+        sys.version_info < (3, 5),
+        reason="pickle file can only be loaded with Python >= 3.5",
+    )
+    def test_unpickle_old_credentials_pickle(self):
+        # make sure a credentials file pickled with an older
+        # library version (google-auth==1.5.1) can be unpickled
+        with open(
+            os.path.join(DATA_DIR, "old_oauth_credentials_py3.pickle"), "rb"
+        ) as f:
+            credentials = pickle.load(f)
+            assert credentials.quota_project_id is None