mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-08-18 14:30:34 +00:00
allaun
d28e9de6ca
Check file sizes before loading bytes into memory and stream large files directly into the tarball to prevent OOM errors on Steam Deck. Skip downloading files that already exist locally. Build: 0 errors (py_compile)
448 lines
17 KiB
Python
448 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
mcp_drive_gccl_compress.py — Download a Google Drive folder via MCP, GCCL-compress it, upload archive.
|
|
|
|
Workflow:
|
|
1. Recursively list all files in a Drive folder using the Google Drive MCP server.
|
|
2. Download each file to a local temp clone.
|
|
3. For each file, attempt GCCL delta compression using a previously-seen similar file as reference.
|
|
4. Pack compressed results into a .tar.zst archive with a JSON manifest.
|
|
5. Upload the archive back to Google Drive via MCP.
|
|
6. Optionally move the original folder to trash via MCP.
|
|
|
|
Usage:
|
|
python3 mcp_drive_gccl_compress.py <drive-folder-id> [options]
|
|
|
|
No rclone is used — only the Google Drive MCP server and gccl_waveprobe.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
# Insert Research Stack shim path for gccl_waveprobe
|
|
REPO = Path("/home/allaun/Research Stack")
|
|
SHIM_DIR = REPO / "4-Infrastructure/shim"
|
|
sys.path.insert(0, str(SHIM_DIR))
|
|
|
|
try:
|
|
import gccl_waveprobe as gw
|
|
except ImportError as e:
|
|
raise RuntimeError(f"Cannot import gccl_waveprobe from {SHIM_DIR}: {e}")
|
|
|
|
|
|
DEFAULT_MCP_COMMAND = ["npx", "-y", "@piotr-agier/google-drive-mcp"]
|
|
|
|
|
|
class DriveMcpClient:
|
|
"""Very thin JSON-RPC stdio client for the Google Drive MCP server."""
|
|
|
|
def __init__(self, env: Optional[Dict[str, str]] = None):
|
|
self._proc = subprocess.Popen(
|
|
DEFAULT_MCP_COMMAND,
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=False,
|
|
env={**os.environ, **(env or {})},
|
|
)
|
|
self._next_id = 1
|
|
self._initialize()
|
|
|
|
def _initialize(self) -> None:
|
|
self._call(
|
|
"initialize",
|
|
{
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {},
|
|
"clientInfo": {"name": "mcp-drive-gccl", "version": "1.0.0"},
|
|
},
|
|
)
|
|
|
|
def _send(self, msg: dict) -> None:
|
|
payload = json.dumps(msg).encode("utf-8") + b"\n"
|
|
self._proc.stdin.write(payload) # type: ignore[union-attr]
|
|
self._proc.stdin.flush() # type: ignore[union-attr]
|
|
|
|
def _recv(self) -> dict:
|
|
line = self._proc.stdout.readline() # type: ignore[union-attr]
|
|
if not line:
|
|
raise RuntimeError("MCP server closed stdout")
|
|
return json.loads(line.decode("utf-8"))
|
|
|
|
def _call(self, method: str, params: dict) -> dict:
|
|
msg = {"jsonrpc": "2.0", "id": self._next_id, "method": method, "params": params}
|
|
self._next_id += 1
|
|
self._send(msg)
|
|
while True:
|
|
resp = self._recv()
|
|
if resp.get("id") == msg["id"]:
|
|
if "error" in resp:
|
|
raise RuntimeError(f"MCP error: {resp['error']}")
|
|
return resp["result"]
|
|
|
|
def list_folder(
|
|
self,
|
|
folder_id: Optional[str] = None,
|
|
page_size: int = 100,
|
|
page_token: Optional[str] = None,
|
|
) -> Tuple[List[Dict[str, Any]], Optional[str]]:
|
|
params: Dict[str, Any] = {"pageSize": page_size}
|
|
if folder_id:
|
|
params["folderId"] = folder_id
|
|
if page_token:
|
|
params["pageToken"] = page_token
|
|
result = self._call("tools/call", {"name": "listFolder", "arguments": params})
|
|
text = self._extract_text(result)
|
|
items = self._parse_folder_listing(text)
|
|
next_token = self._parse_next_page_token(text)
|
|
return items, next_token
|
|
|
|
def list_folder_all(self, folder_id: Optional[str] = None) -> List[Dict[str, Any]]:
|
|
all_items: List[Dict[str, Any]] = []
|
|
token: Optional[str] = None
|
|
while True:
|
|
items, token = self.list_folder(folder_id=folder_id, page_token=token)
|
|
all_items.extend(items)
|
|
if not token:
|
|
break
|
|
return all_items
|
|
|
|
def download_file(self, file_id: str, local_path: Path) -> None:
|
|
result = self._call(
|
|
"tools/call",
|
|
{"name": "downloadFile", "arguments": {"fileId": file_id, "localPath": str(local_path), "overwrite": True}},
|
|
)
|
|
text = self._extract_text(result)
|
|
if "Downloaded" not in text and "Saved" not in text:
|
|
raise RuntimeError(f"Download failed: {text}")
|
|
|
|
def upload_file(self, local_path: Path, parent_folder_id: Optional[str] = None) -> str:
|
|
args: Dict[str, Any] = {"localPath": str(local_path)}
|
|
if parent_folder_id:
|
|
args["parentFolderId"] = parent_folder_id
|
|
result = self._call("tools/call", {"name": "uploadFile", "arguments": args})
|
|
text = self._extract_text(result)
|
|
if "Uploaded" not in text and "uploaded" not in text:
|
|
raise RuntimeError(f"Upload failed: {text}")
|
|
return text
|
|
|
|
def delete_item(self, item_id: str) -> None:
|
|
self._call("tools/call", {"name": "deleteItem", "arguments": {"itemId": item_id}})
|
|
|
|
def close(self) -> None:
|
|
if self._proc.poll() is None:
|
|
self._proc.stdin.close() # type: ignore[union-attr]
|
|
self._proc.wait(timeout=5)
|
|
|
|
@staticmethod
|
|
def _parse_next_page_token(text: str) -> Optional[str]:
|
|
for line in text.splitlines():
|
|
line = line.strip()
|
|
if line.lower().startswith("nextpagetoken:") or line.lower().startswith("next page token:"):
|
|
return line.split(":", 1)[1].strip()
|
|
if "nextpagetoken" in line.lower() and ":" in line:
|
|
return line.split(":", 1)[1].strip()
|
|
return None
|
|
|
|
@staticmethod
|
|
def _extract_text(result: dict) -> str:
|
|
for item in result.get("content", []):
|
|
if item.get("type") == "text":
|
|
return item.get("text", "")
|
|
return ""
|
|
|
|
@staticmethod
|
|
def _parse_folder_listing(text: str) -> List[Dict[str, Any]]:
|
|
items: List[Dict[str, Any]] = []
|
|
for line in text.splitlines():
|
|
line = line.strip()
|
|
if not line or line.startswith("Contents"):
|
|
continue
|
|
# Format: 📄 name (mime) [id: ID, path: ...] [created: ..., modified: ...]
|
|
# or 📁 name (ID: ID)
|
|
if line.startswith("📄"):
|
|
rest = line[2:].strip()
|
|
# Format: "name (ID: FILE_ID, path: ...)" or "name (mime) (ID: ...)"
|
|
if "(ID: " in rest:
|
|
name = rest.split("(ID: ")[0].strip()
|
|
# Remove trailing mime like "(text/plain)" from name
|
|
if name.endswith(")") and "(" in name:
|
|
name = name.rsplit("(", 1)[0].strip()
|
|
fid = rest.split("(ID: ")[1].split(",", 1)[0].rstrip(")").strip()
|
|
items.append({"type": "file", "name": name, "id": fid})
|
|
elif line.startswith("📁"):
|
|
# Format: name (ID: ID)
|
|
rest = line[2:].strip()
|
|
if "(ID: " in rest:
|
|
name = rest.split("(ID: ")[0].strip()
|
|
fid = rest.split("(ID: ")[1].rstrip(")").strip()
|
|
items.append({"type": "folder", "name": name, "id": fid})
|
|
return items
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
while chunk := f.read(64 * 1024):
|
|
h.update(chunk)
|
|
return h.hexdigest()
|
|
|
|
|
|
def gccl_compress_file(
|
|
src_path: Path,
|
|
reference_path: Optional[Path] = None,
|
|
threshold_q16: int = 32768,
|
|
max_gccl_size: int = 100 * 1024 * 1024,
|
|
) -> Tuple[bytes, Dict[str, Any]]:
|
|
"""Compress a single file with GCCL. Optionally use reference for delta.
|
|
Files larger than max_gccl_size are stored raw to avoid memory issues."""
|
|
original_size = src_path.stat().st_size
|
|
|
|
if original_size > max_gccl_size:
|
|
return b"", {
|
|
"mode": "raw-oversize",
|
|
"reference": str(reference_path) if reference_path else None,
|
|
"original_size": original_size,
|
|
"compressed_size": original_size,
|
|
"ratio": 1.0,
|
|
"decision": "OVERSIZE",
|
|
}
|
|
|
|
data = src_path.read_bytes()
|
|
reference = reference_path.read_bytes() if reference_path else None
|
|
|
|
result = gw.gccl_delta_compress(
|
|
data,
|
|
reference=reference,
|
|
gccl_threshold_q16=threshold_q16,
|
|
wave_probe=gw.WaveProbe(id=1, sample_rate=48000, buffer_size=256),
|
|
)
|
|
|
|
if result.gccl_decision == gw.Decision.ACCEPT and result.compressed_size < result.original_size:
|
|
return result.compressed, {
|
|
"mode": "delta" if reference else "raw-gccl",
|
|
"reference": str(reference_path) if reference_path else None,
|
|
"original_size": result.original_size,
|
|
"compressed_size": result.compressed_size,
|
|
"ratio": result.compression_ratio,
|
|
"decision": str(result.gccl_decision),
|
|
}
|
|
else:
|
|
# Compression rejected or expanded — store raw
|
|
return data, {
|
|
"mode": "raw",
|
|
"reference": str(reference_path) if reference_path else None,
|
|
"original_size": len(data),
|
|
"compressed_size": len(data),
|
|
"ratio": 1.0,
|
|
"decision": str(result.gccl_decision) if result else "N/A",
|
|
}
|
|
|
|
|
|
def download_folder_recursive(
|
|
client: DriveMcpClient,
|
|
folder_id: str,
|
|
local_root: Path,
|
|
rename_map: Optional[Dict[str, str]] = None,
|
|
) -> List[Dict[str, Any]]:
|
|
"""Recursively download a Drive folder. Returns manifest of downloaded files."""
|
|
files: List[Dict[str, Any]] = []
|
|
_download_folder_recursive(client, folder_id, local_root, files, rename_map or {})
|
|
return files
|
|
|
|
|
|
def _download_folder_recursive(
|
|
client: DriveMcpClient,
|
|
folder_id: str,
|
|
local_dir: Path,
|
|
files: List[Dict[str, Any]],
|
|
rename_map: Dict[str, str],
|
|
) -> None:
|
|
local_dir.mkdir(parents=True, exist_ok=True)
|
|
items = client.list_folder_all(folder_id)
|
|
for item in items:
|
|
safe_name = rename_map.get(item["id"], item["name"])
|
|
local_path = local_dir / safe_name
|
|
if item["type"] == "folder":
|
|
_download_folder_recursive(client, item["id"], local_path, files, rename_map)
|
|
else:
|
|
if local_path.exists():
|
|
print(f" [skip] {safe_name} already exists locally")
|
|
else:
|
|
print(f" downloading {safe_name} ...")
|
|
client.download_file(item["id"], local_path)
|
|
files.append({
|
|
"drive_id": item["id"],
|
|
"name": safe_name,
|
|
"local_path": str(local_path),
|
|
"sha256": sha256_file(local_path),
|
|
})
|
|
|
|
|
|
def build_reference_index(
|
|
files: List[Dict[str, Any]],
|
|
local_root: Path,
|
|
) -> Dict[str, Optional[Path]]:
|
|
"""For each file, pick a reference: same basename elsewhere, or largest previously seen file."""
|
|
by_basename: Dict[str, Path] = {}
|
|
refs: Dict[str, Optional[Path]] = {}
|
|
for f in files:
|
|
local_path = Path(f["local_path"])
|
|
basename = local_path.name
|
|
if basename in by_basename and by_basename[basename] != local_path:
|
|
refs[f["local_path"]] = by_basename[basename]
|
|
else:
|
|
# Use largest previously seen file as reference
|
|
refs[f["local_path"]] = max(by_basename.values(), key=lambda p: p.stat().st_size) if by_basename else None
|
|
by_basename[basename] = local_path
|
|
return refs
|
|
|
|
|
|
def compress_folder(
|
|
files: List[Dict[str, Any]],
|
|
local_root: Path,
|
|
archive_path: Path,
|
|
threshold_q16: int = 32768,
|
|
max_gccl_size: int = 100 * 1024 * 1024,
|
|
) -> Dict[str, Any]:
|
|
"""GCCL-compress downloaded files and pack into a .tar.zst archive."""
|
|
refs = build_reference_index(files, local_root)
|
|
manifest: Dict[str, Any] = {
|
|
"schema": "mcp_drive_gccl_archive_v1",
|
|
"created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
|
"base_dir": str(local_root),
|
|
"files": [],
|
|
}
|
|
|
|
temp_compressed = archive_path.parent / f"{archive_path.name}.tmp"
|
|
with tarfile.open(temp_compressed, "w") as tar:
|
|
for f in files:
|
|
src = Path(f["local_path"])
|
|
rel = src.relative_to(local_root)
|
|
ref = refs.get(f["local_path"])
|
|
|
|
compressed_bytes, meta = gccl_compress_file(src, ref, threshold_q16, max_gccl_size)
|
|
|
|
arc_name = f"{rel}.gccl" if meta["mode"] not in ("raw", "raw-oversize") else str(rel)
|
|
if meta["mode"] in ("raw", "raw-oversize"):
|
|
tar.add(src, arcname=arc_name)
|
|
else:
|
|
info = tarfile.TarInfo(name=arc_name)
|
|
info.size = len(compressed_bytes)
|
|
info.mtime = time.time()
|
|
tar.addfile(info, io.BytesIO(compressed_bytes))
|
|
|
|
manifest["files"].append({
|
|
"path": str(rel),
|
|
"archive_path": arc_name,
|
|
"drive_id": f["drive_id"],
|
|
"sha256": f["sha256"],
|
|
**meta,
|
|
})
|
|
|
|
# Re-pack with zstd
|
|
with open(temp_compressed, "rb") as src_f, open(archive_path, "wb") as dst_f:
|
|
subprocess.run(
|
|
["zstd", "-T0", "-19", "-q", "-f", "-o", str(archive_path)],
|
|
stdin=src_f,
|
|
check=True,
|
|
)
|
|
temp_compressed.unlink()
|
|
|
|
# Write manifest alongside
|
|
manifest_path = archive_path.with_suffix(".manifest.json")
|
|
manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
|
|
|
total_original = sum(m["original_size"] for m in manifest["files"])
|
|
total_compressed = archive_path.stat().st_size + manifest_path.stat().st_size
|
|
manifest["summary"] = {
|
|
"file_count": len(files),
|
|
"total_original_bytes": total_original,
|
|
"archive_bytes": archive_path.stat().st_size,
|
|
"manifest_bytes": manifest_path.stat().st_size,
|
|
"total_output_bytes": total_compressed,
|
|
"si_compression_ratio": total_original / total_compressed if total_compressed else 0,
|
|
}
|
|
manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
|
return manifest
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("folder_id", help="Google Drive folder ID to compress")
|
|
parser.add_argument("--work-dir", default="/tmp/mcp-drive-gccl", help="Local scratch directory")
|
|
parser.add_argument("--archive-name", help="Output archive name (default: folder_id.tar.zst)")
|
|
parser.add_argument("--upload-parent", help="Drive folder ID to upload archive into")
|
|
parser.add_argument("--trash-original", action="store_true", help="Move original folder to trash after upload")
|
|
parser.add_argument("--threshold-q16", type=int, default=32768, help="GCCL threshold in Q16.16 (default 32768 = 0.5)")
|
|
parser.add_argument("--max-gccl-size", type=int, default=100*1024*1024, help="Max bytes for GCCL compression; larger files stored raw (default 100 MiB)")
|
|
parser.add_argument("--dry-run", action="store_true", help="Download and compress but do not upload or trash")
|
|
args = parser.parse_args()
|
|
|
|
archive_name = args.archive_name or f"{args.folder_id}.tar.zst"
|
|
work_dir = Path(args.work_dir)
|
|
local_clone = work_dir / "clone" / args.folder_id
|
|
archive_path = work_dir / "out" / archive_name
|
|
manifest_path = archive_path.with_suffix(".manifest.json")
|
|
|
|
work_dir.mkdir(parents=True, exist_ok=True)
|
|
local_clone.parent.mkdir(parents=True, exist_ok=True)
|
|
archive_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
env = {"GOOGLE_DRIVE_OAUTH_CREDENTIALS": "/home/allaun/gcp-oauth.keys.json"}
|
|
client = DriveMcpClient(env=env)
|
|
|
|
try:
|
|
print(f"[*] Downloading folder {args.folder_id} to {local_clone}")
|
|
files = download_folder_recursive(client, args.folder_id, local_clone)
|
|
print(f"[+] Downloaded {len(files)} files")
|
|
|
|
if not files:
|
|
print("[-] No files found; aborting")
|
|
return 1
|
|
|
|
print(f"[*] GCCL-compressing into {archive_path}")
|
|
manifest = compress_folder(files, local_clone, archive_path, args.threshold_q16, args.max_gccl_size)
|
|
print(f"[+] Archive: {archive_path} ({archive_path.stat().st_size} bytes)")
|
|
print(f"[+] Manifest: {manifest_path}")
|
|
print(f"[+] SI compression ratio: {manifest['summary']['si_compression_ratio']:.2f}")
|
|
|
|
if not args.dry_run:
|
|
print("[*] Uploading archive to Google Drive")
|
|
archive_upload_text = client.upload_file(archive_path, args.upload_parent)
|
|
print(f"[+] Archive upload result: {archive_upload_text}")
|
|
|
|
print("[*] Uploading manifest to Google Drive")
|
|
manifest_upload_text = client.upload_file(manifest_path, args.upload_parent)
|
|
print(f"[+] Manifest upload result: {manifest_upload_text}")
|
|
|
|
if args.trash_original:
|
|
print(f"[*] Trashing original folder {args.folder_id}")
|
|
client.delete_item(args.folder_id)
|
|
print("[+] Original trashed")
|
|
else:
|
|
print("[DRY RUN] No upload or trash performed")
|
|
|
|
# Print manifest summary
|
|
print(json.dumps(manifest["summary"], indent=2))
|
|
return 0
|
|
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|