Merge pull request #13 from pangeo-gallery/pass-secrets

Pass env vars
pangeo-gallery · Apr 17, 2020 · 0a915da · 0a915da
2 parents ae12964 + 09bd06c
commit 0a915da
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 48 deletions.
diff --git a/binderbot/binderbot.py b/binderbot/binderbot.py
@@ -16,12 +16,17 @@
 import time
 import json
 import textwrap
+import re
 
 import nbformat
 from nbconvert.preprocessors import ClearOutputPreprocessor
 
 logger = structlog.get_logger()
 
+# https://stackoverflow.com/questions/14693701/how-can-i-remove-the-ansi-escape-sequences-from-a-string-in-python
+def _ansi_escape(text):
+    return re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])').sub('', text)
+
 
 class OperationError(Exception):
     pass
@@ -35,7 +40,7 @@ class States(Enum):
         KERNEL_STARTED = 4
 
     async def __aenter__(self):
-        self.session = aiohttp.ClientSession()
+        self.session = aiohttp.ClientSession(headers={'User-Agent': 'BinderBot-cli v0.1'})
         return self
 
     async def __aexit__(self, exc_type, exc, tb):
@@ -225,8 +230,10 @@ async def run_code(self, code):
                             #if msg_type == 'execute_result':
                             #    response = msg['content']['data']['text/plain']
                             if msg_type == 'error':
-                                traceback = msg['content']['traceback']
-                                self.log.msg('Code Execute: Error', action='code-execute', phase='error', traceback=traceback)
+                                traceback = _ansi_escape('\n'.join(msg['content']['traceback']))
+                                self.log.msg('Code Execute: Error', action='code-execute',
+                                             phase='error',
+                                             traceback=traceback)
                                 raise OperationError()
                             elif msg_type == 'stream':
                                 response = msg['content']['text']
@@ -262,10 +269,14 @@ async def list_notebooks(self):
         stdout, stderr = await self.run_code(code)
         return json.loads(stdout)
 
-    async def execute_notebook(self, notebook_filename, timeout=600):
+    async def execute_notebook(self, notebook_filename, timeout=600,
+                               env_vars={}):
+        env_var_str = str(env_vars)
         # https://nbconvert.readthedocs.io/en/latest/execute_api.html
         code = f"""
+        import os
         import nbformat
+        os.environ.update({env_var_str})
         from nbconvert.preprocessors import ExecutePreprocessor
         ep = ExecutePreprocessor(timeout={timeout})
         print("Processing {notebook_filename}")

diff --git a/binderbot/cli.py b/binderbot/cli.py
@@ -30,10 +30,12 @@ def wrapper(*args, **kwargs):
               help="Maximum execution time (in second) for each notebook.")
 @click.option("--binder-start-timeout", default=600,
               help="Maximum time (in seconds) to wait for binder to start.")
+@click.option("--pass-env-var", "-e", multiple=True,
+              help="Environment variables to pass to the binder execution environment.")
 @click.argument('filenames', nargs=-1, type=click.Path(exists=True))
 @coro
 async def main(binder_url, repo, ref, output_dir, nb_timeout,
-               binder_start_timeout, filenames):
+               binder_start_timeout, pass_env_var, filenames):
     """Run local notebooks on a remote binder."""
 
     # validate filename inputs
@@ -49,6 +51,8 @@ async def main(binder_url, repo, ref, output_dir, nb_timeout,
                f"     repo: {repo}\n"
                f"     ref: {ref}")
 
+    extra_env_vars = {k: os.environ[k] for k in pass_env_var}
+
     # inputs look good, start up binder
     async with BinderUser(binder_url, repo, ref) as jovyan:
         await jovyan.start_binder(timeout=binder_start_timeout)
@@ -63,7 +67,8 @@ async def main(binder_url, repo, ref, output_dir, nb_timeout,
                 await jovyan.upload_local_notebook(fname)
                 click.echo("✅")
                 click.echo(f"⌛️ Executing {fname}...", nl=False)
-                await jovyan.execute_notebook(fname, timeout=nb_timeout)
+                await jovyan.execute_notebook(fname, timeout=nb_timeout,
+                                              env_vars=extra_env_vars)
                 click.echo("✅")
                 click.echo(f"⌛️ Downloading and saving {fname}...", nl=False)
                 nb_data = await jovyan.get_contents(fname)

diff --git a/tests/test_binderbot.py b/tests/test_binderbot.py
@@ -14,45 +14,28 @@
 
 @pytest.fixture()
 def example_nb_data():
-    nbdata = {'cells': [{'cell_type': 'markdown',
-               'metadata': {},
-               'source': '# Test Pangeo Gallery'},
-              {'cell_type': 'code',
-               'execution_count': 1,
-               'metadata': {},
-               'outputs': [{'name': 'stdout',
-                 'output_type': 'stream',
-                 'text': "Today's date: 2020-03-12\n"}],
-               'source': 'from datetime import date\ntoday = date.today()\nprint("Today\'s date:", today)'},
-              {'cell_type': 'code',
-               'execution_count': 2,
-               'metadata': {},
-               'outputs': [{'name': 'stdout',
-                 'output_type': 'stream',
-                 'text': 'Ryans-MacBook-Pro.local\n'}],
-               'source': 'import socket\nprint(socket.gethostname())'},
-              {'cell_type': 'code',
-               'execution_count': 0,
-               'metadata': {},
-               'outputs': [],
-               'source': 'import time\ntime.sleep(2)'},
-              {'cell_type': 'code',
-               'execution_count': None,
-               'metadata': {},
-               'outputs': [],
-               'source': ''}],
-             'metadata': {'kernelspec': {'display_name': 'Python 3',
-               'language': 'python',
-               'name': 'python3'},
-              'language_info': {'codemirror_mode': {'name': 'ipython', 'version': 3},
-               'file_extension': '.py',
-               'mimetype': 'text/x-python',
-               'name': 'python',
-               'nbconvert_exporter': 'python',
-               'pygments_lexer': 'ipython3',
-               'version': '3.6.7'}},
-             'nbformat': 4,
-             'nbformat_minor': 4}
+    nbdata = {'cells': [{'cell_type': 'code',
+                       'execution_count': None,
+                       'metadata': {},
+                       'outputs': [],
+                       'source': 'import socket\nprint(socket.gethostname())'},
+                      {'cell_type': 'code',
+                       'execution_count': None,
+                       'metadata': {},
+                       'outputs': [],
+                       'source': 'import os\nprint(os.environ["MY_VAR"])'}],
+                     'metadata': {'kernelspec': {'display_name': 'Python 3',
+                       'language': 'python',
+                       'name': 'python3'},
+                      'language_info': {'codemirror_mode': {'name': 'ipython', 'version': 3},
+                       'file_extension': '.py',
+                       'mimetype': 'text/x-python',
+                       'name': 'python',
+                       'nbconvert_exporter': 'python',
+                       'pygments_lexer': 'ipython3',
+                       'version': '3.8.1'}},
+                     'nbformat': 4,
+                     'nbformat_minor': 4}
     return nbformat.from_dict(nbdata)
 
     return tmpdir
@@ -66,20 +49,23 @@ def test_cli_upload_execute_download(tmp_path, example_nb_data):
     with open(fname, 'w', encoding='utf-8') as f:
         nbformat.write(example_nb_data, f)
 
-    runner = CliRunner()
+    env = {"MY_VAR": "SECRET"}
+    runner = CliRunner(env=env)
     args = ["--binder-url", "http://mybinder.org",
             "--repo", "binder-examples/requirements",
             "--ref", "master", "--nb-timeout", "10",
+            "--pass-env-var",  "MY_VAR",
             fname]
     result = runner.invoke(cli.main, args)
     assert result.exit_code == 0, result.output
-    # assert 'binderbot.cli.main' in result.output
 
     with open(fname) as f:
         nb = nbformat.read(f, as_version=4)
 
-    hostname = nb['cells'][2]['outputs'][0]['text']
+    hostname = nb['cells'][0]['outputs'][0]['text']
     assert hostname.startswith('jupyter-binder-')
+    remote_env_var_value = nb['cells'][1]['outputs'][0]['text']
+    assert remote_env_var_value.rstrip() == env['MY_VAR']
 
     # help_result = runner.invoke(cli.main, ['--help'])
     # assert help_result.exit_code == 0