Merge pull request #143 from mfulleratlassian/Issue142

Issue 142 - Allow users to provide a botocore session on creation of S3FileSystem

Author: martindurant

s3fs/core.py

@@ -119,6 +119,9 @@ class S3FileSystem(object):
         objects.
     config_kwargs : dict of parameters passed to ``botocore.client.Config``
     kwargs : other parameters for boto3 session
+    session : botocore Session object to be used for all connections.
+         This session will be used inplace of creating a new session inside S3FileSystem.
+
 
     Examples
     --------
@@ -140,9 +143,12 @@ def __init__(self, anon=False, key=None, secret=None, token=None,
                  use_ssl=True, client_kwargs=None, requester_pays=False,
                  default_block_size=None, default_fill_cache=True,
                  version_aware=False, config_kwargs=None,
-                 s3_additional_kwargs=None, **kwargs):
+                 s3_additional_kwargs=None, session=None, **kwargs):
         self.anon = anon
         self.session = None
+        self.passed_in_session = session
+        if self.passed_in_session:
+            self.session = self.passed_in_session
         self.key = key
         self.secret = secret
         self.token = token
@@ -205,7 +211,7 @@ def connect(self, refresh=False):
         anon, key, secret, kwargs, ckwargs, token, ssl = (
               self.anon, self.key, self.secret, self.kwargs,
               self.client_kwargs, self.token, self.use_ssl)
-
+        
         # Include the current PID in the connection key so that different
         # SSL connections are made for each process.
         tok = tokenize(anon, key, secret, kwargs, ckwargs, token,
@@ -214,24 +220,30 @@ def connect(self, refresh=False):
             self._conn.pop(tok, None)
         if tok not in self._conn:
             logger.debug("Open S3 connection.  Anonymous: %s", self.anon)
+
             if self.anon:
                 from botocore import UNSIGNED
                 conf = Config(connect_timeout=self.connect_timeout,
                               read_timeout=self.read_timeout,
                               signature_version=UNSIGNED, **self.config_kwargs)
-                self.session = boto3.Session(**self.kwargs)
+                if not self.passed_in_session:
+                    self.session = boto3.Session(**self.kwargs)
             else:
                 conf = Config(connect_timeout=self.connect_timeout,
                               read_timeout=self.read_timeout,
                               **self.config_kwargs)
-                self.session = boto3.Session(self.key, self.secret, self.token,
-                                             **self.kwargs)
+                if not self.passed_in_session:
+                    self.session = boto3.Session(self.key, self.secret, self.token,
+                                                 **self.kwargs)
+
             s3 = self.session.client('s3', config=conf, use_ssl=ssl,
                                      **self.client_kwargs)
             self._conn[tok] = (s3, self.session)
         else:
             s3, session = self._conn[tok]
-            self.session = session
+            if not self.passed_in_session:
+                self.session = session
+
         return s3
 
     def get_delegated_s3pars(self, exp=3600):
@@ -260,6 +272,8 @@ def get_delegated_s3pars(self, exp=3600):
                 'token': cred['SessionToken'], 'anon': False}
 
     def __getstate__(self):
+        if self.passed_in_session:
+            raise NotImplementedError
         d = self.__dict__.copy()
         del d['s3']
         del d['session']

s3fs/tests/test_s3fs.py

@@ -9,6 +9,7 @@
 from s3fs.core import S3FileSystem
 from s3fs.utils import seek_delimiter, ignoring, SSEParams
 import moto
+import boto3
 
 from botocore.exceptions import NoCredentialsError
 
@@ -1114,3 +1115,38 @@ def test_change_defaults_only_subsequent(monkeypatch):
     # Test the other file systems created to see if their block sizes changed
     assert fs_overridden.default_block_size == 64 * (1024 ** 2)
     assert fs_default.default_block_size == 5 * (1024 ** 2)
+
+
+
+def test_passed_in_session_set_correctly(s3):
+    session = boto3.session.Session()
+    s3 = S3FileSystem(session=session)
+    assert s3.passed_in_session is session
+    client = s3.connect()
+    assert s3.session is session
+
+
+def test_without_passed_in_session_set_unique(s3):
+    session = boto3.session.Session()
+    s3 = S3FileSystem()
+    assert s3.passed_in_session is None
+    client = s3.connect()
+    assert s3.session is not session
+
+
+def test_pickle_without_passed_in_session(s3):
+    s3 = S3FileSystem()
+    try:
+        s3.__getstate__()
+    except NotImplementedError:
+        pytest.fail("Unexpected NotImplementedError")
+
+
+def test_pickle_with_passed_in_session(s3):
+    session = boto3.session.Session()
+    s3 = S3FileSystem(session=session)
+    with pytest.raises(NotImplementedError):
+        s3.__getstate__()
+
+
+