SSH tunnel support for S3Store

src/maggma/stores/aws.py

@@ -13,6 +13,7 @@
 from monty.msgpack import default as monty_default
 
 from maggma.core import Sort, Store
+from maggma.stores.ssh_tunnel import SSHTunnel
 from maggma.utils import grouper, to_isoformat_ceil_ms
 
 try:
@@ -40,6 +41,7 @@
         sub_dir: Optional[str] = None,
         s3_workers: int = 1,
         s3_resource_kwargs: Optional[dict] = None,
+        ssh_tunnel: Optional[SSHTunnel] = None,
         key: str = "fs_id",
         store_hash: bool = True,
         unpack_data: bool = True,
@@ -59,11 +61,16 @@
                     aws_session_token (string) -- AWS temporary session token
                     region_name (string) -- Default region when creating new connections
             compress: compress files inserted into the store
-            endpoint_url: endpoint_url to allow interface to minio service
-            sub_dir: (optional)  subdirectory of the s3 bucket to store the data
+            endpoint_url: endpoint_url to allow interface to minio service; ignored if
+                `ssh_tunnel` is provided, in which case the endpoint_url is inferred
+            sub_dir: (optional) subdirectory of the s3 bucket to store the data
             s3_workers: number of concurrent S3 puts to run
+            s3_resource_kwargs: additional kwargs to pass to the boto3 session resource
+            ssh_tunnel: optional SSH tunnel to use for the S3 connection
+            key: main key to index on
             store_hash: store the sha1 hash right before insertion to the database.
-            unpack_data: whether to decompress and unpack byte data when querying from the bucket.
+            unpack_data: whether to decompress and unpack byte data when querying from
+                the bucket
             searchable_fields: fields to keep in the index store
         """
         if boto3 is None:
@@ -79,6 +86,7 @@
         self.s3_bucket: Any = None
         self.s3_workers = s3_workers
         self.s3_resource_kwargs = s3_resource_kwargs if s3_resource_kwargs is not None else {}
+        self.ssh_tunnel = ssh_tunnel
         self.unpack_data = unpack_data
         self.searchable_fields = searchable_fields if searchable_fields is not None else []
         self.store_hash = store_hash
@@ -107,7 +115,8 @@
     def connect(self, *args, **kwargs):  # lgtm[py/conflicting-attributes]
         """Connect to the source data."""
         session = self._get_session()
-        resource = session.resource("s3", endpoint_url=self.endpoint_url, **self.s3_resource_kwargs)
+        endpoint_url = self._get_endpoint_url()
+        resource = session.resource("s3", endpoint_url=endpoint_url, **self.s3_resource_kwargs)
 
         if not self.s3:
             self.s3 = resource
@@ -127,6 +136,9 @@
         self.s3 = None
         self.s3_bucket = None
 
+        if self.ssh_tunnel is not None:
+            self.ssh_tunnel.stop()
+
     @property
     def _collection(self):
         """
@@ -266,7 +278,7 @@
 
         Args:
             key: single key to index
-            unique: Whether or not this index contains only unique keys
+            unique: whether or not this index contains only unique keys
 
         Returns:
             bool indicating if the index exists/was created
@@ -322,19 +334,30 @@
         self.index.update(search_docs, key=self.key)
 
     def _get_session(self):
+        if self.ssh_tunnel is not None:
+            self.ssh_tunnel.start()
+
         if not hasattr(self._thread_local, "s3_bucket"):
             if isinstance(self.s3_profile, dict):
                 return Session(**self.s3_profile)
             return Session(profile_name=self.s3_profile)
         return None
 
+    def _get_endpoint_url(self):
+        if self.ssh_tunnel is None:
+            return self.endpoint_url
+        else:
+            host, port = self.ssh_tunnel.local_address
+            return f"http://{host}:{port}"
+
     def _get_bucket(self):
         """If on the main thread return the bucket created above, else create a new bucket on each thread."""
         if threading.current_thread().name == "MainThread":
             return self.s3_bucket
         if not hasattr(self._thread_local, "s3_bucket"):
             session = self._get_session()
-            resource = session.resource("s3", endpoint_url=self.endpoint_url)
+            endpoint_url = self._get_endpoint_url()
+            resource = session.resource("s3", endpoint_url=endpoint_url)
             self._thread_local.s3_bucket = resource.Bucket(self.bucket)
         return self._thread_local.s3_bucket
 

src/maggma/stores/gridfs.py

@@ -18,7 +18,8 @@
 from ruamel import yaml
 
 from maggma.core import Sort, Store, StoreError
-from maggma.stores.mongolike import MongoStore, SSHTunnel
+from maggma.stores.mongolike import MongoStore
+from maggma.stores.ssh_tunnel import SSHTunnel
 
 # https://github.com/mongodb/specifications/
 #   blob/master/source/gridfs/gridfs-spec.rst#terms

src/maggma/stores/mongolike.py

@@ -7,10 +7,11 @@
 import warnings
 from itertools import chain, groupby
 from pathlib import Path
-from socket import socket
 
 from ruamel import yaml
 
+from maggma.stores.ssh_tunnel import SSHTunnel
+
 try:
     from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Tuple, Union
 except ImportError:
@@ -23,12 +24,11 @@
 import orjson
 from monty.dev import requires
 from monty.io import zopen
-from monty.json import MSONable, jsanitize
+from monty.json import jsanitize
 from monty.serialization import loadfn
 from pydash import get, has, set_
 from pymongo import MongoClient, ReplaceOne, uri_parser
 from pymongo.errors import ConfigurationError, DocumentTooLarge, OperationFailure
-from sshtunnel import SSHTunnelForwarder
 
 from maggma.core import Sort, Store, StoreError
 from maggma.utils import confirm_field_index, to_dt
@@ -39,79 +39,6 @@
     MontyClient = None
 
 
-class SSHTunnel(MSONable):
-    __TUNNELS: Dict[str, SSHTunnelForwarder] = {}
-
-    def __init__(
-        self,
-        tunnel_server_address: str,
-        remote_server_address: str,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
-        private_key: Optional[str] = None,
-        **kwargs,
-    ):
-        """
-        Args:
-            tunnel_server_address: string address with port for the SSH tunnel server
-            remote_server_address: string address with port for the server to connect to
-            username: optional username for the ssh tunnel server
-            password: optional password for the ssh tunnel server; If a private_key is
-                supplied this password is assumed to be the private key password
-            private_key: ssh private key to authenticate to the tunnel server
-            kwargs: any extra args passed to the SSHTunnelForwarder
-        """
-
-        self.tunnel_server_address = tunnel_server_address
-        self.remote_server_address = remote_server_address
-        self.username = username
-        self.password = password
-        self.private_key = private_key
-        self.kwargs = kwargs
-
-        if remote_server_address in SSHTunnel.__TUNNELS:
-            self.tunnel = SSHTunnel.__TUNNELS[remote_server_address]
-        else:
-            open_port = _find_free_port("127.0.0.1")
-            local_bind_address = ("127.0.0.1", open_port)
-
-            ssh_address, ssh_port = tunnel_server_address.split(":")
-            ssh_port = int(ssh_port)  # type: ignore
-
-            remote_bind_address, remote_bind_port = remote_server_address.split(":")
-            remote_bind_port = int(remote_bind_port)  # type: ignore
-
-            if private_key is not None:
-                ssh_password = None
-                ssh_private_key_password = password
-            else:
-                ssh_password = password
-                ssh_private_key_password = None
-
-            self.tunnel = SSHTunnelForwarder(
-                ssh_address_or_host=(ssh_address, ssh_port),
-                local_bind_address=local_bind_address,
-                remote_bind_address=(remote_bind_address, remote_bind_port),
-                ssh_username=username,
-                ssh_password=ssh_password,
-                ssh_private_key_password=ssh_private_key_password,
-                ssh_pkey=private_key,
-                **kwargs,
-            )
-
-    def start(self):
-        if not self.tunnel.is_active:
-            self.tunnel.start()
-
-    def stop(self):
-        if self.tunnel.tunnel_is_up:
-            self.tunnel.stop()
-
-    @property
-    def local_address(self) -> Tuple[str, int]:
-        return self.tunnel.local_bind_address
-
-
 class MongoStore(Store):
     """
     A Store that connects to a Mongo collection
@@ -742,7 +669,7 @@ def __init__(
 
         super().__init__(**kwargs)
 
-    def connect(self, force_reset: bool =False):
+    def connect(self, force_reset: bool = False):
         """
         Loads the files into the collection in memory
 
@@ -1005,9 +932,3 @@ def update(self, docs: Union[List[Dict], Dict], key: Union[List, str, None] = No
                 search_doc = {k: d[k] for k in key} if isinstance(key, list) else {key: d[key]}
 
                 self._collection.replace_one(search_doc, d, upsert=True)
-
-
-def _find_free_port(address="0.0.0.0"):
-    s = socket()
-    s.bind((address, 0))  # Bind to a free port provided by the host.
-    return s.getsockname()[1]  # Return the port number assigned.

src/maggma/stores/ssh_tunnel.py

@@ -0,0 +1,89 @@
+from socket import socket
+from typing import Dict, Optional, Tuple
+
+from monty.json import MSONable
+from sshtunnel import SSHTunnelForwarder
+
+
+class SSHTunnel(MSONable):
+    __TUNNELS: Dict[str, SSHTunnelForwarder] = {}
+
+    def __init__(
+        self,
+        tunnel_server_address: str,
+        remote_server_address: str,
+        local_port: Optional[int] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        private_key: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            tunnel_server_address: string address with port for the SSH tunnel server
+            remote_server_address: string address with port for the server to connect to
+            local_port: optional port to use for the local address (127.0.0.1);
+                if `None`, a random open port will be automatically selected
+            username: optional username for the ssh tunnel server
+            password: optional password for the ssh tunnel server; If a private_key is
+                supplied this password is assumed to be the private key password
+            private_key: ssh private key to authenticate to the tunnel server
+            kwargs: any extra args passed to the SSHTunnelForwarder
+        """
+
+        self.tunnel_server_address = tunnel_server_address
+        self.remote_server_address = remote_server_address
+        self.local_port = local_port
+        self.username = username
+        self.password = password
+        self.private_key = private_key
+        self.kwargs = kwargs
+
+        if remote_server_address in SSHTunnel.__TUNNELS:
+            self.tunnel = SSHTunnel.__TUNNELS[remote_server_address]
+        else:
+            if local_port is None:
+                local_port = _find_free_port("127.0.0.1")
+            local_bind_address = ("127.0.0.1", local_port)
+
+            ssh_address, ssh_port = tunnel_server_address.split(":")
+            ssh_port = int(ssh_port)  # type: ignore
+
+            remote_bind_address, remote_bind_port = remote_server_address.split(":")
+            remote_bind_port = int(remote_bind_port)  # type: ignore
+
+            if private_key is not None:
+                ssh_password = None
+                ssh_private_key_password = password
+            else:
+                ssh_password = password
+                ssh_private_key_password = None
+
+            self.tunnel = SSHTunnelForwarder(
+                ssh_address_or_host=(ssh_address, ssh_port),
+                local_bind_address=local_bind_address,
+                remote_bind_address=(remote_bind_address, remote_bind_port),
+                ssh_username=username,
+                ssh_password=ssh_password,
+                ssh_private_key_password=ssh_private_key_password,
+                ssh_pkey=private_key,
+                **kwargs,
+            )
+
+    def start(self):
+        if not self.tunnel.is_active:
+            self.tunnel.start()
+
+    def stop(self):
+        if self.tunnel.tunnel_is_up:
+            self.tunnel.stop()
+
+    @property
+    def local_address(self) -> Tuple[str, int]:
+        return self.tunnel.local_bind_address
+
+
+def _find_free_port(address="0.0.0.0"):
+    s = socket()
+    s.bind((address, 0))  # Bind to a free port provided by the host.
+    return s.getsockname()[1]  # Return the port number assigned.

tests/stores/test_ssh_tunnel.py

@@ -1,7 +1,8 @@
 import paramiko
 import pymongo
 import pytest
-from maggma.stores.mongolike import MongoStore, SSHTunnel
+from maggma.stores.mongolike import MongoStore
+from maggma.stores.ssh_tunnel import SSHTunnel
 from monty.serialization import dumpfn, loadfn
 from paramiko.ssh_exception import AuthenticationException, NoValidConnectionsError, SSHException