Repository: planetlabs/planet-client-python (python)

Issue:
atomic file support
when downloading files, ensure files are written atomically

Pull request:
support atomic file writing/creation
closes #57

Patch:
diff --git a/planet/api/_fatomic.py b/planet/api/_fatomic.py
new file mode 100644
index 0000000..9ce6e2e
--- /dev/null
+++ b/planet/api/_fatomic.py
@@ -0,0 +1,76 @@
+# Copyright 2015 Planet Labs, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+''' Provide atomic file support - a subset of `fatomic` included here
+(https://github.com/abarnert/fatomic) as upstream is not in pypi
+'''
+import contextlib
+import errno
+import os
+import shutil
+import sys
+import tempfile
+import types
+
+# If we're on 3.3+, just use os.replace; if we're on POSIX, rename
+# and replace do the same thing.
+try:
+    _replace_file = os.replace
+except AttributeError:
+    if sys.platform != 'win32':
+        _replace_file = os.rename
+    else:
+        # This requires PyWin32 if you're on Windows. If that's not
+        # accepted, you can write a ctypes solution, but then you'll
+        # have to handle unicode-vs.-bytes strings and creating an
+        # OSError from GetLastError and so on yourself, which I don't
+        # feel like doing. (I'll accept a pull request from anyone
+        # else who does...)
+        import win32api
+        import win32con
+
+        def _replace_file(src, dst):
+            win32api.MoveFileEx(src, dst, win32con.MOVEFILE_REPLACE_EXISTING)
+
+
+@contextlib.contextmanager
+def atomic_open(filename, mode, *args, **kwargs):
+    if mode[0] not in 'wxa' or len(mode) > 1 and mode[1] == '+':
+        raise ValueError("invalid mode: '{}'".format(mode))
+    f = tempfile.NamedTemporaryFile(mode=mode,
+                                    prefix=os.path.basename(filename),
+                                    suffix='.tmp',
+                                    delete=False)
+    _discard = [False]
+    try:
+        if mode[0] == 'a':
+            try:
+                with open(filename, 'r'+mode[1:], *args, **kwargs) as fin:
+                    shutil.copyfileobj(fin, f)
+            except (OSError, IOError) as e:
+                if e.errno == errno.ENOENT:
+                    pass
+
+        def discard(self, _discard=_discard):
+            _discard[0] = True
+        f.discard = types.MethodType(discard, f)
+        yield f
+    # ischneider addition to fatomic - catch/raise any exception thrown in with
+    # block and force discarding
+    except:
+        _discard = [True]
+        raise
+    finally:
+        f.close()
+        if not _discard[0]:
+            _replace_file(f.name, filename)
diff --git a/planet/api/models.py b/planet/api/models.py
index c63ef75..09ab8d7 100644
--- a/planet/api/models.py
+++ b/planet/api/models.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ._fatomic import atomic_open
 from .utils import get_filename
 from .utils import check_status
 from datetime import datetime
@@ -110,7 +111,7 @@ class Body(object):
         if hasattr(file, 'write'):
             self._write(file, callback)
         else:
-            with open(file, 'wb') as fp:
+            with atomic_open(file, 'wb') as fp:
                 self._write(fp, callback)
 
 
diff --git a/planet/api/sync.py b/planet/api/sync.py
index b387507..9a98f2b 100644
--- a/planet/api/sync.py
+++ b/planet/api/sync.py
@@ -16,6 +16,7 @@ import json
 import os
 from os import path
 import threading
+from ._fatomic import atomic_open
 from .utils import write_to_file
 from .utils import strp_timestamp
 from .utils import strf_timestamp
@@ -98,7 +99,7 @@ class _SyncTool(object):
         if summary.latest:
             sync = self._read_sync_file()
             sync['latest'] = strf_timestamp(summary.latest)
-            with open(self.sync_file, 'wb') as fp:
+            with atomic_open(self.sync_file, 'wb') as fp:
                 fp.write(json.dumps(sync, indent=2).encode('utf-8'))
 
         return summary
@@ -144,7 +145,7 @@ class _SyncHandler(object):
         # write out metadata
         metadata = os.path.join(self.destination,
                                 '%s_metadata.json' % self.metadata['id'])
-        with open(metadata, 'wb') as fp:
+        with atomic_open(metadata, 'wb') as fp:
             fp.write(json.dumps(self.metadata, indent=2).encode('utf-8'))
 
         # summarize
diff --git a/planet/api/utils.py b/planet/api/utils.py
index 6a68f75..ed49829 100644
--- a/planet/api/utils.py
+++ b/planet/api/utils.py
@@ -17,6 +17,7 @@ from . import exceptions
 import json
 import os
 import re
+from ._fatomic import atomic_open
 
 _ISO_FMT = '%Y-%m-%dT%H:%M:%S.%f+00:00'
 
@@ -63,7 +64,7 @@ def build_conditions(workspace):
 
 def write_planet_json(contents):
     fname = _planet_json_file()
-    with open(fname, 'w') as fp:
+    with atomic_open(fname, 'w') as fp:
         fp.write(json.dumps(contents))
 
 


Test patch:
diff --git a/tests/test_atomic.py b/tests/test_atomic.py
new file mode 100644
index 0000000..615b777
--- /dev/null
+++ b/tests/test_atomic.py
@@ -0,0 +1,43 @@
+# Copyright 2015 Planet Labs, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from planet.api._fatomic import atomic_open
+
+
+def test_atomic_open(tmpdir):
+    outfile = str(tmpdir.join('foo'))
+
+    def assert_content_is(expected):
+        with open(outfile, 'r') as fp:
+            assert fp.read() == expected
+
+    # success case
+    with atomic_open(outfile, 'w') as fp:
+        fp.write('bar')
+    assert_content_is('bar')
+
+    # exception during write, assert file remains untouched
+    try:
+        with atomic_open(outfile, 'w') as fp:
+            fp.write('bazzy')
+            raise Exception('drat')
+    except:
+        assert_content_is('bar')
+    else:
+        assert False
+
+    # manual discarding
+    with atomic_open(outfile, 'w') as fp:
+            fp.write('bazzy')
+            fp.discard()
+    assert_content_is('bar')