Use updated authentication method for IAM Bolt queries (#438)

* Use updated authentication method for IAM Bolt queries

* update changelog

Co-authored-by: Michael Chin <chnmch@amazon.com>

Author: michaelnchin

ChangeLog.md

@@ -9,6 +9,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 - Added support for list/tuple element access in cell variable injection ([Link to PR](https://github.com/aws/graph-notebook/pull/409))
 - Fixed failing endpoint creation step in [01-People-Analytics/People-Analytics-using-Neptune-ML](https://github.com/aws/graph-notebook/blob/main/src/graph_notebook/notebooks/04-Machine-Learning/Sample-Applications/01-People-Analytics/People-Analytics-using-Neptune-ML.ipynb) ([Link to PR](https://github.com/aws/graph-notebook/pull/411))
 - Fixed browser-specific issues with fullscreen graph widget ([Link to PR](https://github.com/aws/graph-notebook/pull/427))
+- Fixed IAM authenticated Bolt queries failing on certain Neptune engine versions ([Link to PR](https://github.com/aws/graph-notebook/pull/438))
 - Pinned `numpy<1.24.0` to fix conflicts with `networkx` dependency during installation ([Link to PR](https://github.com/aws/graph-notebook/pull/416))
 - Excluded certain `itables` versions causing errors in query magics ([PR #1](https://github.com/aws/graph-notebook/pull/429)) ([PR #2](https://github.com/aws/graph-notebook/pull/429))
 - Pinned version ceiling for all dependencies ([Link to PR](https://github.com/aws/graph-notebook/pull/431))

src/graph_notebook/neptune/bolt_auth_token.py

@@ -0,0 +1,43 @@
+import json
+
+from neo4j import Auth
+from botocore.awsrequest import AWSRequest
+from botocore.credentials import Credentials
+from botocore.compat import urlsplit
+from botocore.auth import SigV4Auth
+
+SCHEME = "basic"
+REALM = "realm"
+SERVICE_NAME = "neptune-db"
+DUMMY_USERNAME = "username"
+HTTP_METHOD_HDR = "HttpMethod"
+HTTP_METHOD = "GET"
+AUTHORIZATION = "Authorization"
+X_AMZ_DATE = "X-Amz-Date"
+X_AMZ_SECURITY_TOKEN = "X-Amz-Security-Token"
+HOST = "Host"
+
+
+class NeptuneBoltAuthToken(Auth):
+    def __init__(
+        self,
+        credentials: Credentials,
+        region: str,
+        url: str,
+        **parameters
+    ):
+        request = AWSRequest(method=HTTP_METHOD, url=url)
+
+        url_parts = urlsplit(request.url)
+        host_part = url_parts.hostname
+        request.headers.add_header("Host", host_part)
+        sigv4 = SigV4Auth(credentials, SERVICE_NAME, region)
+        sigv4.add_auth(request)
+
+        auth_obj = {
+          hdr: request.headers[hdr]
+          for hdr in [AUTHORIZATION, X_AMZ_DATE, X_AMZ_SECURITY_TOKEN, HOST]
+        }
+        auth_obj[HTTP_METHOD_HDR] = request.method
+        creds: str = json.dumps(auth_obj)
+        super().__init__(SCHEME, DUMMY_USERNAME, creds, REALM, **parameters)

src/graph_notebook/neptune/client.py

@@ -21,6 +21,8 @@
 from base64 import b64encode
 import nest_asyncio
 
+from graph_notebook.neptune.bolt_auth_token import NeptuneBoltAuthToken
+
 
 # This patch is no longer needed when graph_notebook is using the a Gremlin Python
 # client >= 3.5.0 as the HashableDict is now part of that client driver.
@@ -391,22 +393,23 @@ def get_opencypher_driver(self):
         url = f'bolt://{self.host}:{self.port}'
 
         if self.is_neptune_domain():
-            user = DEFAULT_NEO4J_USERNAME
             if self._session:
-                method = 'POST'
-                headers = {
-                    'HttpMethod': method,
-                    'Host': url
-                }
-                aws_request = self._get_aws_request('POST', url)
-                for item in aws_request.headers.items():
-                    headers[item[0]] = item[1]
-
-                auth_str = json.dumps(headers)
-                password = auth_str
+                # check engine version via status API to determine if we need the OC endpoint path
+                status_res = self.status()
+                status_res.raise_for_status()
+                status_res_json = status_res.json()
+                engine_version_raw = status_res_json["dbEngineVersion"]
+                engine_version = int(engine_version_raw.rsplit('.', 1)[0].replace('.', ''))
+                if engine_version >= 1200:
+                    url += "/opencypher"
+
+                credentials = self._session.get_credentials()
+                frozen_creds = credentials.get_frozen_credentials()
+                auth_final = NeptuneBoltAuthToken(frozen_creds, self.region, url)
             else:
+                user = 'username'
                 password = DEFAULT_NEO4J_PASSWORD
-            auth_final = (user, password)
+                auth_final = (user, password)
         else:
             if self.neo4j_auth:
                 auth_final = (self.neo4j_username, self.neo4j_password)