feat: Add drive-sync feature with push and pull commands

Co-authored-by: khnumdev <[email protected]>

Author: Copilot

cli.py

@@ -9,19 +9,30 @@
 from commands.analyze_kinds import analyze_kinds, print_summary_table
 from commands.analyze_entity_fields import analyze_field_contributions, print_field_summary
 from commands.cleanup_expired import cleanup_expired
+from commands.drive_sync import push_to_drive, pull_from_drive
 
-app = typer.Typer(help="Utilities for analyzing and managing local Datastore/Firestore (Datastore mode)", no_args_is_help=True)
+app = typer.Typer(
+    help="Utilities for analyzing and managing local Datastore/Firestore (Datastore mode)",
+    no_args_is_help=True,
+)
 
 # Aliases with flags only — no defaults here
 ConfigOpt = Annotated[Optional[str], typer.Option("--config", help="Path to config.yaml")]
 ProjectOpt = Annotated[Optional[str], typer.Option("--project", help="GCP/Emulator project id")]
-EmulatorHostOpt = Annotated[Optional[str], typer.Option("--emulator-host", help="Emulator host, e.g. localhost:8010")]
+EmulatorHostOpt = Annotated[
+    Optional[str], typer.Option("--emulator-host", help="Emulator host, e.g. localhost:8010")
+]
 LogLevelOpt = Annotated[Optional[str], typer.Option("--log-level", help="Logging level")]
 KindsOpt = Annotated[
     Optional[List[str]],
-    typer.Option("--kind", "-k", help="Kinds to process (omit or empty to process all in each namespace)")
+    typer.Option(
+        "--kind", "-k", help="Kinds to process (omit or empty to process all in each namespace)"
+    ),
+]
+SingleKindOpt = Annotated[
+    Optional[str], typer.Option("--kind", "-k", help="Kind to analyze (falls back to config.kind)")
 ]
-SingleKindOpt = Annotated[Optional[str], typer.Option("--kind", "-k", help="Kind to analyze (falls back to config.kind)")]
+
 
 def _load_cfg(
     config_path: Optional[str],
@@ -38,6 +49,7 @@ def _load_cfg(
         overrides["log_level"] = log_level
     return load_config(config_path, overrides)
 
+
 @app.command("analyze-kinds")
 def cmd_analyze_kinds(
     config: ConfigOpt = None,
@@ -64,17 +76,31 @@ def cmd_analyze_kinds(
     else:
         print_summary_table(rows)
 
+
 @app.command("analyze-fields")
 def cmd_analyze_fields(
     kind: SingleKindOpt = None,
-    namespace: Annotated[Optional[str], typer.Option("--namespace", "-n", help="Namespace to query (omit to use all)")] = None,
-    group_by: Annotated[Optional[str], typer.Option("--group-by", help="Group results by this field value (falls back to config.group_by_field)")] = None,
-    only_field: Annotated[Optional[List[str]], typer.Option("--only-field", help="Only consider these fields")] = None,
+    namespace: Annotated[
+        Optional[str],
+        typer.Option("--namespace", "-n", help="Namespace to query (omit to use all)"),
+    ] = None,
+    group_by: Annotated[
+        Optional[str],
+        typer.Option(
+            "--group-by",
+            help="Group results by this field value (falls back to config.group_by_field)",
+        ),
+    ] = None,
+    only_field: Annotated[
+        Optional[List[str]], typer.Option("--only-field", help="Only consider these fields")
+    ] = None,
     config: ConfigOpt = None,
     project: ProjectOpt = None,
     emulator_host: EmulatorHostOpt = None,
     log_level: LogLevelOpt = None,
-    output_json: Annotated[Optional[str], typer.Option("--output-json", help="Write raw JSON results to file")] = None,
+    output_json: Annotated[
+        Optional[str], typer.Option("--output-json", help="Write raw JSON results to file")
+    ] = None,
 ):
     cfg = _load_cfg(config, project, emulator_host, log_level)
 
@@ -100,17 +126,32 @@ def cmd_analyze_fields(
     else:
         print_field_summary(result)
 
+
 @app.command("cleanup")
 def cmd_cleanup(
     config: ConfigOpt = None,
     project: ProjectOpt = None,
     emulator_host: EmulatorHostOpt = None,
     log_level: LogLevelOpt = None,
     kind: KindsOpt = None,
-    ttl_field: Annotated[Optional[str], typer.Option("--ttl-field", help="TTL field name (falls back to config.ttl_field)")] = None,
-    delete_missing_ttl: Annotated[Optional[bool], typer.Option("--delete-missing-ttl", help="Delete when TTL field is missing (falls back to config.delete_missing_ttl)")] = None,
-    batch_size: Annotated[Optional[int], typer.Option("--batch-size", help="Delete batch size (falls back to config.batch_size)")] = None,
-    dry_run: Annotated[bool, typer.Option("--dry-run", help="Only report counts; do not delete")] = False,
+    ttl_field: Annotated[
+        Optional[str],
+        typer.Option("--ttl-field", help="TTL field name (falls back to config.ttl_field)"),
+    ] = None,
+    delete_missing_ttl: Annotated[
+        Optional[bool],
+        typer.Option(
+            "--delete-missing-ttl",
+            help="Delete when TTL field is missing (falls back to config.delete_missing_ttl)",
+        ),
+    ] = None,
+    batch_size: Annotated[
+        Optional[int],
+        typer.Option("--batch-size", help="Delete batch size (falls back to config.batch_size)"),
+    ] = None,
+    dry_run: Annotated[
+        bool, typer.Option("--dry-run", help="Only report counts; do not delete")
+    ] = False,
 ):
     cfg = _load_cfg(config, project, emulator_host, log_level)
 
@@ -127,6 +168,46 @@ def cmd_cleanup(
     deleted_sum = sum(totals.values())
     typer.echo(f"Total entities {'to delete' if dry_run else 'deleted'}: {deleted_sum}")
 
+
+@app.command("push")
+def cmd_push(
+    version: Annotated[
+        Optional[str], typer.Argument(help="Version name (defaults to today's date YYYY-mm-DD)")
+    ] = None,
+    overwrite: Annotated[
+        bool, typer.Option("-o", "--overwrite", help="Overwrite existing file with same name")
+    ] = False,
+    local_db: Annotated[
+        Optional[str],
+        typer.Option(
+            "--local-db", help="Path to local-db binary (falls back to config.local_db_path)"
+        ),
+    ] = None,
+    config: ConfigOpt = None,
+    log_level: LogLevelOpt = None,
+):
+    cfg = _load_cfg(config, None, None, log_level)
+    push_to_drive(cfg, version, overwrite, local_db)
+
+
+@app.command("pull")
+def cmd_pull(
+    version: Annotated[
+        Optional[str], typer.Argument(help="Version name (omit to download latest)")
+    ] = None,
+    local_db: Annotated[
+        Optional[str],
+        typer.Option(
+            "--local-db", help="Path to local-db binary (falls back to config.local_db_path)"
+        ),
+    ] = None,
+    config: ConfigOpt = None,
+    log_level: LogLevelOpt = None,
+):
+    cfg = _load_cfg(config, None, None, log_level)
+    pull_from_drive(cfg, version, local_db)
+
+
 if __name__ == "__main__":
     import sys
 

commands/__init__.py

@@ -2,20 +2,23 @@
 from .analyze_kinds import analyze_kinds, get_kind_stats, estimate_entity_count_and_size
 from .analyze_entity_fields import analyze_field_contributions, print_field_summary
 from .cleanup_expired import cleanup_expired
+from .drive_sync import push_to_drive, pull_from_drive
 from . import config as config
 
 __all__ = [
-	"AppConfig",
-	"load_config",
-	"build_client",
-	"list_namespaces",
-	"list_kinds",
-	"format_size",
-	"analyze_kinds",
-	"get_kind_stats",
-	"estimate_entity_count_and_size",
-	"analyze_field_contributions",
-	"print_field_summary",
-	"cleanup_expired",
-	"config",
+    "AppConfig",
+    "load_config",
+    "build_client",
+    "list_namespaces",
+    "list_kinds",
+    "format_size",
+    "analyze_kinds",
+    "get_kind_stats",
+    "estimate_entity_count_and_size",
+    "analyze_field_contributions",
+    "print_field_summary",
+    "cleanup_expired",
+    "push_to_drive",
+    "pull_from_drive",
+    "config",
 ]

commands/config.py

@@ -37,6 +37,9 @@ class AppConfig:
     # Logging
     log_level: str = "INFO"
 
+    # Drive sync settings
+    local_db_path: Optional[str] = None
+
 
 def _as_list(value: Optional[Iterable[str]]) -> List[str]:
     if value is None:
@@ -90,6 +93,8 @@ def load_config(path: Optional[str] = None, overrides: Optional[Dict] = None) ->
 
     config.log_level = str(merged.get("log_level", config.log_level)).upper()
 
+    config.local_db_path = merged.get("local_db_path", config.local_db_path)
+
     _configure_logging(config.log_level)
     return config
 
@@ -155,4 +160,3 @@ def format_size(bytes_size: int) -> str:
             return f"{size:.2f} {unit}"
         size /= 1024
     return f"{size:.2f} PB"
-

commands/drive_sync.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import logging
+import os
+import subprocess
+from datetime import datetime
+from typing import Optional
+
+from pydrive2.auth import GoogleAuth
+from pydrive2.drive import GoogleDrive
+
+from .config import AppConfig
+
+logger = logging.getLogger(__name__)
+
+
+def _authenticate_drive() -> GoogleDrive:
+    gauth = GoogleAuth()
+    gauth.LocalWebserverAuth()
+    return GoogleDrive(gauth)
+
+
+def _get_local_db_path(config: AppConfig, local_db_override: Optional[str]) -> str:
+    local_db = local_db_override or config.local_db_path
+    if not local_db:
+        raise ValueError("local-db path must be provided via --local-db or config.local_db_path")
+    if not os.path.exists(local_db):
+        raise FileNotFoundError(f"local-db binary not found at: {local_db}")
+    return local_db
+
+
+def _get_or_create_datastore_folder(drive: GoogleDrive) -> str:
+    file_list = drive.ListFile(
+        {
+            "q": "title='datastore' and mimeType='application/vnd.google-apps.folder' and trashed=false"
+        }
+    ).GetList()
+    if file_list:
+        return file_list[0]["id"]
+    folder = drive.CreateFile(
+        {"title": "datastore", "mimeType": "application/vnd.google-apps.folder"}
+    )
+    folder.Upload()
+    logger.info("Created /datastore folder in Google Drive")
+    return folder["id"]
+
+
+def _run_local_db_command(local_db_path: str, args: list[str]) -> None:
+    cmd = [local_db_path] + args
+    logger.info(f"Running: {' '.join(cmd)}")
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"Command failed: {result.stderr}")
+    if result.stdout:
+        logger.info(result.stdout)
+
+
+def push_to_drive(
+    config: AppConfig, version: Optional[str], overwrite: bool, local_db: Optional[str]
+) -> None:
+    local_db_path = _get_local_db_path(config, local_db)
+
+    if not version:
+        version = datetime.now().strftime("%Y-%m-%d")
+
+    backup_file = f"local-db-{version}.bin"
+
+    _run_local_db_command(local_db_path, ["stash", version])
+
+    if not os.path.exists(backup_file):
+        raise FileNotFoundError(f"Stash did not produce expected file: {backup_file}")
+
+    drive = _authenticate_drive()
+    folder_id = _get_or_create_datastore_folder(drive)
+
+    existing = drive.ListFile(
+        {"q": f"title='{backup_file}' and '{folder_id}' in parents and trashed=false"}
+    ).GetList()
+    if existing:
+        if overwrite:
+            logger.info(f"Overwriting existing file: {backup_file}")
+            file_to_upload = existing[0]
+        else:
+            raise FileExistsError(
+                f"File {backup_file} already exists in /datastore. Use -o to overwrite."
+            )
+    else:
+        file_to_upload = drive.CreateFile({"title": backup_file, "parents": [{"id": folder_id}]})
+
+    file_to_upload.SetContentFile(backup_file)
+    file_to_upload.Upload()
+    logger.info(f"Successfully uploaded {backup_file} to Google Drive /datastore")
+
+
+def pull_from_drive(config: AppConfig, version: Optional[str], local_db: Optional[str]) -> None:
+    local_db_path = _get_local_db_path(config, local_db)
+
+    drive = _authenticate_drive()
+    folder_id = _get_or_create_datastore_folder(drive)
+
+    if version:
+        backup_file = f"local-db-{version}.bin"
+        files = drive.ListFile(
+            {"q": f"title='{backup_file}' and '{folder_id}' in parents and trashed=false"}
+        ).GetList()
+        if not files:
+            raise FileNotFoundError(f"No backup found with version: {version}")
+        file_to_download = files[0]
+    else:
+        files = drive.ListFile(
+            {
+                "q": f"'{folder_id}' in parents and trashed=false and title contains 'local-db-' and title contains '.bin'",
+                "orderBy": "modifiedDate desc",
+                "maxResults": 1,
+            }
+        ).GetList()
+        if not files:
+            raise FileNotFoundError("No backups found in /datastore folder")
+        file_to_download = files[0]
+        backup_file = file_to_download["title"]
+
+    logger.info(f"Downloading {backup_file} from Google Drive")
+    file_to_download.GetContentFile(backup_file)
+
+    version_to_restore = backup_file.replace("local-db-", "").replace(".bin", "")
+    _run_local_db_command(local_db_path, ["restore", version_to_restore])
+    logger.info(f"Successfully restored backup: {version_to_restore}")

pyproject.toml

@@ -27,6 +27,7 @@ dependencies = [
   "google-cloud-datastore>=2.19.0",
   "PyYAML>=6.0.1",
   "typer>=0.12.3",
+  "pydrive2>=1.20.0",
 ]
 
 [project.optional-dependencies]

requirements.txt

@@ -1,3 +1,4 @@
 google-cloud-datastore>=2.19.0
 PyYAML>=6.0.1
-typer>=0.12.3
+typer>=0.12.3
+pydrive2>=1.20.0