Add support to get secrets from taskcluster (#294)

.isort.cfg

@@ -4,4 +4,4 @@ include_trailing_comma=True
 force_grid_wrap=0
 use_parentheses=True
 line_length=88
-known_third_party = dateutil,flask,gensim,hglib,imblearn,keras,libmozdata,numpy,pandas,parsepatch,pytest,requests,setuptools,shap,sklearn,spacy,tqdm,xgboost,zstandard
+known_third_party = dateutil,flask,gensim,hglib,imblearn,keras,libmozdata,numpy,pandas,parsepatch,pytest,requests,setuptools,shap,sklearn,spacy,taskcluster,tqdm,xgboost,zstandard

bugbug/utils.py

@@ -3,11 +3,16 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this file,
 # You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import os
+
 import numpy as np
+import taskcluster
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.compose import ColumnTransformer
 from sklearn.preprocessing import OrdinalEncoder
 
+TASKCLUSTER_DEFAULT_URL = "https://taskcluster.net"
+
 
 def numpy_to_dict(array):
     return {name: array[name].squeeze(axis=1) for name in array.dtype.names}
@@ -50,3 +55,46 @@ class MissingOrdinalEncoder(OrdinalEncoder):
     def transform(self, X):
         X_int, _ = self._transform(X, handle_unknown="ignore")
         return X_int.astype(self.dtype, copy=False)
+
+
+def get_taskcluster_options():
+    """
+    Helper to get the Taskcluster setup options
+    according to current environment (local or Taskcluster)
+    """
+    options = taskcluster.optionsFromEnvironment()
+    proxy_url = os.environ.get("TASKCLUSTER_PROXY_URL")
+
+    if proxy_url is not None:
+        # Always use proxy url when available
+        options["rootUrl"] = proxy_url
+
+    if "rootUrl" not in options:
+        # Always have a value in root url
+        options["rootUrl"] = TASKCLUSTER_DEFAULT_URL
+
+    return options
+
+
+def get_secret(secret_id):
+    """ Return the secret value
+    """
+    env_variable_name = f"BUGBUG_{secret_id}"
+
+    # Try in the environment first
+    secret_from_env = os.environ.get(env_variable_name)
+
+    if secret_from_env:
+        return secret_from_env
+
+    # If not in env, try with TC if we have the secret id
+    tc_secret_id = os.environ.get("TC_SECRET_ID")
+
+    if tc_secret_id:
+        secrets = taskcluster.Secrets(get_taskcluster_options())
+        secret_bucket = secrets.get(tc_secret_id)
+
+        return secret_bucket["secret"][secret_id]
+
+    else:
+        raise ValueError("Failed to find secret {}".format(secret_id))

requirements.txt

@@ -17,3 +17,4 @@ keras==2.2.4
 tqdm==4.31.1
 python-dateutil==2.8.0
 zstandard==0.11.0
+taskcluster==7.0.1

scripts/bug_retriever.py

@@ -2,7 +2,6 @@
 
 import argparse
 import lzma
-import os
 import shutil
 from datetime import datetime
 from logging import INFO, basicConfig, getLogger
@@ -10,21 +9,12 @@ from logging import INFO, basicConfig, getLogger
 from dateutil.relativedelta import relativedelta
 
 from bugbug import bug_snapshot, bugzilla, labels
+from bugbug.utils import get_secret
 
 basicConfig(level=INFO)
 logger = getLogger(__name__)
 
 
-def get_secret(secret_id):
-    """ Return the secret value
-
-    TODO: Support task-cluster secret API
-    """
-    env_variable_name = f"BUGBUG_{secret_id}"
-
-    return os.environ[env_variable_name]
-
-
 class Retriever(object):
     def retrieve_bugs(self):
         bugzilla.set_token(get_secret("BUGZILLA_TOKEN"))