Repository: seung-lab/cloud-files (language: python)

Issue:
Add Resumable Uploads for transfer_to



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index ffb2df4..eee1ed1 100644
--- a/README.md
+++ b/README.md
@@ -53,6 +53,7 @@ CloudFiles was developed to access files from object storage without ever touchi
 8. High speed gzip decompression using libdeflate (compared with zlib).
 9. Bundled CLI tool.
 10. Accepts iterator and generator input.
+11. Resumable transfers.
 
 ## Installation 
 
@@ -335,6 +336,9 @@ cloudfiles cp -c br s3://bkt/file.txt gs://bkt2/
 cloudfiles cp -c none s3://bkt/file.txt gs://bkt2/
 # pass from stdin (use "-" for source argument)
 find some_dir | cloudfiles cp - s3://bkt/
+# resumable transfers
+cloudfiles xfer init SRC DEST --db JOBNAME.db
+cloudfiles xfer execute JOBNAME.db --progress # can quit and resume
 # Get human readable file sizes from anywhere
 cloudfiles du -shc ./tmp gs://bkt/dir s3://bkt/dir
 # remove files
diff --git a/cloudfiles/__init__.py b/cloudfiles/__init__.py
index 9cc0a42..3c0958a 100644
--- a/cloudfiles/__init__.py
+++ b/cloudfiles/__init__.py
@@ -10,4 +10,5 @@
 """
 
 from .cloudfiles import CloudFile, CloudFiles, dl
-from .interfaces import reset_connection_pools
\ No newline at end of file
+from .interfaces import reset_connection_pools
+from .resumable_tools import ResumableTransfer
\ No newline at end of file
diff --git a/cloudfiles/resumable_tools.py b/cloudfiles/resumable_tools.py
new file mode 100644
index 0000000..22e40be
--- /dev/null
+++ b/cloudfiles/resumable_tools.py
@@ -0,0 +1,215 @@
+from typing import (
+  Any, Dict, Optional, 
+  Union, List, Tuple, 
+  Callable, Generator, 
+  Iterable, cast
+)
+
+import datetime
+import os
+import sqlite3
+
+from tqdm import tqdm
+
+from .lib import sip
+from .cloudfiles import CloudFiles
+
+# the maximum value of a host parameter number is 
+# SQLITE_MAX_VARIABLE_NUMBER, which defaults to 999 
+# for SQLite versions prior to 3.32.0 (2020-05-22) or 
+# 32766 for SQLite versions after 3.32.0. 
+# https://www.sqlite.org/limits.html
+SQLITE_MAX_PARAMS = 999
+
+# syntax that changes between sqlite and mysql
+# or easy adjustment if we ever need it
+BIND = '?'
+AUTOINC = "AUTOINCREMENT"
+INTEGER = "INTEGER"
+
+def now_msec():
+  return int(datetime.datetime.utcnow().timestamp() * 1000)
+
+class ResumableFileSet:
+  """
+  An interface to an sqlite database for starting and resuming
+  resumable uploads or downloads.
+  """
+  def __init__(self, db_path, lease_msec=0):
+    self.conn = sqlite3.connect(db_path)
+    self.lease_msec = int(lease_msec)
+
+  def __del__(self):
+    self.conn.close()
+
+  def delete(self):
+    cur = self.conn.cursor()
+    cur.execute("""DROP TABLE IF EXISTS filelist""")
+    cur.execute("""DROP TABLE IF EXISTS xfermeta""")
+    cur.close()
+
+  def create(self, src, dest, reencode=None):
+    cur = self.conn.cursor()
+
+    cur.execute("""DROP TABLE IF EXISTS filelist""")
+    cur.execute("""DROP TABLE IF EXISTS xfermeta""")
+
+    cur.execute(f"""
+      CREATE TABLE xfermeta (
+        id {INTEGER} PRIMARY KEY {AUTOINC},
+        source TEXT NOT NULL,
+        dest TEXT NOT NULL,
+        reencode TEXT,
+        created INTEGER NOT NULL
+      )
+    """)
+
+    cur.execute(
+      "INSERT INTO xfermeta VALUES (?,?,?,?,?)", [ 1, src, dest, reencode, now_msec() ]
+    );
+
+    cur.execute(f"""
+      CREATE TABLE filelist (
+        id {INTEGER} PRIMARY KEY {AUTOINC},
+        filename TEXT NOT NULL,
+        finished INTEGER NOT NULL,
+        lease INTEGER NOT NULL
+      )
+    """)
+    cur.execute("CREATE INDEX resumableidxfin ON filelist(finished,lease)")
+    cur.execute("CREATE INDEX resumableidxfile ON filelist(filename)")
+    cur.close()
+
+  def insert(self, fname_iter):
+    cur = self.conn.cursor()
+
+    # cur.execute("PRAGMA journal_mode = MEMORY")
+    # cur.execute("PRAGMA synchronous = OFF")
+
+    for filenames in sip(fname_iter, SQLITE_MAX_PARAMS):
+      bindlist = ",".join([f"({BIND},0,0)"] * len(filenames))
+      cur.execute(f"INSERT INTO filelist(filename,finished,lease) VALUES {bindlist}", filenames)
+      cur.execute("commit")
+
+    cur.close()    
+
+  def metadata(self):
+    cur = self.conn.cursor()
+    cur.execute("SELECT source, dest, reencode, created FROM xfermeta LIMIT 1")
+    row = cur.fetchone()
+
+    return {
+      "source": row[0],
+      "dest": row[1],
+      "reencode": row[2],
+      "created": row[3],
+    }
+
+  def mark_finished(self, fname_iter):
+    cur = self.conn.cursor()
+
+    for filenames in sip(fname_iter, SQLITE_MAX_PARAMS):
+      bindlist = ",".join([f"{BIND}"] * len(filenames))
+      cur.execute(f"UPDATE filelist SET finished = 1 WHERE filename in ({bindlist})", filenames)
+      cur.execute("commit")
+    cur.close()
+
+  def next(self, limit=None, block_size=200):
+    cur = self.conn.cursor()
+
+    N = 0
+
+    while True:
+      ts = now_msec() + self.lease_msec
+      cur.execute(f"""SELECT filename FROM filelist WHERE finished = 0 AND lease <= {ts} LIMIT {int(block_size)}""")
+      rows = cur.fetchmany(block_size)
+      N += len(rows)
+      if len(rows) == 0:
+        break
+      
+      filenames = [ x[0] for x in rows ]
+      bindlist = ",".join([f"{BIND}"] * len(filenames))
+      ts = now_msec() + self.lease_msec
+      cur.execute(f"UPDATE filelist SET lease = {ts} WHERE filename in ({bindlist})", filenames)
+      cur.execute("commit")
+
+      yield from filenames
+      
+      if limit and N >= limit:
+        break
+
+    cur.close()
+
+  def total(self):
+    cur = self.conn.cursor()
+    cur.execute(f"SELECT count(filename) FROM filelist")
+    res = cur.fetchone()
+    cur.close()
+    return int(res[0])
+
+  def remaining(self):
+    cur = self.conn.cursor()
+    cur.execute(f"SELECT count(filename) FROM filelist WHERE finished = 0")
+    res = cur.fetchone()
+    cur.close()
+    return int(res[0])
+
+  def available(self):
+    cur = self.conn.cursor()
+    ts = int(now_msec() + self.lease_msec)
+    cur.execute(f"SELECT count(filename) FROM filelist WHERE finished = 0 AND lease < {ts}")
+    res = cur.fetchone()
+    cur.close()
+    return int(res[0])
+
+  def release(self):
+    cur.execute(f"UPDATE filelist SET lease = 0")
+    cur.execute("commit")
+
+  def __len__(self):
+    return self.remaining()
+
+  def __iter__(self):
+    return self.next()
+
+class ResumableTransfer:
+  def __init__(self, db_path, lease_msec=0):
+    self.db_path = db_path
+    self.rfs = ResumableFileSet(db_path, lease_msec)
+
+  def __len__(self):
+    return len(self.rfs)
+
+  def init(self, src, dest, paths, reencode=None):
+    if isinstance(paths, str):
+      paths = list(CloudFiles(paths))
+    elif isinstance(paths, CloudFiles):
+      paths = list(paths)
+
+    self.rfs.create(src, dest, reencode)
+    self.rfs.insert(paths)
+
+  def execute(self, progress=False, block_size=200):
+    meta = self.rfs.metadata()
+
+    cf_src = CloudFiles(meta["source"])
+
+    total = self.rfs.total()
+    pbar = tqdm(total=total, desc="Transfer", disable=(not progress))
+    pbar.n = total - self.rfs.remaining()
+
+    with pbar:
+      pbar.refresh()
+      for paths in sip(self.rfs, block_size):
+        cf_src.transfer_to(meta["dest"], paths=paths, reencode=meta["reencode"])
+        self.rfs.mark_finished(paths)
+        
+        pbar.n = total - self.rfs.remaining()
+        pbar.refresh()
+
+  def close(self):
+    self.rfs.delete()
+    try:
+      os.remove(self.db_path)
+    except FileNotFoundError:
+      pass
diff --git a/cloudfiles_cli/cloudfiles_cli.py b/cloudfiles_cli/cloudfiles_cli.py
index 5ffcef8..5c3e9af 100644
--- a/cloudfiles_cli/cloudfiles_cli.py
+++ b/cloudfiles_cli/cloudfiles_cli.py
@@ -24,6 +24,7 @@
 import cloudfiles
 import cloudfiles.paths
 from cloudfiles import CloudFiles
+from cloudfiles.resumable_tools import ResumableTransfer
 from cloudfiles.compression import transcode
 from cloudfiles.paths import extract, get_protocol
 from cloudfiles.lib import (
@@ -165,9 +166,9 @@ def get_mfp(path, recursive):
 @click.argument("source", nargs=-1)
 @click.argument("destination", nargs=1)
 @click.option('-r', '--recursive', is_flag=True, default=False, help='Recursive copy.')
-@click.option('-c', '--compression', default='same', help="Destination compression type. Options: same (default), none, gzip, br, zstd")
-@click.option('--progress', is_flag=True, default=False, help="Show transfer progress.")
-@click.option('-b', '--block-size', default=128, help="Number of files to download at a time.")
+@click.option('-c', '--compression', default='same', help="Destination compression type. Options: same, none, gzip, br, zstd", show_default=True)
+@click.option('--progress', is_flag=True, default=False, help="Show transfer progress.", show_default=True)
+@click.option('-b', '--block-size', default=128, help="Number of files to download at a time.", show_default=True)
 @click.pass_context
 def cp(ctx, source, destination, recursive, compression, progress, block_size):
   """
@@ -305,6 +306,60 @@ def _cp_stdout(src, paths):
     content = res["content"].decode("utf8")
     sys.stdout.write(content)
 
+@main.group("xfer")
+def xfergroup():
+  """
+  Create named resumable transfers.
+
+  This is a more reliable version of
+  the cp command for large transfers.
+
+  Resumable transfers can be performed
+  in parallel by multiple clients. They
+  work by saving filenames to a sqlite3
+  database and checking them off.
+
+  To use run:
+
+  1. cloudfiles xfer init ... --db NAME
+
+  2. cloudfiles xfer execute NAME
+  """
+  pass
+
+@xfergroup.command("init")
+@click.argument("source")
+@click.argument("destination")
+@click.option('-c', '--compression', default='same', help="Destination compression type. Options: same, none, gzip, br, zstd", show_default=True)
+@click.option('--db', default=None, required=True, help="Filepath of the sqlite database used for tracking progress. Different databases should be used for each job.")
+def xferinit(source, destination, compression, db):
+  """(1) Create db of files from the source."""
+  if compression == "same":
+    compression = None
+  elif compression == "none":
+    compression = False
+
+  source = normalize_path(source)
+  destination = normalize_path(destination)
+
+  rt = ResumableTransfer(db)
+  rt.init(source, destination, source, compression)
+
+@xfergroup.command("execute")
+@click.argument("db")
+@click.option('--progress', is_flag=True, default=False, help="Show transfer progress.")
+@click.option('--lease-msec', default=0, help="(distributed transfers) Number of milliseconds to lease each task for.", show_default=True)
+@click.option('-b', '--block-size', default=200, help="Number of files to process at a time.", show_default=True)
+def xferexecute(db, progress, lease_msec, block_size):
+  """(2) Perform the transfer using the database.
+
+  Multiple clients can use the same database
+  for execution.
+  """
+  rt = ResumableTransfer(db, lease_msec)
+  rt.execute(progress=progress, block_size=block_size)
+  rt.close()
+
 @main.command()
 @click.argument("sources", nargs=-1)
 @click.option('-r', '--range', 'byte_range', default=None, help='Retrieve start-end bytes.')


Test patch:
diff --git a/automated_test.py b/automated_test.py
index 85d2516..cee082a 100644
--- a/automated_test.py
+++ b/automated_test.py
@@ -801,6 +801,43 @@ def mkfiles(mkdirname):
   except FileNotFoundError:
     pass
 
+def test_xfer():
+  import subprocess
+  from cloudfiles.lib import mkdir, touch
+  test_dir = os.path.dirname(os.path.abspath(__file__))
+  srcdir = os.path.join(test_dir, "testfiles_src")
+  destdir = os.path.join(test_dir, "testfiles_dest")
+  N = 100
+
+  try:
+    shutil.rmtree(srcdir)
+  except FileNotFoundError:
+    pass
+
+  if os.path.isfile(f"{destdir}"):
+    os.remove(destdir)
+
+  try:
+    shutil.rmtree(destdir)
+  except FileNotFoundError:
+    pass
+
+  def mkfiles(mkdirname):
+    try:
+      shutil.rmtree(mkdirname)
+    except FileNotFoundError:
+      pass
+    mkdir(mkdirname)
+    for i in range(N):
+      touch(os.path.join(mkdirname, str(i)))
+
+  mkfiles(srcdir)
+
+  subprocess.run(["cloudfiles", "xfer", "init", srcdir, destdir, "--db", "bananas.db"])
+  subprocess.run(["cloudfiles", "xfer", "execute", "bananas.db"])
+  assert len(os.listdir(srcdir)) == N
+  assert os.listdir(srcdir) == os.listdir(destdir)
+
 def test_cli_cat():
   import subprocess
   from cloudfiles.lib import mkdir, touch