add support for sts regional interface endpoints

neptune-python-utils/neptune_python_utils/endpoints.py

@@ -60,8 +60,8 @@ def __init__(self, uri, querystring, headers):
         self.headers = headers
 
 class Endpoint:
-    
-    def __init__(self, protocol, neptune_endpoint, neptune_port, suffix, region, credentials=None, role_arn=None, proxy_dns=None, proxy_port=8182, remove_host_header=False): 
+
+    def __init__(self, protocol, neptune_endpoint, neptune_port, suffix, region, credentials=None, role_arn=None, proxy_dns=None, proxy_port=8182, remove_host_header=False, endpoint_url=None):
 
         self.protocol = protocol
         self.neptune_endpoint = neptune_endpoint
@@ -71,7 +71,8 @@ def __init__(self, protocol, neptune_endpoint, neptune_port, suffix, region, cre
         self.proxy_dns = proxy_dns
         self.proxy_port = proxy_port
         self.remove_host_header = remove_host_header
-
+        self.endpoint_url = endpoint_url
+
         if role_arn:
             self.role_arn = role_arn
             self.credentials = None
@@ -93,7 +94,10 @@ def _get_session_credentials(self):
     def _get_credentials(self):        
 
         if self.credentials is None:
-            sts = boto3.client('sts', region_name=self.region)
+            if self.endpoint_url:
+                sts = boto3.client('sts', region_name=self.region, endpoint_url=self.endpoint_url)
+            else:
+                sts = boto3.client('sts', region_name=self.region)
 
             role = sts.assume_role(
                 RoleArn=self.role_arn,
@@ -110,7 +114,10 @@ def _get_credentials(self):
         return self.credentials.get_frozen_credentials()
 
     def _new_credentials(self):
-        sts = boto3.client('sts', region_name=self.region)
+        if self.endpoint_url:
+            sts = boto3.client('sts', region_name=self.region, endpoint_url=self.endpoint_url)
+        else:
+            sts = boto3.client('sts', region_name=self.region)
 
         role = sts.assume_role(
             RoleArn=self.role_arn,
@@ -169,8 +176,8 @@ def get_headers():
 
 
 class Endpoints:
-    
-    def __init__(self, neptune_endpoint=None, neptune_port=None, region_name=None, credentials=None, role_arn=None, proxy_dns=None, proxy_port=8182, remove_host_header=False):
+
+    def __init__(self, neptune_endpoint=None, neptune_port=None, region_name=None, credentials=None, role_arn=None, proxy_dns=None, proxy_port=8182, remove_host_header=False, endpoint_url=None):
 
         if neptune_endpoint is None:
             assert ('NEPTUNE_CLUSTER_ENDPOINT' in os.environ), 'neptune_endpoint is missing.'
@@ -188,12 +195,13 @@ def __init__(self, neptune_endpoint=None, neptune_port=None, region_name=None, c
         else:
             session = boto3.session.Session()
             self.region = session.region_name
-            
+
         self.credentials = credentials
         self.role_arn = role_arn
         self.proxy_dns = proxy_dns
         self.proxy_port = proxy_port
         self.remove_host_header = remove_host_header
+        self.endpoint_url = endpoint_url
 
 
     def gremlin_endpoint(self):
@@ -218,5 +226,4 @@ def sparql_stream_endpoint(self):
         return self.__endpoint('https', self.neptune_endpoint, self.neptune_port, 'sparql/stream')
 
     def __endpoint(self, protocol, neptune_endpoint, neptune_port, suffix):
-        return Endpoint(protocol, neptune_endpoint, neptune_port, suffix, self.region, self.credentials, self.role_arn, self.proxy_dns, self.proxy_port, self.remove_host_header)
-
+        return Endpoint(protocol, neptune_endpoint, neptune_port, suffix, self.region, self.credentials, self.role_arn, self.proxy_dns, self.proxy_port, self.remove_host_header, self.endpoint_url)

neptune-python-utils/readme.md

@@ -80,6 +80,16 @@ from neptune_python_utils.endpoints import Endpoints
 endpoints = Endpoints(role_arn='arn:aws:iam::...')
 ```
 
+If your Amazon VPC configuration doesn't have a public subnet, AWS services' endpoints like STS can be accessed via VPC endpoints (InterfaceEndpoints or GatewayEndpoints). AWS STS has regional endpoints, listed here - [Using AWS STS interface VPC endpoints](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_sts_vpce.html)
+
+To enable the library to successfully connect to STS, pass STS regional endpoint value in `endpoint_url` parameter.
+
+```
+from neptune_python_utils.endpoints import Endpoints
+
+endpoints = Endpoints(**other_kwargs, endpoint_url='https://sts.eu-west-1.amazonaws.com')
+```
+
 #### Proxies
 
 If you want to connect to Neptune via a proxy – a bastion host, [application load balancer or network load balancer](https://github.com/aws-samples/aws-dbs-refarch-graph/tree/master/src/connecting-using-a-load-balancer) – you must supply the proxy DNS and port to an `Endpoints` instance: