Spaces:
Running
Running
from collections import deque | |
class Transaction: | |
"""Filesystem transaction write context | |
Gathers files for deferred commit or discard, so that several write | |
operations can be finalized semi-atomically. This works by having this | |
instance as the ``.transaction`` attribute of the given filesystem | |
""" | |
def __init__(self, fs, **kwargs): | |
""" | |
Parameters | |
---------- | |
fs: FileSystem instance | |
""" | |
self.fs = fs | |
self.files = deque() | |
def __enter__(self): | |
self.start() | |
return self | |
def __exit__(self, exc_type, exc_val, exc_tb): | |
"""End transaction and commit, if exit is not due to exception""" | |
# only commit if there was no exception | |
self.complete(commit=exc_type is None) | |
if self.fs: | |
self.fs._intrans = False | |
self.fs._transaction = None | |
self.fs = None | |
def start(self): | |
"""Start a transaction on this FileSystem""" | |
self.files = deque() # clean up after previous failed completions | |
self.fs._intrans = True | |
def complete(self, commit=True): | |
"""Finish transaction: commit or discard all deferred files""" | |
while self.files: | |
f = self.files.popleft() | |
if commit: | |
f.commit() | |
else: | |
f.discard() | |
self.fs._intrans = False | |
self.fs._transaction = None | |
self.fs = None | |
class FileActor: | |
def __init__(self): | |
self.files = [] | |
def commit(self): | |
for f in self.files: | |
f.commit() | |
self.files.clear() | |
def discard(self): | |
for f in self.files: | |
f.discard() | |
self.files.clear() | |
def append(self, f): | |
self.files.append(f) | |
class DaskTransaction(Transaction): | |
def __init__(self, fs): | |
""" | |
Parameters | |
---------- | |
fs: FileSystem instance | |
""" | |
import distributed | |
super().__init__(fs) | |
client = distributed.default_client() | |
self.files = client.submit(FileActor, actor=True).result() | |
def complete(self, commit=True): | |
"""Finish transaction: commit or discard all deferred files""" | |
if commit: | |
self.files.commit().result() | |
else: | |
self.files.discard().result() | |
self.fs._intrans = False | |
self.fs = None | |