Some sideload improvements

[skip ci]

Author: khemarato

scripts/gdrive_cache.py

@@ -70,6 +70,9 @@ def decorator(func: Callable[[], list[dict]]):
   return decorator
 
 add_backup_level(0,  "Cache Only", "Don't proactively fill the cache at all")
+
+# @backup_level(1, "lone published", "Non-video files included but without a website source")
+
 add_backup_level(10, "High", "All valuable items in need of backing up")
 add_backup_level(30, "Medium", "All items in active need of backing up")
 
@@ -292,35 +295,50 @@ def run_backup_level(level: BackupLevel, parallelism=14):
   graceful_threadmap(download_file_to_cache, files, unit='f', max_workers=parallelism)
   print(f"Done backing up to level {level.level}!")
 
-def sideload_file(file: Path, cache_dir: Path, parent_folder: str | None, move: bool):
+def sideload_file(file: Path, cache_dir: Path, parent_folder: str | None, move: bool, check: bool):
   """moves (or copies, if not `move`) `file` into `cache_dir`
   
   If the file doesn't exist in `gdrive.gcache` then it's uploaded to `parent_folder` (else skipped)"""
   assert cache_dir.is_dir()
   hashval = md5(file)
-  remote_files = gdrive.gcache.get_items_with_md5(hashval)
-  if not remote_files:
+  target_path = gdrive.gcache.get_cache_path_for_md5(hashval)
+  if not target_path:
     if not parent_folder:
       print(f"WARNING: Skipping untracked file {file}")
       return
     newid = gdrive.gcache.upload_file(file, folder_id=parent_folder)
-    remote_files = [gdrive.gcache.get_item(newid)]
-  target_path = cache_dir / hashval[:2] / f"{hashval[2:]}{file.suffix.lower()}"
+    target_path = gdrive.gcache.get_cache_path_for_md5(hashval)
+    assert target_path is not None
+  assert target_path.suffix == file.suffix.lower(), f"How did we get a different extension {target_path.suffix} for {file}?"
+  is_in_trash = target_path.parent.parent.parent.name == 'trash'
+  if is_in_trash:
+    if target_path.exists() and md5(target_path) != hashval:
+      new_path = target_path.with_stem(file.stem)
+      if new_path.exists() and md5(new_path) != hashval:
+        raise FileExistsError(f"{new_path} also exists with a different file. Idk what to do now")
+      target_path = new_path
+    print(f"WARNING: File was trashed. Placing in {target_path}")
   if target_path.exists():
-    if md5(target_path) == hashval:
+    if not check or md5(target_path) == hashval:
       if move:
         file.unlink()
       return
-    print(f"WARNING: Overwriting old, corrupted {target_path}")
+    print(f"Found corrupted: {target_path}")
     target_path.unlink()
-  target_path.parent.mkdir(exist_ok=True)
+  target_path.parent.mkdir(exist_ok=True, parents=is_in_trash)
   if move:
     file.rename(target_path)
   else:
     shutil.copy2(file, target_path)
 
 
-def sideload_main(files: Collection[Path], parent_folder: str | None = None, move: bool = True):
+def sideload_main(
+  files: Collection[Path],
+  parent_folder: str | None = None,
+  move: bool = True,
+  recurse: bool = False,
+  check: bool = False,
+):
   if parent_folder:
     if parent_folder.startswith(gdrive.FOLDER_LINK_PREFIX):
       parent_folder = gdrive.folderlink_to_id(parent_folder)
@@ -329,17 +347,23 @@ def sideload_main(files: Collection[Path], parent_folder: str | None = None, mov
       raise ValueError(f"Folder with ID {parent_folder} not found")
     if folder['mimeType'] != 'application/vnd.google-apps.folder':
       raise ValueError(f"{parent_folder} is not a Google Drive Folder, but a {folder['mimeType']}")
+  to_remove = set()
+  for file in files:
+    if not file.exists():
+      raise FileNotFoundError(file)
+    if file.is_dir():
+      if not recurse:
+        raise ValueError(f"{file} is a directory! Please specify files or use -r")
+      to_remove.add(file)
+      for child in file.iterdir():
+        files.append(child)
+  files = [f for f in files if f not in to_remove]
   if len(files) > 100:
     file_iter = tqdm(files)
   else:
     file_iter = iter(files)
   for file in file_iter:
-    if not file.exists():
-      print(f"WARNING: {file} does not exist!")
-      continue
-    if file.is_dir():
-      raise ValueError(f"{file} is a directory! Please only specify specific files")
-    sideload_file(file, gdrive.gcache.file_cache_dir, parent_folder, move)
+    sideload_file(file, gdrive.gcache.file_cache_dir, parent_folder, move, check)
 
 def get_saved_backup_level() -> int | None:
   with gdrive.gcache._lock:
@@ -358,6 +382,29 @@ def save_backup_level(level: int):
     )
     gdrive.gcache.conn.commit()
 
+def backup_main(new_max_level: int | None=None, parallelism: int=0):
+  import sys
+  if new_max_level is not None:
+    save_backup_level(new_max_level)
+    max_level = new_max_level
+  else:
+    max_level = get_saved_backup_level()
+    if max_level is None:
+      print("ERROR: No backup level supplied and no previous level found in the database. Please provide a --level.", file=sys.stderr)
+      sys.exit(1)
+  if max_level == 0:
+    print('The cache is set to "cach only" mode. Nothing further to do.')
+    sys.exit(0)
+  print(f"Will now back up GDrive to a level {max_level}")
+  for level in BACKUP_LEVELS.values():
+    if not level.finder:
+      continue
+    if level.level > max_level:
+      break
+    if parallelism < 1:
+      parallelism = 14
+    run_backup_level(level, parallelism=parallelism)
+  print(f"All files with priority <= {max_level} are now saved locally!")
 
 if __name__ == "__main__":
     import argparse
@@ -412,14 +459,26 @@ def backup_level(value):
         help="Copy files in (default: move)",
         default=False,
     )
+    sideload.add_argument(
+      "--recursive", "-r",
+      action="store_true",
+      help="Allow sideload to crawl directories",
+      default=False,
+    )
+    sideload.add_argument(
+      '--replace', '-f',
+      action="store_true",
+      default=False,
+      help="Don't assume the existing cache files are good",
+    )
 
     args = parser.parse_args()
 
     if not gdrive.gcache.file_cache_dir:
       gdrive.gcache.set_file_cache_dir()
 
     if args.command == "sideload":
-      sideload_main(args.files, args.parent_folder, move=(not args.copy))
+      sideload_main(args.files, args.parent_folder, move=(not args.copy), recurse=args.recursive, check=args.replace)
     elif args.command == "backup":
       if args.list_levels:
         print("Available Backup Levels:")
@@ -429,27 +488,6 @@ def backup_level(value):
             else:
                 print(f"  {lvl:3d}: {bl.name:<15} - {bl.description}")
       else:
-        import sys
-        if args.level is not None:
-          save_backup_level(args.level)
-        else:
-          args.level = get_saved_backup_level()
-          if args.level is None:
-            print("ERROR: No backup level supplied and no previous level found in the database. Please provide a --level.", file=sys.stderr)
-            sys.exit(1)
-        if args.level == 0:
-          print('The cache is set to "cach only" mode. Nothing further to do.')
-          sys.exit(0)
-        print(f"Will now back up GDrive to a level {args.level}")
-        for level in BACKUP_LEVELS.values():
-          if not level.finder:
-            continue
-          if level.level > args.level:
-            break
-          parallelism = args.threads
-          if parallelism < 1:
-            parallelism = 14
-          run_backup_level(level, parallelism=parallelism)
-        print(f"All files with priority <= {args.level} are now saved locally!")
+        backup_main(new_max_level=args.level, parallelism=args.threads)
     else:
       parser.print_help()

scripts/local_gdrive.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 from typing import List, Dict, Any, Optional, Callable, TypedDict, TypeVar, ParamSpec, Concatenate
 from time import sleep
+from mimetypes import guess_extension
 
 import gdrive_base
 
@@ -171,6 +172,12 @@ def _create_table(self):
         self.cursor.execute(create_trashed_items_sql)
         self.cursor.execute(create_properties_table_sql)
         self.cursor.executescript(create_index_sql)
+
+        try:
+            self.cursor.execute(f"ALTER TABLE trashed_drive_items ADD COLUMN trashed_time TEXT")
+        except sqlite3.OperationalError:
+            pass # Already exists
+
         self.conn.commit()
 
     def _prompt_for_file_cache_dir(self) -> Path:
@@ -373,7 +380,7 @@ def update(self):
         with yaspin(text="Pulling latest data from GDrive...") as ys:
             changes_page = changes_page['value']
             file_ids_to_fetch = set()
-            file_ids_removed = set()
+            file_ids_removed = {} # map id -> time
             while True:
                 changelist = gdrive_base.session().changes().list(includeRemoved=True, restrictToMyDrive=False, pageToken=changes_page, pageSize=1000).execute()
                 for change in changelist['changes']:
@@ -385,7 +392,7 @@ def update(self):
                           file = self.get_trashed_item(change['fileId'])
                           if file:
                             print(f"Trashed item was permanently deleted: \"{file['name']}\"")
-                        file_ids_removed.add(change['fileId'])
+                        file_ids_removed[change['fileId']] = change.get('time')
                     else:
                         file_ids_to_fetch.add(change['fileId'])
                 if 'nextPageToken' in changelist: # nextPageToken signals there is more to fetch
@@ -394,9 +401,9 @@ def update(self):
                 changes_page = changelist['newStartPageToken'] # newStartPageToken says come back later for more
                 break
         if len(file_ids_removed):
-            for fileId in file_ids_removed:
-                self._move_to_trash(fileId)
-        file_ids_to_fetch = list(file_ids_to_fetch - file_ids_removed)
+            for fileId, time in file_ids_removed.items():
+                self._move_to_trash(fileId, trashed_time=time)
+        file_ids_to_fetch = list(file_ids_to_fetch - set(file_ids_removed.keys()))
         if len(file_ids_to_fetch):
             all_items = gdrive_base.batch_get_files_by_id(file_ids_to_fetch, FILE_FIELDS)
             for item in tqdm(all_items, total=len(file_ids_to_fetch), desc="Fetching updated files"):
@@ -683,6 +690,25 @@ def find_duplicate_md5s(self) -> List[str]:
         self.cursor.execute(sql)
         return [row['md5_checksum'] for row in self.cursor.fetchall()]
 
+    def get_cache_path_for_md5(self, hashval: str) -> Path | None:
+        """Returns None if the hashval is unknown to me"""
+        assert len(hashval) == 32
+        remote_files = self.get_items_with_md5(hashval)
+        if not remote_files:
+            remote_files = self.get_trashed_items_with_md5(hashval)
+            if remote_files:
+                rm_date = max(f['trashed_time'] or f['modifiedTime'] for f in remote_files)
+                rm_date = datetime.fromisoformat(rm_date)
+                return self.file_cache_dir / 'trash' / str(rm_date.year) / f"{rm_date.month:02d}" / remote_files[0]['name']
+            return None
+        file = remote_files[0]
+        extension = ''
+        if '.' in file['name']:
+            extension = '.' + str(file['name']).split('.')[-1].lower()
+        if len(extension) < 1 or len(extension) > 6:
+            extension = guess_extension(file['mimeType']) or ''
+        return self.file_cache_dir / hashval[:2] / f"{hashval[2:]}{extension}"
+
     ########
     # Write-through Functions
     #
@@ -701,14 +727,26 @@ def trash_file(self, file_id: str):
 
         gdrive_base.trash_drive_file(file_id)
         with self._lock:
-            self._move_to_trash(file_id)
+            self._move_to_trash(file_id, trashed_time=UTC_NOW())
             self.conn.commit()
 
     @locked
-    def _move_to_trash(self, file_id: str):
+    def _move_to_trash(self, file_id: str, trashed_time: str = None):
+        # If we get a removal event from the API for a file already in the trash,
+        # only add the timestamp to the trash table if the item had a NULL trashed time before.
+        self.cursor.execute("SELECT trashed_time FROM trashed_drive_items WHERE id = ?", (file_id,))
+        row = self.cursor.fetchone()
+        if row:
+            if trashed_time and row['trashed_time'] is None:
+                self.cursor.execute("UPDATE trashed_drive_items SET trashed_time = ? WHERE id = ?", (trashed_time, file_id))
+            return
+
         self.cursor.execute("INSERT INTO trashed_drive_items SELECT * FROM drive_items WHERE id = ?", (file_id,))
         self.cursor.execute("DELETE FROM drive_items WHERE id = ?", (file_id,))
 
+        if trashed_time:
+            self.cursor.execute("UPDATE trashed_drive_items SET trashed_time = ? WHERE id = ?", (trashed_time, file_id))
+
     def move_file(self, file_id: str, folder: str, previous_parents=None, verbose=True):
         folder = gdrive_base.folderlink_to_id(folder) if folder.startswith("http") else folder
         with self._lock: