[BUG] Fix issue causing export magic commands to use neptune's auth s…

…etting

when calling the exporter, pick up the iam setting of the export magic and build a new client off it to ensure we sign requests when needed, no matter the auth setting on the neptune cluster.

ChangeLog.md

@@ -4,6 +4,10 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 
 ## Upcoming
 
+## Release 2.1.1 (April 22, 2021)
+
+- Fix bug in `%neptune_ml export ...` logic where the iam setting for the exporter endpoint wasn't getting picked up properly
+
 ## Release 2.1.0 (April 15, 2021)
 
 - Add support for Mode, queueRequest, and Dependencies parameters when running %load command ([Link to PR](https://github.com/aws/graph-notebook/pull/91))

src/graph_notebook/__init__.py

@@ -3,4 +3,4 @@
 SPDX-License-Identifier: Apache-2.0
 """
 
-__version__ = '2.1.0'
+__version__ = '2.1.1'

src/graph_notebook/magics/ml.py

@@ -4,10 +4,11 @@
 import logging
 import time
 from IPython.core.display import display
+from botocore.session import get_session
 from ipywidgets import widgets
 
 from graph_notebook.magics.parsing import str_to_namespace_var
-from graph_notebook.neptune.client import Client
+from graph_notebook.neptune.client import Client, ClientBuilder
 
 logger = logging.getLogger("neptune_ml_magic_handler")
 
@@ -189,22 +190,34 @@ def wait_for_export(client: Client, export_url: str, job_id: str, output: widget
 
 def neptune_ml_export(args: argparse.Namespace, client: Client, output: widgets.Output,
                       cell: str):
+    # since the exporter is a different host than Neptune, it's IAM auth setting can be different
+    # than the client we're using to connect to Neptune. Because of this, need o recreate out client before calling
+    # any exporter urls to ensure we're signing any requests that we need to.
+    builder = ClientBuilder().with_host(client.host)\
+                            .with_port(client.port)\
+                            .with_region(client.region)\
+                            .with_tls(client.ssl)
+
+    if args.export_iam:
+        builder = builder.with_iam(get_session())
+    export_client = builder.build()
+
     export_ssl = not args.export_no_ssl
     if args.which_sub == 'start':
         if cell == '':
             return 'Cell body must have json payload or reference notebook variable using syntax ${payload_var}'
-        export_job = neptune_ml_export_start(client, cell, args.export_url, export_ssl)
+        export_job = neptune_ml_export_start(export_client, cell, args.export_url, export_ssl)
         if args.wait:
-            return wait_for_export(client, args.export_url, export_job['jobId'],
+            return wait_for_export(export_client, args.export_url, export_job['jobId'],
                                    output, export_ssl, args.wait_interval, args.wait_timeout)
         else:
             return export_job
     elif args.which_sub == 'status':
         if args.wait:
-            status = wait_for_export(client, args.export_url, args.job_id, output, export_ssl, args.wait_interval,
+            status = wait_for_export(export_client, args.export_url, args.job_id, output, export_ssl, args.wait_interval,
                                      args.wait_timeout)
         else:
-            status_res = client.export_status(args.export_url, args.job_id, export_ssl)
+            status_res = export_client.export_status(args.export_url, args.job_id, export_ssl)
             status_res.raise_for_status()
             status = status_res.json()
         return status

src/graph_notebook/neptune/client.py

@@ -67,13 +67,12 @@
 
 class Client(object):
     def __init__(self, host: str, port: int = DEFAULT_PORT, ssl: bool = True, region: str = DEFAULT_REGION,
-                 sparql_path: str = '/sparql', auth=None, session: Session = None):
+                 sparql_path: str = '/sparql', session: Session = None):
         self.host = host
         self.port = port
         self.ssl = ssl
         self.sparql_path = sparql_path
         self.region = region
-        self._auth = auth
         self._session = session
 
         self._http_protocol = 'https' if self.ssl else 'http'
@@ -434,7 +433,7 @@ def _query_status(self, language: str, *, query_id: str = '', **kwargs) -> reque
 
     def _prepare_request(self, method, url, *, data=None, params=None, headers=None, service=NEPTUNE_SERVICE_NAME):
         self._ensure_http_session()
-        request = requests.Request(method=method, url=url, data=data, params=params, headers=headers, auth=self._auth)
+        request = requests.Request(method=method, url=url, data=data, params=params, headers=headers)
         if self._session is not None:
             credentials = self._session.get_credentials()
             frozen_creds = credentials.get_frozen_credentials()

src/graph_notebook/widgets/package.json

@@ -1,6 +1,6 @@
 {
   "name": "graph_notebook_widgets",
-  "version": "2.1.0",
+  "version": "2.1.1",
   "author": "amazon",
   "description": "A Custom Jupyter Library for rendering NetworkX MultiDiGraphs using vis-network",
   "dependencies": {