diff --git a/benchmarks/download_files_parallel.py b/benchmarks/download_files_parallel.py index 64d74cd7aac..217abfc82f0 100644 --- a/benchmarks/download_files_parallel.py +++ b/benchmarks/download_files_parallel.py @@ -44,8 +44,8 @@ def benchmark_gdrive_load_state(): connection = manager_do1._connection_router.connections[0] - personal_syftbox_folder_id = connection.get_personal_syftbox_folder_id() - file_metadatas = connection.get_file_metadatas_from_folder( + personal_syftbox_folder_id = connection.layout.get_personal_syftbox_folder_id() + file_metadatas = connection.layout.get_file_metadatas_from_folder( personal_syftbox_folder_id ) # valid_fname_objs = connection._get_valid_events_from_file_metadatas(file_metadatas) diff --git a/packages/syft-enclave/src/syft_enclaves/utils.py b/packages/syft-enclave/src/syft_enclaves/utils.py index 553c74bf1eb..6403696865f 100644 --- a/packages/syft-enclave/src/syft_enclaves/utils.py +++ b/packages/syft-enclave/src/syft_enclaves/utils.py @@ -3,7 +3,7 @@ MockDriveBackingStore, MockDriveService, ) -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.syftbox_manager import SyftboxManager from syft.sync.transport.protocols import Transport @@ -66,7 +66,7 @@ def create_clients(configs) -> tuple: def _mock_transport(backing_store: MockDriveBackingStore, email: str) -> Transport: service = MockDriveService(backing_store, email) return SyftboxManager._drive_transport( - GDriveConnection.from_service(email, service) + DriveCachedClient.from_service(email, service) ) diff --git a/packages/syft-rds/src/syft_rds/client.py b/packages/syft-rds/src/syft_rds/client.py index a27519c6392..f34a6b41c5c 100644 --- a/packages/syft-rds/src/syft_rds/client.py +++ b/packages/syft-rds/src/syft_rds/client.py @@ -394,9 +394,8 @@ def _share_any_datasets_with_peer(self, peer_email: str) -> None: ) except Exception: # One dataset failing (missing folder, quota, network) must not - # stop us sharing the rest with this peer. "alreadyShared" is - # already handled in _batch_add_permissions, so anything - # reaching here is a real failure worth a traceback. + # stop us sharing the rest with this peer. Anything reaching + # here is a real failure worth a traceback. logger.exception("Failed to share dataset %r with %s", tag, peer_email) # ------------------------------------------------------------------ # @@ -818,9 +817,7 @@ def _create_and_upload_collection( Returns ``(folder_id, content_hash)``. """ - from syft.sync.connections.drive.gdrive_transport import ( - CollectionFolder, - ) + from syft.sync.connections.drive.drive_collections import CollectionFolder content_hash = CollectionFolder.compute_hash(files) folder_id = self.sync_engine.create_collection_folder( diff --git a/packages/syft-rds/tests/test_dataset_upload_private.py b/packages/syft-rds/tests/test_dataset_upload_private.py index 4f3b6a057a6..a015c9bab6d 100644 --- a/packages/syft-rds/tests/test_dataset_upload_private.py +++ b/packages/syft-rds/tests/test_dataset_upload_private.py @@ -1,14 +1,11 @@ -from syft.sync.connections.drive.gdrive_transport import ( - GDriveConnection, - collection_name_query, -) +from syft.sync.connections.drive.drive_collections import collection_name_query from syft_rds import SyftRDSClient from syft_datasets.dataset_manager import ( DATASET_COLLECTION_PREFIX, PRIVATE_DATASET_COLLECTION_PREFIX, ) from dataset_test_utils import create_tmp_dataset_files -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of class TestDatasetUploadPrivate: @@ -146,11 +143,11 @@ def test_ds_cannot_find_private_folders_via_gdrive_query(self): ) # Get the DS's raw GDrive connection and search for private folders - ds_connection: GDriveConnection = gdrive_connection_of( + ds_cached_client = drive_cached_client_of( ds_manager.sync_engine._connection_router.transport ) results = ( - ds_connection.drive_service.files() + ds_cached_client.client.drive_service.files() .list( q=f"{collection_name_query(PRIVATE_DATASET_COLLECTION_PREFIX)} and trashed=false", fields="files(id,name)", @@ -160,11 +157,11 @@ def test_ds_cannot_find_private_folders_via_gdrive_query(self): assert len(results.get("files", [])) == 0 # DO should find it with the same query - do_connection: GDriveConnection = gdrive_connection_of( + do_cached_client = drive_cached_client_of( do_manager.sync_engine._connection_router.transport ) do_results = ( - do_connection.drive_service.files() + do_cached_client.client.drive_service.files() .list( q=f"{collection_name_query(PRIVATE_DATASET_COLLECTION_PREFIX)} and trashed=false", fields="files(id,name)", diff --git a/packages/syft-rds/tests/test_sync_manager.py b/packages/syft-rds/tests/test_sync_manager.py index 6d83bb46852..df9abb69af0 100644 --- a/packages/syft-rds/tests/test_sync_manager.py +++ b/packages/syft-rds/tests/test_sync_manager.py @@ -8,7 +8,7 @@ import pytest from syft.sync.connections.drive import mock_drive_service -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.syftbox_manager import SyftboxManager from syft_datasets import Dataset from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX @@ -22,7 +22,7 @@ create_tmp_dataset_files, create_tmp_dataset_files_with_parquet, ) -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of def path_for_job(do_email: str, ds_email: str, filename: str = "test.job") -> str: @@ -233,9 +233,9 @@ def test_dataset_only_mock_data_uploaded(): # Sync so DS receives the dataset ds_manager.sync() - files = gdrive_connection_of( + files = drive_cached_client_of( ds_manager.sync_engine._connection_router.transport - ).drive_service._backing_store.files + ).client.drive_service._backing_store.files list(files.keys()) file_objs = list(files.values()) file_objs_ex_dataset_yaml = [ @@ -284,11 +284,13 @@ def test_jobs(): # we do this by not always syncing on a file change, currently this logic is a bit of # a short cut, but we could do this based on timing eventually (if there are items in the # queue for longer than a certain time we start pushing) - connection_do = gdrive_connection_of( + do_cached_client = drive_cached_client_of( do_manager.sync_engine._connection_router.transport ) - inbox_folder_id = connection_do._get_own_datasite_inbox_id(ds_manager.email) - inbox_file_metadatas = connection_do.get_file_metadatas_from_folder(inbox_folder_id) + inbox_folder_id = do_cached_client.get_own_datasite_inbox_id(ds_manager.email) + inbox_file_metadatas = do_cached_client.get_file_metadatas_from_folder( + inbox_folder_id + ) assert len(inbox_file_metadatas) == 1 do_manager.sync() @@ -1040,9 +1042,9 @@ def test_ds_dataset_cache_aware_sync(): assert remote_hash is not None # Get the mock backing store and directories for creating second pair - mock_backing_store = gdrive_connection_of( + mock_backing_store = drive_cached_client_of( ds_manager.sync_engine._connection_router.transport - ).drive_service._backing_store + ).client.drive_service._backing_store ds_folder = ds_manager.syftbox_folder do_folder = do_manager.syftbox_folder ds_email = ds_manager.email @@ -1059,12 +1061,12 @@ def test_ds_dataset_cache_aware_sync(): ) # Replace mock backing store to share dataset collections - gdrive_connection_of( + drive_cached_client_of( ds_manager2.sync_engine._connection_router.transport - ).drive_service._backing_store = mock_backing_store - gdrive_connection_of( + ).client.drive_service._backing_store = mock_backing_store + drive_cached_client_of( do_manager2.sync_engine._connection_router.transport - ).drive_service._backing_store = mock_backing_store + ).client.drive_service._backing_store = mock_backing_store # Load peers (already approved in shared backing store) ds_manager2.load_peers() @@ -1140,9 +1142,9 @@ def test_do_dataset_cache_aware_sync(): assert len(do_manager.datasets.get_all()) == 1 # Get the mock backing store and directories for creating second pair - mock_backing_store = gdrive_connection_of( + mock_backing_store = drive_cached_client_of( ds_manager.sync_engine._connection_router.transport - ).drive_service._backing_store + ).client.drive_service._backing_store ds_folder = ds_manager.syftbox_folder do_folder = do_manager.syftbox_folder ds_email = ds_manager.email @@ -1159,12 +1161,12 @@ def test_do_dataset_cache_aware_sync(): ) # Replace mock backing store to share dataset collections - gdrive_connection_of( + drive_cached_client_of( ds_manager2.sync_engine._connection_router.transport - ).drive_service._backing_store = mock_backing_store - gdrive_connection_of( + ).client.drive_service._backing_store = mock_backing_store + drive_cached_client_of( do_manager2.sync_engine._connection_router.transport - ).drive_service._backing_store = mock_backing_store + ).client.drive_service._backing_store = mock_backing_store # Load peers (already approved in shared backing store) ds_manager2.load_peers() @@ -1207,9 +1209,9 @@ def test_in_memory_connection_load_state(): ) # Get the backing store that will persist across "restarts" - backing_store = gdrive_connection_of( + backing_store = drive_cached_client_of( do_manager1.sync_engine._connection_router.transport - ).drive_service._backing_store + ).client.drive_service._backing_store ds_folder = ds_manager1.syftbox_folder do_folder = do_manager1.syftbox_folder ds_email = ds_manager1.email @@ -1269,16 +1271,16 @@ def test_in_memory_connection_load_state(): # Connect to the same backing store service_do = mock_drive_service.MockDriveService(backing_store, do_email) - do_connection2 = GDriveConnection.from_service(do_email, service_do) + do_cached_client2 = DriveCachedClient.from_service(do_email, service_do) service_ds = mock_drive_service.MockDriveService(backing_store, ds_email) - ds_connection2 = GDriveConnection.from_service(ds_email, service_ds) + ds_cached_client2 = DriveCachedClient.from_service(ds_email, service_ds) do_manager2 = SyftRDSClient.from_config( - do_config2, SyftboxManager._drive_transport(do_connection2) + do_config2, SyftboxManager._drive_transport(do_cached_client2) ) ds_manager2 = SyftRDSClient.from_config( - ds_config2, SyftboxManager._drive_transport(ds_connection2) + ds_config2, SyftboxManager._drive_transport(ds_cached_client2) ) # Load peers diff --git a/packages/syft-rds/tests/test_version_mismatch_flow.py b/packages/syft-rds/tests/test_version_mismatch_flow.py index f41865e1a36..7b4728ac12d 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -7,11 +7,12 @@ from unittest.mock import patch -from syft.sync.connections.drive.gdrive_transport import ( +from syft.sync.connections.drive.drive_peering import DrivePeering +from syft.sync.connections.drive.gdrive_client import GOOGLE_FOLDER_MIME_TYPE +from syft.sync.connections.drive.drive_cached_client import ( GDRIVE_P2P_FOLDER_DATASITE_PREFIX, - GOOGLE_FOLDER_MIME_TYPE, - GDriveConnection, ) +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.mock_drive_service import ( MockDriveService, ) @@ -20,13 +21,13 @@ from syft.version import SYFT_VERSION from dataset_test_utils import create_test_project_folder, create_tmp_dataset_files -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of from syft.sync.syftbox_manager import SyftboxManager NEW_VERSION = "99.0.0" -def _find_p2p_folders(connection, peer_email): +def _find_p2p_folders(cached_client, peer_email): """Find all P2P folders involving a specific peer on the mock drive.""" q = ( f"mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" @@ -35,21 +36,25 @@ def _find_p2p_folders(connection, peer_email): " and trashed=false" ) results = ( - connection.drive_service.files().list(q=q, fields="files(id, name)").execute() + cached_client.client.drive_service.files() + .list(q=q, fields="files(id, name)") + .execute() ) return results.get("files", []) -def _find_versioned_p2p_folders(connection, peer_email, version): +def _find_versioned_p2p_folders(cached_client, peer_email, version): """Find P2P folders for a specific peer and version.""" - folders = _find_p2p_folders(connection, peer_email) + folders = _find_p2p_folders(cached_client, peer_email) return [f for f in folders if f"#{version}#" in f["name"]] def _get_backing_store(manager): """Extract the mock backing store from a manager's connection.""" - conn = gdrive_connection_of(manager.peer_manager.connection_router.transport) - return conn.drive_service._backing_store + cached_client = drive_cached_client_of( + manager.peer_manager.connection_router.transport + ) + return cached_client.client.drive_service._backing_store def _reinitialize_manager( @@ -69,8 +74,10 @@ def _reinitialize_manager( syftbox_folder=syftbox_folder, ) mock_service = MockDriveService(backing_store, email) - conn = GDriveConnection.from_service(email, mock_service) - manager = SyftRDSClient.from_config(config, SyftboxManager._drive_transport(conn)) + cached_client = DriveCachedClient.from_service(email, mock_service) + manager = SyftRDSClient.from_config( + config, SyftboxManager._drive_transport(cached_client) + ) if has_ds_role: manager.sync_engine.file_writer.add_callback( @@ -94,10 +101,10 @@ def _simulate_continue_on_mismatch(manager, backing_store): syftbox_folder = manager.syftbox_folder mock_service = MockDriveService(backing_store, email) - mock_conn = GDriveConnection.from_service(email, mock_service) + cached_client = DriveCachedClient.from_service(email, mock_service) def read_remote(e, t): - return mock_conn.read_own_peer_info_file().version_info + return DrivePeering(cached_client).read_own_peer_info().version_info with ( patch( @@ -172,15 +179,17 @@ def test_version_mismatch_continues_and_repairs(): assert do_manager.jobs[0].status == "done" # -- Step 4: Record P2P folders and personal folder id before upgrade -- - do_conn = gdrive_connection_of(do_manager.peer_manager.connection_router.transport) + do_cached_client = drive_cached_client_of( + do_manager.peer_manager.connection_router.transport + ) do_p2p_current = _find_versioned_p2p_folders( - do_conn, ds_manager.email, SYFT_VERSION + do_cached_client, ds_manager.email, SYFT_VERSION ) assert len(do_p2p_current) > 0 old_personal_name = f"{SYFT_VERSION}#{do_manager.email}" - old_personal_id = do_conn._find_folder_by_name( + old_personal_id = do_cached_client.client.find_folder_by_name( old_personal_name, - parent_id=do_conn.get_syftbox_folder_id(), + parent_id=do_cached_client.get_syftbox_folder_id(), owner_email=do_manager.email, ) assert old_personal_id is not None @@ -194,7 +203,11 @@ def test_version_mismatch_continues_and_repairs(): with ( patch("syft.version.SYFT_VERSION", NEW_VERSION), patch( - "syft.sync.connections.drive.gdrive_transport.SYFT_VERSION", + "syft.sync.connections.drive.drive_owner_store.SYFT_VERSION", + NEW_VERSION, + ), + patch( + "syft.sync.connections.drive.drive_cached_client.SYFT_VERSION", NEW_VERSION, ), patch("syft.sync.login_utils.SYFT_VERSION", NEW_VERSION), @@ -211,21 +224,21 @@ def test_version_mismatch_continues_and_repairs(): ) # Personal folder is adopted (same Drive id, new name), not recreated. - do_conn_new = gdrive_connection_of( + do_cached_client_new = drive_cached_client_of( do_manager.peer_manager.connection_router.transport ) new_personal_name = f"{NEW_VERSION}#{do_email}" - new_personal_id = do_conn_new._find_folder_by_name( + new_personal_id = do_cached_client_new.client.find_folder_by_name( new_personal_name, - parent_id=do_conn_new.get_syftbox_folder_id(), + parent_id=do_cached_client_new.get_syftbox_folder_id(), owner_email=do_email, ) assert new_personal_id is not None assert new_personal_id == old_personal_id assert ( - do_conn_new._find_folder_by_name( + do_cached_client_new.client.find_folder_by_name( old_personal_name, - parent_id=do_conn_new.get_syftbox_folder_id(), + parent_id=do_cached_client_new.get_syftbox_folder_id(), owner_email=do_email, ) is None @@ -249,13 +262,13 @@ def test_version_mismatch_continues_and_repairs(): syftbox_folder=ds_syftbox, ) - ds_conn_new = gdrive_connection_of( + ds_cached_client_new = drive_cached_client_of( ds_manager.peer_manager.connection_router.transport ) ds_personal_name = f"{NEW_VERSION}#{ds_email}" - ds_personal_id = ds_conn_new._find_folder_by_name( + ds_personal_id = ds_cached_client_new.client.find_folder_by_name( ds_personal_name, - parent_id=ds_conn_new.get_syftbox_folder_id(), + parent_id=ds_cached_client_new.get_syftbox_folder_id(), owner_email=ds_email, ) assert ds_personal_id is not None @@ -265,9 +278,13 @@ def test_version_mismatch_continues_and_repairs(): # P2P folders of the old version are reused, not replaced. A peer that # has not upgraded still looks for the old name. - do_p2p_new = _find_versioned_p2p_folders(do_conn_new, ds_email, NEW_VERSION) + do_p2p_new = _find_versioned_p2p_folders( + do_cached_client_new, ds_email, NEW_VERSION + ) assert len(do_p2p_new) == 0 - do_p2p_old = _find_versioned_p2p_folders(do_conn_new, ds_email, SYFT_VERSION) + do_p2p_old = _find_versioned_p2p_folders( + do_cached_client_new, ds_email, SYFT_VERSION + ) assert len(do_p2p_old) > 0 assert len(do_p2p_old) == len(do_p2p_current) @@ -331,29 +348,25 @@ def test_login_writes_the_remote_version_file_too(): use_in_memory_cache=False, check_versions=True, ) - conn = gdrive_connection_of(do_manager.peer_manager.connection_router.transport) - assert ( - conn.read_own_peer_info_file().version_info.syft_client_version == SYFT_VERSION - ) + peering = do_manager.peer_manager.connection_router.transport.peering + assert peering.read_own_peer_info().version_info.syft_client_version == SYFT_VERSION with ( patch("syft.version.SYFT_VERSION", NEW_VERSION), patch("syft.sync.version.version_info.SYFT_VERSION", NEW_VERSION), ): upgraded = _upgraded_manager(do_manager) - new_conn = gdrive_connection_of( - upgraded.peer_manager.connection_router.transport - ) + new_peering = upgraded.peer_manager.connection_router.transport.peering # Still the pre-upgrade version: nothing has refreshed it yet. assert ( - new_conn.read_own_peer_info_file().version_info.syft_client_version + new_peering.read_own_peer_info().version_info.syft_client_version == SYFT_VERSION ) _init_client_login(upgraded.sync_engine, sync=False, load_peers=False) assert ( - new_conn.read_own_peer_info_file().version_info.syft_client_version + new_peering.read_own_peer_info().version_info.syft_client_version == NEW_VERSION ) local = read_local_version(upgraded.syftbox_folder) @@ -394,10 +407,10 @@ def test_the_mismatch_prompt_does_not_return_after_a_login(): handle_potential_version_mismatches_on_login, ) - conn = gdrive_connection_of(do_manager.peer_manager.connection_router.transport) + peering = do_manager.peer_manager.connection_router.transport.peering with patch( "syft.sync.login_utils._read_remote_version", - side_effect=lambda e, t: conn.read_own_peer_info_file().version_info, + side_effect=lambda e, t: peering.read_own_peer_info().version_info, ): # The check runs before the client exists, so it reads the files # the previous client version left behind. diff --git a/scripts/check_token.py b/scripts/check_token.py index 0501a05ba9a..a9e3b095f36 100644 --- a/scripts/check_token.py +++ b/scripts/check_token.py @@ -9,7 +9,7 @@ from google.oauth2.credentials import Credentials from googleapiclient.errors import HttpError -from syft.sync.connections.drive.gdrive_transport import ( +from syft.sync.connections.drive.gdrive_client import ( SCOPES, build_drive_service, ) diff --git a/scripts/reproduce_stale_socket.py b/scripts/reproduce_stale_socket.py index bb92477077e..2594ba0f161 100644 --- a/scripts/reproduce_stale_socket.py +++ b/scripts/reproduce_stale_socket.py @@ -32,7 +32,7 @@ from google.oauth2.credentials import Credentials from syft.sync.connections.drive.gdrive_retry import execute_with_retries -from syft.sync.connections.drive.gdrive_transport import build_drive_service +from syft.sync.connections.drive.gdrive_client import build_drive_service SCOPES = ["https://www.googleapis.com/auth/drive"] REPO_ROOT = Path(__file__).resolve().parent.parent diff --git a/syft/gdrive_utils.py b/syft/gdrive_utils.py index b84201edb55..0139a2c63c9 100644 --- a/syft/gdrive_utils.py +++ b/syft/gdrive_utils.py @@ -4,29 +4,29 @@ import re import time from pathlib import Path -from urllib.parse import urlparse, parse_qs +from urllib.parse import parse_qs, urlparse from google.oauth2.credentials import Credentials as GoogleCredentials from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.errors import HttpError from googleapiclient.http import MediaIoBaseDownload -from syft.sync.connections.drive.gdrive_transport import ( - build_drive_service, +from syft.sync.config.config import settings +from syft.sync.connections.drive.gdrive_client import ( GOOGLE_FOLDER_MIME_TYPE, + build_drive_service, ) -from syft.sync.connections.drive.gdrive_transport import SCOPES as GDRIVE_SCOPES +from syft.sync.connections.drive.gdrive_client import SCOPES as GDRIVE_SCOPES from syft.sync.connections.drive.gdrive_retry import ( execute_with_retries, next_chunk_with_retries, ) +from syft.sync.environments.environment import Environment from syft.sync.utils.syftbox_utils import ( - check_env, _resolve_email, _resolve_token_path, + check_env, ) -from syft.sync.environments.environment import Environment -from syft.sync.config.config import settings # Matches strings that could plausibly be a Drive ID: # - Only base64url chars (alphanumeric, dash, underscore) @@ -301,7 +301,7 @@ def delete_remote_syftbox( """Delete all SyftBox state from Google Drive. This is a standalone utility that does not require a full SyftboxManager. - It creates a temporary GDriveConnection to find and delete all SyftBox + It logs in to Drive on its own to find and delete all SyftBox files/folders from Google Drive. Note: This function does NOT broadcast ``is_deleted`` events to peers. @@ -313,22 +313,26 @@ def delete_remote_syftbox( this is auto-detected if not provided. verbose: If True (default), print deletion progress. """ - from syft.sync.connections.drive.gdrive_transport import GDriveConnection + from syft.sync.connections.drive.drive_cached_client import DriveCachedClient + from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore email = _resolve_email(email) token_path = _resolve_token_path(token_path) - conn = GDriveConnection.from_token_path(email=email, token_path=token_path) + cached_client = DriveCachedClient.from_token_path( + email=email, token_path=token_path + ) + store = DriveOwnerStore(cached_client) # Gather file IDs via folder hierarchy - folder_file_ids = set(conn.gather_all_file_and_folder_ids()) + folder_file_ids = set(store.gather_all_file_and_folder_ids()) # Find the orphaned syft files in My Drive root - orphaned_file_ids = set(conn.find_orphaned_syft_files()) + orphaned_file_ids = set(store.find_orphaned_syft_files()) all_file_ids = list(folder_file_ids | orphaned_file_ids) start = time.time() - conn.delete_multiple_files_by_ids(all_file_ids) + cached_client.client.delete_multiple_files_by_ids(all_file_ids) elapsed = time.time() - start if verbose: @@ -342,4 +346,4 @@ def delete_remote_syftbox( else: print() - conn.reset_caches() + store.reset_caches() diff --git a/syft/sync/connections/base_connection.py b/syft/sync/connections/base_connection.py index 5ed17aae862..d1b297089fa 100644 --- a/syft/sync/connections/base_connection.py +++ b/syft/sync/connections/base_connection.py @@ -1,6 +1,4 @@ -from typing import ClassVar, Type from pydantic import BaseModel -from syft.sync.messages.proposed_filechange import ProposedFileChangesMessage class FileCollection(BaseModel): @@ -17,145 +15,5 @@ class FileCollection(BaseModel): class ConnectionConfig(BaseModel): - connection_type: ClassVar[Type["SyftboxPlatformConnection"]] - - -class SyftboxPlatformConnection(BaseModel): - config: ConnectionConfig | None = None - - def get_authenticated_email(self) -> str: - """Return the email of the account behind the underlying credentials/transport.""" - raise NotImplementedError() - - def watcher_send_proposed_file_changes_message( - self, proposed_file_change_message: ProposedFileChangesMessage - ): - raise NotImplementedError() - - @classmethod - def from_config(cls, config: ConnectionConfig): - return config.connection_type.from_config(config) - - def clone_for_thread(self) -> "SyftboxPlatformConnection": - """A connection that another thread may use. - - This one, unless the backend's client cannot be shared across threads. - """ - return self - - # ========================================================================= - # GENERIC PREFIX-PARAMETERIZED COLLECTION METHODS - # ========================================================================= - - def owner_create_collection_folder( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> str: - raise NotImplementedError() - - def owner_tag_collection_as_any( - self, prefix: str, tag: str, content_hash: str - ) -> None: - raise NotImplementedError() - - def owner_share_collection( - self, prefix: str, tag: str, content_hash: str, users: list[str] - ) -> None: - raise NotImplementedError() - - def owner_upload_collection_files( - self, prefix: str, tag: str, content_hash: str, files: dict[str, bytes] - ) -> None: - raise NotImplementedError() - - def owner_list_collections(self, prefix: str) -> list[str]: - raise NotImplementedError() - - def owner_list_all_collections_with_permissions( - self, prefix: str - ) -> list[FileCollection]: - raise NotImplementedError() - - def owner_delete_collection(self, prefix: str, tag: str) -> None: - """Delete every layout of ``tag`` published under ``prefix``.""" - raise NotImplementedError() - - def watcher_list_collections(self, prefix: str) -> list[dict]: - """Collections a peer shared with us, in every layout they published. - - Each dict carries owner_email, tag, content_hash and variant. ``prefix`` - matches all layouts, so the caller picks the one it can read. - """ - raise NotImplementedError() - - def watcher_download_collection( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> dict[str, bytes]: - raise NotImplementedError() - - def watcher_get_collection_file_metadatas( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> list[dict]: - raise NotImplementedError() - - def watcher_download_collection_file(self, file_id: str) -> bytes: - raise NotImplementedError() - - # ========================================================================= - # RAW BYTES TRANSPORT (used by ConnectionRouter for encryption) - # ========================================================================= - - def watcher_send_raw_bytes_to_inbox( - self, recipient: str, filename: str, data: bytes - ) -> None: - raise NotImplementedError() - - def owner_get_inbox_file_metadatas(self, sender_email: str) -> list[dict]: - """The messages in the inbox from ``sender_email``, oldest first, each - with ``file_id`` and ``file_name``.""" - raise NotImplementedError() - - def owner_write_raw_bytes_to_outbox( - self, recipient: str, filename: str, data: bytes - ) -> None: - raise NotImplementedError() - - def watcher_download_raw_events_from_outbox( - self, peer_email: str, since_timestamp: float | None - ) -> list[bytes]: - raise NotImplementedError() - - # ========================================================================= - # RAW BYTES TRANSPORT — owner's own storage (events, checkpoints, rolling state) - # ========================================================================= - - def owner_write_raw_bytes_to_syftbox(self, filename: str, data: bytes) -> str: - raise NotImplementedError() - - def owner_download_raw_bytes_by_id(self, file_id: str) -> bytes: - raise NotImplementedError() - - def owner_download_all_raw_events_from_syftbox(self) -> list[bytes]: - raise NotImplementedError() - - def upload_raw_checkpoint(self, filename: str, data: bytes) -> str: - raise NotImplementedError() - - def download_raw_latest_checkpoint(self) -> bytes | None: - raise NotImplementedError() - - def upload_raw_incremental_checkpoint(self, filename: str, data: bytes) -> str: - raise NotImplementedError() - - def download_all_raw_incremental_checkpoints(self) -> list[bytes]: - raise NotImplementedError() - - def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: - raise NotImplementedError() - - def download_raw_rolling_state(self) -> bytes | None: - raise NotImplementedError() - - def download_raw_events_since_timestamp( - self, since_timestamp: float - ) -> list[bytes]: - raise NotImplementedError() + """How to log in to one platform. Each platform has its own subclass, such + as ``GdriveConnectionConfig``.""" diff --git a/syft/sync/connections/drive/drive_cached_client.py b/syft/sync/connections/drive/drive_cached_client.py new file mode 100644 index 00000000000..4d6d31d0bcb --- /dev/null +++ b/syft/sync/connections/drive/drive_cached_client.py @@ -0,0 +1,633 @@ +"""One Google Drive account, and where syft keeps its folders in it. + +``DriveCachedClient`` logs in, finds and creates syft's folders and caches their +ids. A folder whose name holds the client version follows the version rules in +this module. +""" + +import re +from dataclasses import dataclass, field +from functools import partial +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable + +from google.oauth2.credentials import Credentials as GoogleCredentials +from pydantic import BaseModel, field_validator + +from syft.sync.connections.drive.gdrive_client import ( + GOOGLE_FOLDER_MIME_TYPE, + SCOPES, + GDriveClient, + build_drive_service, +) +from syft.sync.connections.drive.gdrive_retry import execute_with_retries +from syft.sync.events.file_change_event import FileChangeEventsMessageFileName +from syft.sync.transport.naming import timestamp_from_message_name +from syft.sync.utils.syftbox_utils import check_env +from syft.sync.version.version_info import _parse_semver +from syft.version import SYFT_VERSION + +if TYPE_CHECKING: + from syft.sync.connections.drive.grdrive_config import GdriveConnectionConfig + from syft.sync.connections.drive.mock_drive_service import MockDriveService + +SYFTBOX_FOLDER = "SyftBox" +GDRIVE_P2P_FOLDER_DATASITE_PREFIX = "syft_datasite" + + +class GdriveArchiveFolder(BaseModel): + sender_email: str + recipient_email: str + + def as_string(self) -> str: + return f"syft_{self.sender_email}_to_{self.recipient_email}_archive" + + +class GdrivePersonalSyftboxFolder(BaseModel): + """Folder for personal SyftBox: {version}#{email} under /SyftBox/""" + + email: str + + def as_string(self) -> str: + return f"{SYFT_VERSION}#{self.email}" + + +# An email part of a folder name goes into Drive queries and local paths. It has +# one '@', a '.' in the domain, and no quote, slash, '#', whitespace or control +# character. It is not RFC 5322. +_EMAIL_PART = r"[^@\s'\"\\/#\x00-\x1f\x7f]+" +_FOLDER_EMAIL_RE = re.compile(rf"{_EMAIL_PART}@{_EMAIL_PART}\.{_EMAIL_PART}") + + +class GdriveP2PFolder(BaseModel): + """Folder for peer communication: syft_datasite#version#datasite_email#inbox|outbox#peer_email""" + + datasite_email: str + folder_type: str # "inbox" or "outbox" + peer_email: str + + @field_validator("datasite_email", "peer_email") + @classmethod + def _check_email(cls, v: str) -> str: + if not _FOLDER_EMAIL_RE.fullmatch(v): + raise ValueError(f"Invalid email in P2P folder name: {v!r}") + return v + + def as_string(self) -> str: + return f"{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#{SYFT_VERSION}#{self.datasite_email}#{self.folder_type}#{self.peer_email}" + + @classmethod + def from_name(cls, name: str) -> "GdriveP2PFolder": + parts = name.split("#") + if len(parts) != 5 or parts[0] != GDRIVE_P2P_FOLDER_DATASITE_PREFIX: + raise ValueError(f"Invalid P2P folder name: {name}") + return cls(datasite_email=parts[2], folder_type=parts[3], peer_email=parts[4]) + + +# Helpers for finding folders whose names embed SYFT_VERSION. Folder +# names use '#' or '-' as field separators with the version as one field, +# so we walk those fields looking for an X.Y.Z-shaped chunk -- no per-format +# parser needed. Patch-compat means same major.minor (matches the protocol +# definition in version_info.is_compatible_with). + + +def _looks_like_version(s: str) -> bool: + """True if s is shaped like 'X.Y.Z' with all numeric parts.""" + parts = s.split(".") + return len(parts) == 3 and all(p.isdigit() for p in parts) + + +def _extract_version_from_name(name: str) -> str | None: + """Return the first '#'/'-'-separated field that looks like a semver.""" + for chunk in name.replace("-", "#").split("#"): + if _looks_like_version(chunk): + return chunk + return None + + +# A folder id and name, with the version from the name. The version fields come +# first, so the default sort puts these in version order. +_VersionedFolder = tuple[int, int, int, str, str] + + +def _sorted_by_version(folders: list[tuple[str, str]]) -> list[tuple[str, str]]: + """Folders from the lowest version to the highest. + + A name with no readable version sorts first, so a versioned folder always + wins when the caller takes the last entry. + """ + + def key(entry: tuple[str, str]) -> tuple[int, int, int]: + version_str = _extract_version_from_name(entry[1]) + if version_str is None: + return (-1, -1, -1) + try: + return _parse_semver(version_str) + except ValueError: + return (-1, -1, -1) + + return sorted(folders, key=key) + + +def _partition_by_version( + folders: list[tuple[str, str]], + current_version: str | None = None, +) -> tuple[list[tuple[str, str]], list[tuple[str, str]], list[tuple[str, str]]]: + """Split folders into (compatible, older, newer) by the version in the name. + + Compatible means the same major and minor as the current version. The function + drops a folder that has no version in its name. Each list starts at the lowest + version. + """ + if current_version is None: + current_version = SYFT_VERSION + try: + current = _parse_semver(current_version) + except ValueError: + return [], [], [] + + compatible: list[_VersionedFolder] = [] + older: list[_VersionedFolder] = [] + newer: list[_VersionedFolder] = [] + for fid, name in folders: + version_str = _extract_version_from_name(name) + if version_str is None: + continue + try: + found = _parse_semver(version_str) + except ValueError: + continue + entry = (*found, fid, name) + if found[:2] == current[:2]: + compatible.append(entry) + elif found < current: + older.append(entry) + else: + newer.append(entry) + + def _ordered(entries: list[_VersionedFolder]) -> list[tuple[str, str]]: + return [(fid, name) for *_, fid, name in sorted(entries)] + + return _ordered(compatible), _ordered(older), _ordered(newer) + + +def valid_events_from_file_metadatas( + file_metadatas: list[dict], +) -> list[FileChangeEventsMessageFileName]: + res = [] + for file_metadata in file_metadatas: + fname = file_metadata["name"] + try: + message_filename = FileChangeEventsMessageFileName.from_string(fname) + except ValueError: + print(f"Warning: invalid event file name: {fname}") + continue + res.append(message_filename) + return res + + +@dataclass +class DriveCachedClient: + """The Drive client of one account, and where syft's folders are in it. + + It logs in, finds and creates syft's folders, follows the version rules for + a versioned folder name, and caches the folder ids so a lookup costs no + Drive call the second time. The transport components share one. + """ + + client: GDriveClient + email: str + # /SyftBox: the top folder, with the inboxes, outboxes and personal folder. + _syftbox_folder_id: str | None = field(default=None, init=False) + # /SyftBox/{version}#{email}: the personal data. + _personal_syftbox_folder_id: str | None = field(default=None, init=False) + # peer email -> id of the folder I made for that peer's datasite. + peer_datasite_inbox_cache: dict[str, str] = field(default_factory=dict, init=False) + peer_datasite_outbox_cache: dict[str, str] = field(default_factory=dict, init=False) + # peer email -> id of the folder that peer made for my datasite. + own_datasite_inbox_cache: dict[str, str] = field(default_factory=dict, init=False) + own_datasite_outbox_cache: dict[str, str] = field(default_factory=dict, init=False) + # sender email -> id of my archive folder for that sender. + archive_folder_id_cache: dict[str, str] = field(default_factory=dict, init=False) + + @classmethod + def from_config(cls, config: "GdriveConnectionConfig") -> "DriveCachedClient": + return cls.from_token_path(config.email, config.token_path) + + @classmethod + def from_token_path( + cls, + email: str, + token_path: Path | None, + warm_syftbox_folder_id_cache: bool = True, + ) -> "DriveCachedClient": + """Log in with the OAuth token at ``token_path``, or with Colab's own + login when ``token_path`` is None.""" + credentials = ( + GoogleCredentials.from_authorized_user_file(token_path, SCOPES) + if token_path + else None + ) + new_service = partial(build_drive_service, credentials, environment=check_env()) + return cls._build(email, new_service, warm_syftbox_folder_id_cache) + + @classmethod + def from_service( + cls, + email: str, + mock_service: "MockDriveService", + warm_syftbox_folder_id_cache: bool = True, + ) -> "DriveCachedClient": + """A client for ``email`` over the backing store of ``mock_service``, + for tests.""" + from syft.sync.connections.drive.mock_drive_service import ( + MockDriveService, + ) + + new_service = partial(MockDriveService, mock_service._backing_store, email) + return cls._build(email, new_service, warm_syftbox_folder_id_cache) + + @classmethod + def _build( + cls, + email: str, + new_service: Callable[[], Any], + warm_syftbox_folder_id_cache: bool, + ) -> "DriveCachedClient": + """A client over a Drive service from ``new_service``, which also builds + the service of each copy for another thread. + + Args: + warm_syftbox_folder_id_cache: If True, eagerly resolve the personal + SyftBox folder id (2-3 API calls). Pass False when the caller + seeds the caches. + """ + cached_client = cls(GDriveClient(new_service(), new_service), email) + if warm_syftbox_folder_id_cache: + cached_client.get_personal_syftbox_folder_id() + return cached_client + + def clone_for_thread(self) -> "DriveCachedClient": + """A copy over a new Drive client, for use from another thread, that + knows the folder ids this one found.""" + clone = DriveCachedClient(self.client.clone_for_thread(), self.email) + self.copy_caches_to(clone) + return clone + + def copy_caches_to(self, other: "DriveCachedClient") -> None: + """Seed the folder ids of ``other`` from this one.""" + other._syftbox_folder_id = self._syftbox_folder_id + other._personal_syftbox_folder_id = self._personal_syftbox_folder_id + other.peer_datasite_inbox_cache = dict(self.peer_datasite_inbox_cache) + other.peer_datasite_outbox_cache = dict(self.peer_datasite_outbox_cache) + other.own_datasite_inbox_cache = dict(self.own_datasite_inbox_cache) + other.own_datasite_outbox_cache = dict(self.own_datasite_outbox_cache) + other.archive_folder_id_cache = dict(self.archive_folder_id_cache) + + def reset_caches(self) -> None: + self._syftbox_folder_id = None + self._personal_syftbox_folder_id = None + self.peer_datasite_inbox_cache.clear() + self.peer_datasite_outbox_cache.clear() + self.own_datasite_inbox_cache.clear() + self.own_datasite_outbox_cache.clear() + self.archive_folder_id_cache.clear() + + def get_syftbox_folder_id(self) -> str: + """/SyftBox""" + # cached + if self._syftbox_folder_id: + return self._syftbox_folder_id + else: + syftbox_folder_id = self.get_syftbox_folder_id_from_drive() + if syftbox_folder_id: + self._syftbox_folder_id = syftbox_folder_id + return self._syftbox_folder_id + else: + return self.create_syftbox_folder() + + def get_syftbox_folder_id_from_drive(self) -> str | None: + query = f"name='{SYFTBOX_FOLDER}' and mimeType='application/vnd.google-apps.folder' and 'me' in owners and trashed=false" + results = execute_with_retries( + self.client.drive_service.files().list(q=query, fields="files(id, name)") + ) + items = results.get("files", []) + return items[0]["id"] if items else None + + def create_syftbox_folder(self) -> str: + """Creates /SyftBox""" + folder = self.client.create_syft_file( + {"name": SYFTBOX_FOLDER, "mimeType": GOOGLE_FOLDER_MIME_TYPE} + ) + return folder["id"] + + def get_personal_syftbox_folder_id(self) -> str: + """/SyftBox/{version}#{email}""" + if self._personal_syftbox_folder_id: + return self._personal_syftbox_folder_id + folders = self.client.find_folders( + name_contains=[f"#{self.email}"], + parent_id=self.get_syftbox_folder_id(), + owner_email=self.email, + ) + # The substring '#{email}' also matches p2p folder names that end in + # '#{peer}#{type}#{email}'. Personal folder shape is exactly + # '{version}#{email}', so require a single '#'. + folders = [(fid, name) for fid, name in folders if name.count("#") == 1] + folder_id = self.find_or_adopt_versioned_folder( + folders, + current_name=GdrivePersonalSyftboxFolder(email=self.email).as_string(), + ) + if folder_id: + self._personal_syftbox_folder_id = folder_id + return folder_id + return self.create_personal_syftbox_folder() + + def create_personal_syftbox_folder(self) -> str: + """Creates /SyftBox/{version}#{email}""" + syftbox_folder_id = self.get_syftbox_folder_id() + folder_name = GdrivePersonalSyftboxFolder(email=self.email).as_string() + return self.client.create_folder(folder_name, syftbox_folder_id) + + def find_or_adopt_versioned_folder( + self, + folders: list[tuple[str, str]], + current_name: str, + current_version: str | None = None, + ) -> str | None: + """Return the id of a PRIVATE folder for this client version, or None. + + A private folder name holds the client version, so a minor upgrade looks + for a name that does not exist yet. This method renames the folder of the + highest earlier version to `current_name` and keeps the data. A new folder + would leave the data of the user on Drive and out of reach. + + Renames the folder, so the caller must own it and no peer may look it up by + name. A P2P folder fails both conditions: use `expect_one` for those. + + Raises RuntimeError if only a folder from a later version exists, or if + more than one compatible folder exists. + """ + compatible, older, newer = _partition_by_version(folders, current_version) + if compatible: + return self.expect_one(compatible) + if newer: + names = [n for _, n in newer] + latest = _extract_version_from_name(names[-1]) + raise RuntimeError( + f"Found a folder from a later client version on Drive: {names}. " + f"This client is {current_version or SYFT_VERSION} and " + f"cannot read that data. Install syft {latest} or later." + ) + if not older: + return None + + folder_id, name = older[-1] + execute_with_retries( + self.client.drive_service.files().update( + fileId=folder_id, body={"name": current_name} + ) + ) + print(f"Adopted the folder of an earlier version: {name} -> {current_name}") + if len(older) > 1: + stale = [n for _, n in older[:-1]] + print( + f"Warning: {len(stale)} folder(s) of earlier versions stay on " + f"Drive: {stale}" + ) + return folder_id + + def expect_one(self, folders: list[tuple[str, str]]) -> str | None: + """Return the single folder id, None if empty, raise if more than one. + + Auto-picking from multiple matches risks orphaning data, so the user + must reconcile on Drive themselves. + """ + if not folders: + return None + if len(folders) == 1: + return folders[0][0] + names = [n for _, n in folders] + raise RuntimeError( + f"Found {len(folders)} compatible folders on Drive: {names}. " + f"Exactly one is expected. This usually indicates a stale folder " + f"left behind by a prior client version. Please delete the stale " + f"folder(s) on Drive (keeping the one with your data) and retry." + ) + + def create_peer_datasite_folder(self, peer_email: str, folder_type: str) -> str: + """Create a datasite folder under /SyftBox for peer's datasite.""" + if folder_type not in ("inbox", "outbox"): + raise ValueError( + f"Invalid folder_type: {folder_type}. Must be 'inbox' or 'outbox'." + ) + parent_id = self.get_syftbox_folder_id() + folder = GdriveP2PFolder( + datasite_email=peer_email, folder_type=folder_type, peer_email=self.email + ) + folder_name = folder.as_string() + folder_id = self.client.create_folder(folder_name, parent_id) + if folder_type == "inbox": + self.peer_datasite_inbox_cache[peer_email] = folder_id + else: + self.peer_datasite_outbox_cache[peer_email] = folder_id + return folder_id + + def find_p2p_folder_id( + self, + datasite_email: str, + folder_type: str, + peer_email: str, + owner_email: str, + ) -> str | None: + folders = self.client.find_folders( + name_contains=[ + f"{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#", + f"#{datasite_email}#{folder_type}#{peer_email}", + ], + owner_email=owner_email, + ) + + # Drive's `name contains` is a fuzzy token/prefix match, not a substring + # match, so it can return folders for other peers when email tokens + # collide + def _is_exact_match(name: str) -> bool: + try: + folder = GdriveP2PFolder.from_name(name) + except ValueError: + return False + return ( + folder.datasite_email == datasite_email + and folder.folder_type == folder_type + and folder.peer_email == peer_email + ) + + # Ignore the version in the name. Each peer builds this name from its own + # client version, so a filter here hides the folder that the peer uses. + # After an upgrade the client therefore finds the old folder and writes to + # it. It makes no second folder, which an older peer would never look for. + candidates = _sorted_by_version( + [(fid, name) for fid, name in folders if _is_exact_match(name)] + ) + if not candidates: + return None + if len(candidates) > 1: + print( + f"Warning: {len(candidates)} P2P folders for {datasite_email} " + f"{folder_type} {peer_email}; using {candidates[-1][1]}" + ) + return candidates[-1][0] + + def get_peer_datasite_inbox_id(self, peer_email: str) -> str | None: + """Get folder: syft_datasite_{peer}_inbox_{self}, owned by self.""" + if peer_email in self.peer_datasite_inbox_cache: + return self.peer_datasite_inbox_cache[peer_email] + folder_id = self.find_p2p_folder_id( + datasite_email=peer_email, + folder_type="inbox", + peer_email=self.email, + owner_email=self.email, + ) + if folder_id is not None: + self.peer_datasite_inbox_cache[peer_email] = folder_id + return folder_id + + def get_peer_datasite_outbox_id(self, peer_email: str) -> str | None: + """Get folder: syft_datasite_{peer}_outbox_{self}, owned by self.""" + if peer_email in self.peer_datasite_outbox_cache: + return self.peer_datasite_outbox_cache[peer_email] + folder_id = self.find_p2p_folder_id( + datasite_email=peer_email, + folder_type="outbox", + peer_email=self.email, + owner_email=self.email, + ) + if folder_id is not None: + self.peer_datasite_outbox_cache[peer_email] = folder_id + return folder_id + + def get_own_datasite_inbox_id(self, peer_email: str) -> str | None: + """Get folder: syft_datasite_{self}_inbox_{peer}, owned by peer.""" + if peer_email in self.own_datasite_inbox_cache: + return self.own_datasite_inbox_cache[peer_email] + folder_id = self.find_p2p_folder_id( + datasite_email=self.email, + folder_type="inbox", + peer_email=peer_email, + owner_email=peer_email, + ) + if folder_id is not None: + self.own_datasite_inbox_cache[peer_email] = folder_id + return folder_id + + def get_own_datasite_outbox_id(self, peer_email: str) -> str | None: + """Get folder: syft_datasite_{self}_outbox_{peer}, owned by peer.""" + if peer_email in self.own_datasite_outbox_cache: + return self.own_datasite_outbox_cache[peer_email] + folder_id = self.find_p2p_folder_id( + datasite_email=self.email, + folder_type="outbox", + peer_email=peer_email, + owner_email=peer_email, + ) + if folder_id is not None: + self.own_datasite_outbox_cache[peer_email] = folder_id + return folder_id + + def get_archive_folder_id(self, sender_email: str) -> str: + if sender_email not in self.archive_folder_id_cache: + archive_folder_id = self.get_archive_folder_id_from_drive( + sender_email + ) or self.create_archive_folder(sender_email) + self.archive_folder_id_cache[sender_email] = archive_folder_id + return self.archive_folder_id_cache[sender_email] + + def get_archive_folder_id_from_drive(self, sender_email: str) -> str | None: + archive_folder = GdriveArchiveFolder( + sender_email=sender_email, recipient_email=self.email + ) + archive_folder_name = archive_folder.as_string() + query = f"name='{archive_folder_name}' and mimeType='application/vnd.google-apps.folder' and 'me' in owners and trashed=false" + results = execute_with_retries( + self.client.drive_service.files().list(q=query, fields="files(id)") + ) + items = results.get("files", []) + return items[0]["id"] if items else None + + def create_archive_folder(self, sender_email: str) -> str: + archive_folder = GdriveArchiveFolder( + sender_email=sender_email, recipient_email=self.email + ) + archive_folder_name = archive_folder.as_string() + syftbox_folder_id = self.get_syftbox_folder_id() + return self.client.create_folder(archive_folder_name, syftbox_folder_id) + + def get_file_metadatas_from_folder( + self, + folder_id: str, + since_timestamp: float | None = None, + page_size: int = 100, + ) -> list[dict]: + """ + Get file metadatas from folder with early termination. + + Args: + folder_id: Google Drive folder ID + since_timestamp: Optional timestamp. If provided, only files with a + newer timestamp are returned, and pagination stops at + a file from an earlier second. + page_size: Number of files to fetch per API call. Default 100. + + Returns: + List of file metadata dicts, sorted by name descending (newest first) + """ + query = f"'{folder_id}' in parents and trashed=false" + all_files = [] + page_token = None + + while True: + results = execute_with_retries( + self.client.drive_service.files().list( + q=query, + fields="files(id, name, size, mimeType, modifiedTime), nextPageToken", + pageSize=page_size, + pageToken=page_token, + orderBy="name desc", + ) + ) + + page_files = results.get("files", []) + + # Early termination: Check if this page contains old files + if since_timestamp is not None and page_files: + should_stop = False + + for file in page_files: + timestamp = timestamp_from_message_name(file["name"]) + + if timestamp is not None: + if timestamp > since_timestamp: + all_files.append(file) + elif int(timestamp) < int(since_timestamp): + # Drive sorts names as strings, so a 0.9 fraction of + # a second lists before 0.95. We stop when the + # second flips. + should_stop = True + break + else: + # No timestamp in filename, include the file + all_files.append(file) + + if should_stop: + # Don't fetch more pages + break + else: + # No early termination check, add all files + all_files.extend(page_files) + + # Check for next page + page_token = results.get("nextPageToken") + if not page_token: + break + + return all_files diff --git a/syft/sync/connections/drive/drive_collections.py b/syft/sync/connections/drive/drive_collections.py new file mode 100644 index 00000000000..2919b2f00ef --- /dev/null +++ b/syft/sync/connections/drive/drive_collections.py @@ -0,0 +1,336 @@ +"""Collections on Google Drive: folders in /SyftBox shared by Drive permission. + +A collection folder is named ``{prefix}{variant}_{tag}_{content_hash}``. +""" + +import re +from functools import lru_cache +from pathlib import Path +from typing import Dict, List + +from pydantic import BaseModel + +from syft.sync.connections.base_connection import FileCollection +from syft.sync.connections.drive.gdrive_client import GOOGLE_FOLDER_MIME_TYPE +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient +from syft.sync.connections.drive.gdrive_retry import execute_with_retries +from syft.sync.transport.protocols import Collections + + +def collection_name_query(prefix: str) -> str: + """A Drive query that finds a collection of ``prefix`` in any layout. + + It has no trailing '_', because a layout writes its variant in that + position. A client that predates multi-layout searches with the '_', and so + never lists a layout it cannot read. + """ + return f"name contains '{prefix}'" + + +@lru_cache(maxsize=None) +def _collection_name_re(prefix: str) -> "re.Pattern[str]": + """Matches '{prefix}{variant}_{tag}_{hash}'. + + The variant and the hash hold no '_'. The tag can hold one, so it takes + every character in between. + """ + return re.compile( + rf"^{re.escape(prefix)}" + r"(?P[^_]*)_(?P.+)_(?P[^_]+)$" + ) + + +class CollectionFolder(BaseModel): + """Naming value object for collection folders (domain-agnostic). + + Collections (datasets, private datasets, or any future collection type) are + named ``{prefix}{variant}_{tag}_{content_hash}`` on the wire. The variant + names the layout (see CollectionSyncSpec) and is "" for the original one, so + that name is unchanged byte for byte. This is the single source of truth for + building and parsing the name; it also computes the content hash from the + folder's file contents. + + A writer normally folds the variant into ``prefix`` (the spec builds the + qualified prefix once) and leaves ``variant`` empty; ``from_name`` splits + them apart again for a reader. + """ + + prefix: str + tag: str + content_hash: str + variant: str = "" + + @property + def folder_name(self) -> str: + """The on-wire folder name. Doubles as the cache key (same string).""" + return f"{self.prefix}{self.variant}_{self.tag}_{self.content_hash}" + + @classmethod + def from_name(cls, prefix: str, name: str) -> "CollectionFolder": + """Parse '{prefix}{variant}_{tag}_{hash}' into a CollectionFolder. + + Raises ValueError if the name does not carry the prefix, so callers can + skip non-matching folders. + """ + match = _collection_name_re(prefix).match(name) + if match is None: + raise ValueError(f"Invalid collection folder name: {name}") + return cls(prefix=prefix, **match.groupdict()) + + @staticmethod + def compute_hash(files: dict[str, bytes]) -> str: + """Compute a content hash from file contents.""" + from syft.sync.file_utils import compute_file_hashes + + return compute_file_hashes(files) + + +class DriveCollections(Collections): + """Collection folders in /SyftBox, shared by Drive permission.""" + + def __init__(self, cached_client: DriveCachedClient): + self._cached_client = cached_client + # folder name -> id of my collection folder + self.collection_folder_id_cache: dict[str, str] = {} + + def clone_for_thread(self) -> "DriveCollections": + """Collections over a new Drive client, for use from another thread, + with these caches.""" + clone = DriveCollections(self._cached_client.clone_for_thread()) + clone.collection_folder_id_cache = dict(self.collection_folder_id_cache) + return clone + + def reset_caches(self) -> None: + self._cached_client.reset_caches() + self.collection_folder_id_cache.clear() + + def create(self, prefix: str, tag: str, content_hash: str, owner_email: str) -> str: + """Create /SyftBox/{prefix}_{tag}_{hash} folder.""" + folder_name = CollectionFolder( + prefix=prefix, tag=tag, content_hash=content_hash + ).folder_name + + # Check cache + if folder_name in self.collection_folder_id_cache: + return self.collection_folder_id_cache[folder_name] + + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + + # Check if exists + folder_id = self._cached_client.client.find_folder_by_name( + folder_name, parent_id=syftbox_folder_id + ) + if folder_id: + self.collection_folder_id_cache[folder_name] = folder_id + return folder_id + + # Create new folder + folder_id = self._cached_client.client.create_folder( + folder_name, syftbox_folder_id + ) + self.collection_folder_id_cache[folder_name] = folder_id + return folder_id + + def mark_shared_with_any(self, prefix: str, tag: str, content_hash: str) -> None: + """Mark collection as shared with 'any' via appProperties.""" + folder_id = self._get_collection_folder_id(prefix, tag, content_hash) + execute_with_retries( + self._cached_client.client.drive_service.files().update( + fileId=folder_id, + body={"appProperties": {"syft_shared_with_any": "true"}}, + ) + ) + + def grant_read( + self, prefix: str, tag: str, content_hash: str, users: list[str] + ) -> None: + """Share collection folder with specific users via batch API.""" + if not users: + return + folder_id = self._get_collection_folder_id(prefix, tag, content_hash) + self._cached_client.client.batch_add_permissions(folder_id, users) + + def upload( + self, prefix: str, tag: str, content_hash: str, files: dict[str, bytes] + ) -> None: + """Upload files to a collection folder.""" + folder_id = self._get_collection_folder_id(prefix, tag, content_hash) + + for file_path, content in files.items(): + file_payload, _ = self._cached_client.client.create_file_payload(content) + file_name = Path(file_path).name + + file_metadata = {"name": file_name, "parents": [folder_id]} + self._cached_client.client.create_syft_file(file_metadata, file_payload) + + def list_mine(self, prefix: str) -> list[str]: + """List collections created by DO (owned by me). Returns list of tags. + + A tag appears once however many layouts hold it. + """ + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + query = ( + f"{collection_name_query(prefix)} and '{syftbox_folder_id}' in parents " + f"and 'me' in owners and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" + ) + result: list[str] = [] + for folder in self._cached_client.client.list_all_files(query, fields="name"): + try: + cf = CollectionFolder.from_name(prefix, folder["name"]) + except ValueError: + continue + if cf.tag not in result: + result.append(cf.tag) + return result + + def list_mine_with_permissions( + self, + prefix: str, + ) -> list[FileCollection]: + """List all DO's collections with permissions info, one per layout.""" + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + query = ( + f"{collection_name_query(prefix)} and '{syftbox_folder_id}' in parents " + f"and 'me' in owners and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" + ) + collections = [] + for folder in self._cached_client.client.list_all_files( + query, fields="id,name,appProperties" + ): + try: + cf = CollectionFolder.from_name(prefix, folder["name"]) + except ValueError: + continue + has_anyone = ( + folder.get("appProperties", {}).get("syft_shared_with_any") == "true" + ) + collections.append( + FileCollection( + folder_id=folder["id"], + tag=cf.tag, + content_hash=cf.content_hash, + has_any_permission=has_anyone, + variant=cf.variant, + ) + ) + + return collections + + def delete(self, prefix: str, tag: str) -> None: + """Delete every layout of ``tag`` published under ``prefix``.""" + collections = self.list_mine_with_permissions(prefix) + for c in collections: + if c.tag == tag: + self._cached_client.client.delete_file_by_id(c.folder_id) + folder_name = CollectionFolder( + prefix=prefix, + tag=c.tag, + content_hash=c.content_hash, + variant=c.variant, + ).folder_name + self.collection_folder_id_cache.pop(folder_name, None) + + def delete_by_handle(self, handle: str) -> None: + self._cached_client.client.delete_file_by_id(handle) + self.collection_folder_id_cache = { + name: folder_id + for name, folder_id in self.collection_folder_id_cache.items() + if folder_id != handle + } + + def list_shared_with_me(self, prefix: str) -> list[dict]: + """List collections shared with DS (not owned by me), one per layout. + + Returns list of dicts with keys: owner_email, tag, content_hash, variant + """ + query = ( + f"{collection_name_query(prefix)} and not 'me' in owners " + f"and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" + ) + result = [] + for folder in self._cached_client.client.list_all_files( + query, fields="name, owners" + ): + try: + cf = CollectionFolder.from_name(prefix, folder["name"]) + except ValueError: + # Skip folders that don't match the expected format + continue + owner_email = folder.get("owners", [{}])[0].get("emailAddress", "unknown") + result.append( + { + "owner_email": owner_email, + "tag": cf.tag, + "content_hash": cf.content_hash, + "variant": cf.variant, + } + ) + return result + + def download( + self, prefix: str, tag: str, content_hash: str, owner_email: str + ) -> dict[str, bytes]: + """Download all files from a collection.""" + folder_name = CollectionFolder( + prefix=prefix, tag=tag, content_hash=content_hash + ).folder_name + # Try to find folder by name (could be owned by someone else) + folder_id = self._cached_client.client.find_folder_by_name( + folder_name, owner_email=owner_email + ) + + if not folder_id: + raise ValueError(f"Collection {tag} with hash {content_hash} not found") + + file_metadatas = self._cached_client.get_file_metadatas_from_folder(folder_id) + files = {} + for file_meta in file_metadatas: + file_id = file_meta["id"] + file_name = file_meta["name"] + files[file_name] = self._cached_client.client.download_file(file_id) + + return files + + def list_files( + self, prefix: str, tag: str, content_hash: str, owner_email: str + ) -> List[Dict]: + """Get file metadata from a collection without downloading.""" + folder_name = CollectionFolder( + prefix=prefix, tag=tag, content_hash=content_hash + ).folder_name + folder_id = self._cached_client.client.find_folder_by_name( + folder_name, owner_email=owner_email + ) + + if not folder_id: + raise ValueError(f"Collection {tag} with hash {content_hash} not found") + + file_metadatas = self._cached_client.get_file_metadatas_from_folder(folder_id) + return [{"file_id": f["id"], "file_name": f["name"]} for f in file_metadatas] + + def fetch_file(self, handle: str) -> bytes: + """Download a single file from a collection.""" + return self._cached_client.client.download_file(handle) + + def _get_collection_folder_id( + self, prefix: str, tag: str, content_hash: str + ) -> str: + """Get folder ID for a collection, with caching.""" + folder_name = CollectionFolder( + prefix=prefix, tag=tag, content_hash=content_hash + ).folder_name + if folder_name in self.collection_folder_id_cache: + return self.collection_folder_id_cache[folder_name] + + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + folder_id = self._cached_client.client.find_folder_by_name( + folder_name, parent_id=syftbox_folder_id + ) + + if not folder_id: + raise ValueError( + f"Collection folder {tag} with hash {content_hash} not found" + ) + + self.collection_folder_id_cache[folder_name] = folder_id + return folder_id diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 75118f968e5..93ec9b1051e 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -1,259 +1,18 @@ -"""The transport jobs on Google Drive. +"""The transport components on Google Drive. -Each class wraps a ``GDriveConnection`` and calls the method that does the job -today, so what is read from and written to Drive does not change. Where a -contract builds a method from simpler ones, the Drive class overrides it to -keep the single Drive call it makes today. The Drive code moves out of -``GDriveConnection`` into these classes, over a shared ``GDriveClient``, in a -later step. +Each component has its own module: ``drive_p2p``, ``drive_owner_store``, +``drive_collections`` and ``drive_peering``. They share one ``DriveCachedClient``, +which holds the Drive client. ``SyftboxManager._drive_transport`` builds them. """ -from __future__ import annotations +from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient +from syft.sync.transport.protocols import Transport, WatcherStore -import threading -import warnings -from concurrent.futures import ThreadPoolExecutor -from typing import TYPE_CHECKING, Self -from syft.sync.connections.drive.gdrive_retry import raise_unavailable -from syft.sync.transport.exceptions import TransportUnavailableError -from syft.sync.transport.protocols import ( - Collections, - OwnerStore, - P2PConnection, - Peering, - Transport, - WatcherStore, -) -from syft.sync.transport.types import MessageKind, RawMessage, SendStatus - -if TYPE_CHECKING: - from syft.sync.connections.base_connection import FileCollection - from syft.sync.connections.drive.gdrive_transport import GDriveConnection - from syft.sync.version.peer_info import PeerInfo - - -# Threads that download the messages of one receive at the same time. -DOWNLOAD_WORKERS = 10 - - -def _download_or_skip(connection: GDriveConnection, listed: dict) -> bytes | None: - """The bytes of one listed inbox file, or None with a warning when this - file cannot be downloaded, for example because it was deleted. - - Raises: - TransportUnavailableError: when Drive cannot be reached. - """ - name = listed["file_name"] - try: - with raise_unavailable(f"Drive is unreachable; inbox message {name} waits"): - return connection.download_file(listed["file_id"]) - except TransportUnavailableError: - raise - except Exception as e: - warnings.warn(f"Skipped inbox message {name}: download failed: {e!r}") - return None - - -def _require_inbox_kind(kind: MessageKind) -> None: - """Only the inbox is read through ``receive`` so far.""" - if kind is not MessageKind.PROPOSED_CHANGES: - raise NotImplementedError( - f"Messages of kind {kind!r} are still read through the outbox " - "methods; they move to receive with send-and-forget." - ) - - -class DriveTransportComponent: - """A transport component over one Google Drive connection.""" - - def __init__(self, connection: GDriveConnection): - self._conn = connection - - def clone_for_thread(self) -> Self: - """A copy safe to use from another thread: the Drive client is not.""" - return type(self)(self._conn.clone_for_thread()) - - def reset_caches(self) -> None: - self._conn.reset_caches() - - -class DriveP2PConnection(DriveTransportComponent, P2PConnection): - """Messages through the shared inbox and outbox folders. - - A proposed change goes to the data owner's inbox and file events go to the - data scientist's outbox. ``msg_id`` is the file name. A written file is - ``DELIVERED``: it waits in the folder until the peer reads it. - - The inbox is read through ``receive`` and a processed message is archived. - The outbox is still read through the Drive-only methods below, from a - timestamp cursor, and kept: the data scientist rebuilds a lost cache by - reading it again. It moves to ``receive`` with send-and-forget, once the - data scientist keeps what it received in its own storage. - """ - - def send( - self, peer: str, msg_id: str, data: bytes, kind: MessageKind - ) -> SendStatus: - if kind is MessageKind.PROPOSED_CHANGES: - self._conn.watcher_send_raw_bytes_to_inbox(peer, msg_id, data) - elif kind is MessageKind.FILE_EVENTS: - self._conn.owner_write_raw_bytes_to_outbox(peer, msg_id, data) - else: - raise ValueError(f"Drive has no folder for messages of kind {kind!r}") - return SendStatus.DELIVERED - - def delivered(self) -> list[str]: - # Every Drive send is DELIVERED when it returns; none waits for a peer. - return [] - - def receive( - self, peer: str, kind: MessageKind, max_messages: int | None = None - ) -> list[RawMessage]: - """One listing for the whole inbox, then the files download in parallel. - - A file that cannot be downloaded is left out with a warning, so it does - not hold back the others. - - Raises: - TransportUnavailableError: when Drive cannot be reached. Every - message stays in the inbox. - """ - _require_inbox_kind(kind) - listed = self._conn.owner_get_inbox_file_metadatas(peer)[:max_messages] - contents = self._download_all(listed) - return [ - RawMessage( - sender=peer, - msg_id=m["file_name"], - kind=kind, - data=data, - handle=m["file_id"], - ) - for m, data in zip(listed, contents) - if data is not None - ] - - def _download_all(self, listed: list[dict]) -> list[bytes | None]: - """The bytes of each listed file, in order; None for a file that could - not be downloaded. Several files download over a thread pool with one - connection copy per thread, because the Drive client is not - thread-safe.""" - if len(listed) <= 1: - return [_download_or_skip(self._conn, m) for m in listed] - per_thread = threading.local() - - def download(m: dict) -> bytes | None: - if not hasattr(per_thread, "conn"): - per_thread.conn = self._conn.clone_for_thread() - return _download_or_skip(per_thread.conn, m) - - workers = min(DOWNLOAD_WORKERS, len(listed)) - with ThreadPoolExecutor(max_workers=workers) as pool: - return list(pool.map(download, listed)) - - def mark_processed(self, message: RawMessage) -> None: - # The archive folder is named after ``message.sender``, which receive - # set from the inbox folder, never from the message body. - _require_inbox_kind(message.kind) - self._conn.owner_archive_inbox_file(message.handle, message.sender) - - def mark_all_processed(self, messages: list[RawMessage]) -> None: - for message in messages: - _require_inbox_kind(message.kind) - self._conn.owner_archive_inbox_files([(m.handle, m.sender) for m in messages]) - - def outbox_messages(self, peer: str, since_timestamp: float | None) -> list[bytes]: - """The messages in ``peer``'s outbox to me, newer than ``since_timestamp``.""" - return self._conn.watcher_download_raw_events_from_outbox(peer, since_timestamp) - - def outbox_message_metadatas( - self, peer: str, since_timestamp: float | None - ) -> list[dict]: - """Like ``outbox_messages``, without downloading: each entry's - ``file_id`` is the handle to ``fetch_message`` it with.""" - return self._conn.watcher_get_outbox_file_metadatas(peer, since_timestamp) - - def fetch_message(self, handle: str) -> bytes: - return self._conn.download_file(handle) - - def own_outbox_message_metadatas(self, peer: str) -> list[dict]: - """The messages in my outbox to ``peer``, without downloading. An - outbox whose folder is not found holds none.""" - folder_id = self._conn._get_own_datasite_outbox_id(peer) - if folder_id is None: - return [] - return self._conn.get_file_metadatas_from_folder( - folder_id, since_timestamp=None - ) - - def delete_messages(self, handles: list[str]) -> None: - self._conn.delete_multiple_files_by_ids(handles) - - -class DriveOwnerStore(DriveTransportComponent, OwnerStore): - """The data owner's own folders in /SyftBox: event log, checkpoints, - rolling state.""" - - def append_event(self, name: str, data: bytes) -> str: - return self._conn.owner_write_raw_bytes_to_syftbox(name, data) - - def list_event_handles(self, since_timestamp: float | None = None) -> list[str]: - return self._conn.owner_get_all_accepted_event_file_ids(since_timestamp) - - def fetch(self, handle: str) -> bytes: - return self._conn.owner_download_raw_bytes_by_id(handle) - - def fetch_events_since(self, since_timestamp: float | None) -> list[bytes]: - if since_timestamp is None: - return self.fetch_all_events() - return self._conn.download_raw_events_since_timestamp(since_timestamp) - - def fetch_all_events(self) -> list[bytes]: - return self._conn.owner_download_all_raw_events_from_syftbox() - - def put_checkpoint(self, name: str, data: bytes) -> str: - return self._conn.upload_raw_checkpoint(name, data) - - def latest_checkpoint(self) -> bytes | None: - return self._conn.download_raw_latest_checkpoint() - - def put_incremental(self, name: str, data: bytes) -> str: - return self._conn.upload_raw_incremental_checkpoint(name, data) - - def all_incrementals(self) -> list[bytes]: - return self._conn.download_all_raw_incremental_checkpoints() - - def incremental_count(self) -> int: - return self._conn.get_incremental_checkpoint_count() - - def next_incremental_sequence(self) -> int: - return self._conn.get_next_incremental_sequence_number() - - def delete_incrementals(self) -> None: - self._conn.delete_all_incremental_checkpoints() - - def put_rolling_state(self, name: str, data: bytes) -> str: - return self._conn.upload_raw_rolling_state(name, data) - - def get_rolling_state(self) -> bytes | None: - return self._conn.download_raw_rolling_state() - - def delete_rolling_state(self) -> None: - self._conn.delete_rolling_state() - - def delete_all(self) -> int: - """Everything in /SyftBox, and the syft files Drive moved to the root - of My Drive when they lost their folder.""" - file_ids = set(self._conn.gather_all_file_and_folder_ids()) - file_ids |= set(self._conn.find_orphaned_syft_files()) - self._conn.delete_multiple_files_by_ids(list(file_ids)) - return len(file_ids) - - -class DriveWatcherStore(DriveTransportComponent, WatcherStore): - """The data scientist's storage on Drive. Not built yet: today the data - scientist keeps what it received in a local cache.""" +class DriveWatcherStore(WatcherStore): + """The DS's storage on Drive. Not built yet: today the DS keeps what it + received in a local cache.""" def put_checkpoint(self, peer: str, name: str, data: bytes) -> str: raise NotImplementedError("Watcher storage on Drive is not built yet") @@ -271,111 +30,11 @@ def reset(self, peer: str) -> None: raise NotImplementedError("Watcher storage on Drive is not built yet") -class DriveCollections(DriveTransportComponent, Collections): - """Collection folders in /SyftBox, shared by Drive permission.""" - - def create(self, prefix: str, tag: str, content_hash: str, owner_email: str) -> str: - return self._conn.owner_create_collection_folder( - prefix, tag, content_hash, owner_email - ) - - def mark_shared_with_any(self, prefix: str, tag: str, content_hash: str) -> None: - self._conn.owner_tag_collection_as_any(prefix, tag, content_hash) - - def grant_read( - self, prefix: str, tag: str, content_hash: str, users: list[str] - ) -> None: - self._conn.owner_share_collection(prefix, tag, content_hash, users) - - def upload( - self, prefix: str, tag: str, content_hash: str, files: dict[str, bytes] - ) -> None: - self._conn.owner_upload_collection_files(prefix, tag, content_hash, files) - - def list_mine_with_permissions(self, prefix: str) -> list[FileCollection]: - return self._conn.owner_list_all_collections_with_permissions(prefix) - - def list_mine(self, prefix: str) -> list[str]: - return self._conn.owner_list_collections(prefix) - - def delete(self, prefix: str, tag: str) -> None: - self._conn.owner_delete_collection(prefix, tag) - - def delete_by_handle(self, handle: str) -> None: - self._conn.delete_file_by_id(handle) - - def list_shared_with_me(self, prefix: str) -> list[dict]: - return self._conn.watcher_list_collections(prefix) - - def list_files( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> list[dict]: - return self._conn.watcher_get_collection_file_metadatas( - prefix, tag, content_hash, owner_email - ) - - def fetch_file(self, handle: str) -> bytes: - return self._conn.watcher_download_collection_file(handle) - - def download( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> dict[str, bytes]: - return self._conn.watcher_download_collection( - prefix, tag, content_hash, owner_email - ) - - -class DrivePeering(DriveTransportComponent, Peering): - """Identity from the Google account; channels, the peer list, key bundles - and version files in /SyftBox.""" - - def whoami(self) -> str: - return self._conn.get_authenticated_email() - - def open_channels(self, peer: str) -> None: - self._conn.add_peer(peer_email=peer) - - def incoming_requests(self) -> list[str]: - return self._conn.get_peer_requests() - - def read_peer_registry(self, force_download: bool = False) -> dict: - return self._conn._get_peers_json(force_download=force_download) - - def update_peer_registry( - self, peer: str, state: str, public_encryption_bundle: dict | None = None - ) -> None: - self._conn._update_peer_state(peer, state, public_encryption_bundle) - - def publish_key_bundle(self, peer: str, bundle_json: str) -> None: - self._conn.write_encryption_bundle(peer, bundle_json) - - def grant_key_bundle_read(self, peer: str) -> None: - self._conn.share_encryption_bundles_folder(peer) - - def read_own_key_bundle(self, peer: str) -> str | None: - return self._conn.read_own_encryption_bundle(peer) - - def read_peer_key_bundle(self, peer: str) -> str | None: - return self._conn.read_peer_encryption_bundle(peer) - - def publish_peer_info(self, peer_info: PeerInfo) -> None: - self._conn.write_peer_info_file(peer_info) - - def read_own_peer_info(self) -> PeerInfo | None: - return self._conn.read_own_peer_info_file() - - def read_peer_info(self, peer: str) -> PeerInfo | None: - return self._conn.read_peer_info_file(peer) - - def grant_peer_info_read(self, peer: str, peer_info: PeerInfo) -> None: - self._conn.share_peer_info_file_with_peer(peer, peer_info) - - -def gdrive_connection_of(transport: Transport) -> GDriveConnection: - """The Google Drive connection under ``transport``, for tests and tools - that need Drive itself. The sync code reaches Drive only through the - components. Raises if the transport is not built on Drive.""" +def drive_cached_client_of(transport: Transport) -> DriveCachedClient: + """The ``DriveCachedClient`` under ``transport``, for tests and tools that + need Drive itself. The sync code reaches Drive only through the components. + Raises if the transport is not built on Drive.""" store = transport.owner_store - if not isinstance(store, DriveTransportComponent): + if not isinstance(store, DriveOwnerStore): raise TypeError("This transport is not built on Google Drive") - return store._conn + return store._cached_client diff --git a/syft/sync/connections/drive/drive_owner_store.py b/syft/sync/connections/drive/drive_owner_store.py new file mode 100644 index 00000000000..b0e20170951 --- /dev/null +++ b/syft/sync/connections/drive/drive_owner_store.py @@ -0,0 +1,571 @@ +"""The DO's own storage on Google Drive. + +The event log is the personal folder /SyftBox/{version}#{email}. Checkpoints +and the rolling state have folders of their own under /SyftBox. +""" + +from syft.sync.checkpoints.checkpoint import ( + CHECKPOINT_FILENAME_PREFIX, + INCREMENTAL_CHECKPOINT_PREFIX, + Checkpoint, + IncrementalCheckpoint, +) +from syft.sync.checkpoints.rolling_state import ( + ROLLING_STATE_FILENAME_PREFIX, + RollingState, +) +from syft.sync.connections.drive.gdrive_client import ( + SYFT_FILE_PROPERTY_KEY, + SYFT_FILE_PROPERTY_VALUE, +) +from syft.sync.connections.drive.drive_cached_client import ( + SYFTBOX_FOLDER, + DriveCachedClient, + valid_events_from_file_metadatas, +) +from syft.sync.connections.drive.gdrive_retry import ( + execute_with_retries, + raise_unavailable, +) +from syft.sync.connections.drive.gdrive_utils import ( + gather_all_file_and_folder_ids_recursive, +) +from syft.sync.events.file_change_event import FileChangeEventsMessageFileName +from syft.sync.transport.protocols import OwnerStore +from syft.version import SYFT_VERSION + + +def _filter_valid_file_metadatas( + file_metadatas: list[dict], +) -> list[dict]: + res = [] + for file_metadata in file_metadatas: + fname = file_metadata["name"] + try: + FileChangeEventsMessageFileName.from_string(fname) + except ValueError: + # The folder holds other files, so a name that is not an event + # name is normal here. This method filters them out. + continue + res.append(file_metadata) + return res + + +class DriveOwnerStore(OwnerStore): + """The DO's own folders in /SyftBox: event log, checkpoints, rolling state.""" + + def __init__(self, cached_client: DriveCachedClient): + self._cached_client = cached_client + # Rolling state caches for single-API-call optimization + self._rolling_state_folder_id: str | None = None + self._rolling_state_file_id: str | None = None + # fname -> gdrive id + self.personal_syftbox_event_id_cache: dict[str, str] = {} + + def clone_for_thread(self) -> "DriveOwnerStore": + """A store over a new Drive client, for use from another thread, with + this store's caches.""" + clone = DriveOwnerStore(self._cached_client.clone_for_thread()) + clone._rolling_state_folder_id = self._rolling_state_folder_id + clone._rolling_state_file_id = self._rolling_state_file_id + clone.personal_syftbox_event_id_cache = dict( + self.personal_syftbox_event_id_cache + ) + return clone + + def reset_caches(self) -> None: + self._cached_client.reset_caches() + self._rolling_state_folder_id = None + self._rolling_state_file_id = None + self.personal_syftbox_event_id_cache.clear() + + def append_event(self, name: str, data: bytes) -> str: + """Write raw bytes to /SyftBox/myemail.""" + personal_syftbox_folder_id = ( + self._cached_client.get_personal_syftbox_folder_id() + ) + file_metadata = { + "name": name, + "parents": [personal_syftbox_folder_id], + } + file_payload, _ = self._cached_client.client.create_file_payload(data) + + res = self._cached_client.client.create_syft_file(file_metadata, file_payload) + gdrive_id = res.get("id") + self.personal_syftbox_event_id_cache[name] = gdrive_id + return gdrive_id + + def fetch(self, handle: str) -> bytes: + """Download raw bytes by file ID.""" + return self._cached_client.client.download_file(handle) + + def list_event_handles(self, since_timestamp: float | None = None) -> list[str]: + personal_syftbox_folder_id = ( + self._cached_client.get_personal_syftbox_folder_id() + ) + file_metadatas = self._cached_client.get_file_metadatas_from_folder( + personal_syftbox_folder_id, since_timestamp=since_timestamp + ) + valid_fname_objs = _filter_valid_file_metadatas(file_metadatas) + oldest_first = sorted( + valid_fname_objs, + key=lambda f: ( + FileChangeEventsMessageFileName.from_string(f["name"]).timestamp + ), + ) + return [f["id"] for f in oldest_first] + + @raise_unavailable("Could not read the event log from Drive.") + def fetch_all_events(self) -> list[bytes]: + """Download all event files from /SyftBox/myemail as raw bytes.""" + personal_syftbox_folder_id = ( + self._cached_client.get_personal_syftbox_folder_id() + ) + file_metadatas = self._cached_client.get_file_metadatas_from_folder( + personal_syftbox_folder_id + ) + valid_fname_objs = valid_events_from_file_metadatas(file_metadatas) + + result = [] + for fname_obj in sorted(valid_fname_objs, key=lambda x: x.timestamp): + gdrive_id = [ + f for f in file_metadatas if f["name"] == fname_obj.as_string() + ][0]["id"] + result.append(self._cached_client.client.download_file(gdrive_id)) + return result + + @raise_unavailable("Could not read the event log from Drive.") + def fetch_events_since(self, since_timestamp: float | None) -> list[bytes]: + """Download event files created after a timestamp as raw bytes.""" + if since_timestamp is None: + return self.fetch_all_events() + personal_folder_id = self._cached_client.get_personal_syftbox_folder_id() + file_metadatas = self._cached_client.get_file_metadatas_from_folder( + personal_folder_id, since_timestamp=since_timestamp + ) + + valid_fname_objs = valid_events_from_file_metadatas(file_metadatas) + + result = [] + for fname_obj in sorted(valid_fname_objs, key=lambda x: x.timestamp): + if fname_obj.timestamp <= since_timestamp: + continue + + gdrive_id = [ + f for f in file_metadatas if f["name"] == fname_obj.as_string() + ][0]["id"] + + result.append(self._cached_client.client.download_file(gdrive_id)) + + return result + + def _get_checkpoints_folder_name(self) -> str: + """Get the checkpoints folder name: {email}-{version}-checkpoints""" + return f"{self._cached_client.email}-{SYFT_VERSION}-checkpoints" + + def _get_checkpoints_folder_id(self) -> str | None: + """Find the checkpoints folder ID from Google Drive.""" + folders = self._cached_client.client.find_folders( + name_contains=[f"{self._cached_client.email}-", "-checkpoints"], + parent_id=self._cached_client.get_syftbox_folder_id(), + ) + return self._cached_client.find_or_adopt_versioned_folder( + folders, current_name=self._get_checkpoints_folder_name() + ) + + def _get_or_create_checkpoints_folder_id(self) -> str: + """Get or create the checkpoints folder.""" + folder_id = self._get_checkpoints_folder_id() + if folder_id is not None: + return folder_id + # Create the folder + folder_name = self._get_checkpoints_folder_name() + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + return self._cached_client.client.create_folder(folder_name, syftbox_folder_id) + + def put_checkpoint(self, name: str, data: bytes) -> str: + """Upload raw checkpoint bytes to Google Drive. + + Uploads the new checkpoint first, then deletes old ones. + """ + folder_id = self._get_or_create_checkpoints_folder_id() + payload, _ = self._cached_client.client.create_file_payload(data) + + file_metadata = { + "name": name, + "parents": [folder_id], + } + + result = self._cached_client.client.create_syft_file(file_metadata, payload) + + # Only delete old checkpoints after successful upload + self._delete_all_checkpoints(exclude_file_id=result.get("id")) + + return result.get("id") + + @raise_unavailable("Could not read the latest checkpoint from Drive.") + def latest_checkpoint(self) -> bytes | None: + """Download the latest checkpoint as raw bytes, or None.""" + folder_id = self._get_checkpoints_folder_id() + if folder_id is None: + return None + + query = ( + f"'{folder_id}' in parents and trashed=false " + f"and name contains '{CHECKPOINT_FILENAME_PREFIX}'" + ) + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id, name)" + ) + ) + items = results.get("files", []) + + if not items: + return None + + latest_file = None + latest_timestamp = -1.0 + for item in items: + timestamp = Checkpoint.filename_to_timestamp(item["name"]) + if timestamp is not None and timestamp > latest_timestamp: + latest_timestamp = timestamp + latest_file = item + + if latest_file is None: + return None + + return self._cached_client.client.download_file(latest_file["id"]) + + def _delete_all_checkpoints(self, exclude_file_id: str | None = None): + """Delete all existing full checkpoints (not incremental ones). + + Args: + exclude_file_id: If provided, skip deleting this file ID + (used to preserve a newly uploaded checkpoint). + """ + folder_id = self._get_checkpoints_folder_id() + if folder_id is None: + return + + # List only full checkpoint files (start with "checkpoint_" not "incremental_checkpoint_") + query = ( + f"'{folder_id}' in parents and trashed=false " + f"and name contains '{CHECKPOINT_FILENAME_PREFIX}'" + ) + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id, name)" + ) + ) + items = results.get("files", []) + + # Delete only full checkpoints (not incremental ones) + for item in items: + if item["name"].startswith(INCREMENTAL_CHECKPOINT_PREFIX): + continue # Skip incremental checkpoints + if item["id"] == exclude_file_id: + continue # Skip the newly uploaded checkpoint + try: + execute_with_retries( + self._cached_client.client.drive_service.files().delete( + fileId=item["id"] + ) + ) + except Exception as e: + print(f"Warning: Failed to delete checkpoint {item['name']}: {e}") + + def put_incremental(self, name: str, data: bytes) -> str: + """Upload raw incremental checkpoint bytes to Google Drive.""" + folder_id = self._get_or_create_checkpoints_folder_id() + payload, _ = self._cached_client.client.create_file_payload(data) + + file_metadata = { + "name": name, + "parents": [folder_id], + } + + result = self._cached_client.client.create_syft_file(file_metadata, payload) + return result.get("id") + + @raise_unavailable("Could not read the incremental checkpoints from Drive.") + def all_incrementals(self) -> list[bytes]: + """Download all incremental checkpoints as raw bytes.""" + folder_id = self._get_checkpoints_folder_id() + if folder_id is None: + return [] + + query = ( + f"'{folder_id}' in parents and trashed=false " + f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" + ) + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id, name)" + ) + ) + items = results.get("files", []) + + if not items: + return [] + + return [self._cached_client.client.download_file(item["id"]) for item in items] + + def incremental_count(self) -> int: + """Get the number of incremental checkpoints on Google Drive.""" + folder_id = self._get_checkpoints_folder_id() + if folder_id is None: + return 0 + + query = ( + f"'{folder_id}' in parents and trashed=false " + f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" + ) + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) + ) + return len(results.get("files", [])) + + def next_incremental_sequence(self) -> int: + """Get the next sequence number for incremental checkpoints.""" + folder_id = self._get_checkpoints_folder_id() + if folder_id is None: + return 1 + + query = ( + f"'{folder_id}' in parents and trashed=false " + f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" + ) + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(name)" + ) + ) + items = results.get("files", []) + + if not items: + return 1 + + # Find highest sequence number + max_seq = 0 + for item in items: + seq = IncrementalCheckpoint.filename_to_sequence_number(item["name"]) + if seq is not None and seq > max_seq: + max_seq = seq + + return max_seq + 1 + + def delete_incrementals(self) -> None: + """Delete all incremental checkpoints (called after compacting).""" + folder_id = self._get_checkpoints_folder_id() + if folder_id is None: + return + + query = ( + f"'{folder_id}' in parents and trashed=false " + f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" + ) + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id, name)" + ) + ) + items = results.get("files", []) + + for item in items: + try: + execute_with_retries( + self._cached_client.client.drive_service.files().delete( + fileId=item["id"] + ) + ) + except Exception as e: + print( + f"Warning: Failed to delete incremental checkpoint {item['name']}: {e}" + ) + + def _get_rolling_state_folder_name(self) -> str: + """Get the rolling state folder name: {email}-{version}-rolling-state""" + return f"{self._cached_client.email}-{SYFT_VERSION}-rolling-state" + + def _get_rolling_state_folder_id(self, use_cache: bool = True) -> str | None: + """ + Find the rolling state folder ID from Google Drive. + + Args: + use_cache: If True, return cached value if available. + """ + if use_cache and self._rolling_state_folder_id is not None: + return self._rolling_state_folder_id + + folders = self._cached_client.client.find_folders( + name_contains=[f"{self._cached_client.email}-", "-rolling-state"], + parent_id=self._cached_client.get_syftbox_folder_id(), + ) + folder_id = self._cached_client.find_or_adopt_versioned_folder( + folders, current_name=self._get_rolling_state_folder_name() + ) + if folder_id is not None: + self._rolling_state_folder_id = folder_id + return folder_id + + def _get_or_create_rolling_state_folder_id(self) -> str: + """Get or create the rolling state folder.""" + folder_id = self._get_rolling_state_folder_id() + if folder_id is not None: + return folder_id + # Create the folder + folder_name = self._get_rolling_state_folder_name() + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + folder_id = self._cached_client.client.create_folder( + folder_name, syftbox_folder_id + ) + self._rolling_state_folder_id = folder_id + return folder_id + + def put_rolling_state(self, name: str, data: bytes) -> str: + """Upload raw rolling state bytes to Google Drive. + + Optimized to use update() if file ID is cached. + """ + payload, _ = self._cached_client.client.create_file_payload(data) + + # Try to update existing file if we have a cached ID + if self._rolling_state_file_id is not None: + try: + execute_with_retries( + self._cached_client.client.drive_service.files().update( + fileId=self._rolling_state_file_id, + media_body=payload, + ) + ) + return self._rolling_state_file_id + except Exception as e: + # The cached file is gone or unreachable. Clear the cache and + # write a new file below. + print( + f"Warning: could not update rolling state " + f"{self._rolling_state_file_id}, writing a new file: {e}" + ) + self._rolling_state_file_id = None + + folder_id = self._get_or_create_rolling_state_folder_id() + + file_metadata = { + "name": name, + "parents": [folder_id], + } + + result = self._cached_client.client.create_syft_file(file_metadata, payload) + file_id = result.get("id") + self._rolling_state_file_id = file_id + return file_id + + @raise_unavailable("Could not read the rolling state from Drive.") + def get_rolling_state(self) -> bytes | None: + """Download the latest rolling state as raw bytes, or None. + + Also populates the folder and file ID caches for subsequent uploads. + """ + folder_id = self._get_rolling_state_folder_id() + if folder_id is None: + return None + + query = ( + f"'{folder_id}' in parents and trashed=false " + f"and name contains '{ROLLING_STATE_FILENAME_PREFIX}'" + ) + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id, name)" + ) + ) + items = results.get("files", []) + + if not items: + return None + + latest_file = None + latest_timestamp = -1.0 + for item in items: + timestamp = RollingState.filename_to_timestamp(item["name"]) + if timestamp is not None and timestamp > latest_timestamp: + latest_timestamp = timestamp + latest_file = item + + if latest_file is None: + return None + + data = self._cached_client.client.download_file(latest_file["id"]) + self._rolling_state_file_id = latest_file["id"] + return data + + def delete_rolling_state(self) -> None: + """Delete all existing rolling state files and clear cache.""" + # Clear the file ID cache + self._rolling_state_file_id = None + + folder_id = self._get_rolling_state_folder_id() + if folder_id is None: + return + + # List all rolling state files + query = ( + f"'{folder_id}' in parents and trashed=false " + f"and name contains '{ROLLING_STATE_FILENAME_PREFIX}'" + ) + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id, name)" + ) + ) + items = results.get("files", []) + + # Delete each rolling state file + for item in items: + try: + execute_with_retries( + self._cached_client.client.drive_service.files().delete( + fileId=item["id"] + ) + ) + except Exception as e: + print(f"Warning: Failed to delete rolling state {item['name']}: {e}") + + # -- the whole account ----------------------------------------------------- + + def delete_all(self) -> int: + """Everything in /SyftBox, and the syft files Drive moved to the root + of My Drive when they lost their folder.""" + file_ids = set(self.gather_all_file_and_folder_ids()) + file_ids |= set(self.find_orphaned_syft_files()) + self._cached_client.client.delete_multiple_files_by_ids(list(file_ids)) + return len(file_ids) + + def gather_all_file_and_folder_ids(self) -> list[str]: + """Every file and folder under /SyftBox, and /SyftBox itself.""" + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + return gather_all_file_and_folder_ids_recursive( + self._cached_client.client.drive_service, syftbox_folder_id + ) + + def find_orphaned_syft_files(self) -> list[str]: + """Find syft files of this account in its My Drive root, except /SyftBox. + + Drive moves a file to the root of its owner when the file loses its + folder. Only /SyftBox belongs in root. + + Returns list of orphaned file IDs. + """ + query = ( + "'root' in parents and 'me' in owners and trashed=false and " + f"name != '{SYFTBOX_FOLDER}' and " + f"properties has {{ key='{SYFT_FILE_PROPERTY_KEY}' " + f"and value='{SYFT_FILE_PROPERTY_VALUE}' }}" + ) + return [ + f["id"] + for f in self._cached_client.client.list_all_files(query, fields="id") + ] diff --git a/syft/sync/connections/drive/drive_p2p.py b/syft/sync/connections/drive/drive_p2p.py new file mode 100644 index 00000000000..9132e94cbbc --- /dev/null +++ b/syft/sync/connections/drive/drive_p2p.py @@ -0,0 +1,318 @@ +"""Messages between peers on Google Drive: the inbox and outbox folders.""" + +import threading +import warnings +from concurrent.futures import ThreadPoolExecutor + +from syft.sync.connections.drive.gdrive_client import DRIVE_BATCH_SIZE, GDriveClient +from syft.sync.connections.drive.drive_cached_client import ( + DriveCachedClient, + valid_events_from_file_metadatas, +) +from syft.sync.connections.drive.gdrive_retry import ( + batch_execute_with_retries, + execute_with_retries, + raise_unavailable, +) +from syft.sync.messages.proposed_filechange import FileNameParseError, MessageFileName +from syft.sync.transport.exceptions import TransportUnavailableError +from syft.sync.transport.protocols import P2PConnection +from syft.sync.transport.types import MessageKind, RawMessage, SendStatus + +# Threads that download the messages of one receive at the same time. +DOWNLOAD_WORKERS = 10 + + +def _download_or_skip(client: GDriveClient, listed: dict) -> bytes | None: + """The bytes of one listed inbox file, or None with a warning when this + file cannot be downloaded, for example because it was deleted. + + Raises: + TransportUnavailableError: when Drive cannot be reached. + """ + name = listed["file_name"] + try: + with raise_unavailable(f"Drive is unreachable; inbox message {name} waits"): + return client.download_file(listed["file_id"]) + except TransportUnavailableError: + raise + except Exception as e: + warnings.warn(f"Skipped inbox message {name}: download failed: {e!r}") + return None + + +def _require_inbox_kind(kind: MessageKind) -> None: + """Only the inbox is read through ``receive`` so far.""" + if kind is not MessageKind.PROPOSED_CHANGES: + raise NotImplementedError( + f"Messages of kind {kind!r} are still read through the outbox " + "methods; they move to receive with send-and-forget." + ) + + +def _valid_message_names(file_metadatas: list[dict]) -> list[MessageFileName]: + res = [] + for file_metadata in file_metadatas: + try: + res.append(MessageFileName.from_string(file_metadata["name"])) + except FileNameParseError: + continue + return res + + +class DriveP2PConnection(P2PConnection): + """Messages through the shared inbox and outbox folders. + + A proposed change goes to the DO's inbox and file events go to the DS's + outbox. ``msg_id`` is the file name. A written file is ``DELIVERED``: it + waits in the folder until the peer reads it. + + The inbox is read through ``receive`` and a processed message is archived. + The outbox is still read through the Drive-only methods below, from a + timestamp cursor, and kept: the DS rebuilds a lost cache by reading it + again. It moves to ``receive`` with send-and-forget, once the DS keeps what + it received in its own storage. + """ + + def __init__(self, cached_client: DriveCachedClient): + self._cached_client = cached_client + + def clone_for_thread(self) -> "DriveP2PConnection": + """A connection over a new Drive client, for use from another thread.""" + return DriveP2PConnection(self._cached_client.clone_for_thread()) + + def reset_caches(self) -> None: + self._cached_client.reset_caches() + + def send( + self, peer: str, msg_id: str, data: bytes, kind: MessageKind + ) -> SendStatus: + if kind is MessageKind.PROPOSED_CHANGES: + self.watcher_send_raw_bytes_to_inbox(peer, msg_id, data) + elif kind is MessageKind.FILE_EVENTS: + self.owner_write_raw_bytes_to_outbox(peer, msg_id, data) + else: + raise ValueError(f"Drive has no folder for messages of kind {kind!r}") + return SendStatus.DELIVERED + + def delivered(self) -> list[str]: + # Every Drive send is DELIVERED when it returns; none waits for a peer. + return [] + + def receive( + self, peer: str, kind: MessageKind, max_messages: int | None = None + ) -> list[RawMessage]: + """One listing for the whole inbox, then the files download in parallel. + + A file that cannot be downloaded is left out with a warning, so it does + not hold back the others. + + Raises: + TransportUnavailableError: when Drive cannot be reached. Every + message stays in the inbox. + """ + _require_inbox_kind(kind) + listed = self._list_inbox(peer)[:max_messages] + contents = self._download_all(listed) + return [ + RawMessage( + sender=peer, + msg_id=m["file_name"], + kind=kind, + data=data, + handle=m["file_id"], + ) + for m, data in zip(listed, contents) + if data is not None + ] + + def _download_all(self, listed: list[dict]) -> list[bytes | None]: + """The bytes of each listed file, in order; None for a file that could + not be downloaded. Several files download over a thread pool with one + client copy per thread, because the Drive client is not thread-safe.""" + if len(listed) <= 1: + return [_download_or_skip(self._cached_client.client, m) for m in listed] + per_thread = threading.local() + + def download(m: dict) -> bytes | None: + if not hasattr(per_thread, "client"): + per_thread.client = self._cached_client.client.clone_for_thread() + return _download_or_skip(per_thread.client, m) + + workers = min(DOWNLOAD_WORKERS, len(listed)) + with ThreadPoolExecutor(max_workers=workers) as pool: + return list(pool.map(download, listed)) + + def mark_processed(self, message: RawMessage) -> None: + # The archive folder is named after ``message.sender``, which receive + # set from the inbox folder, never from the message body. + _require_inbox_kind(message.kind) + execute_with_retries(self._archive_request(message.handle, message.sender)) + + def mark_all_processed(self, messages: list[RawMessage]) -> None: + """Archive every message, up to ``DRIVE_BATCH_SIZE`` moves per request. + + A move that fails is warned about and its file stays in the inbox. + """ + for message in messages: + _require_inbox_kind(message.kind) + + def callback(request_id, response, exception): + if exception is not None: + warnings.warn( + f"Could not archive inbox file {request_id}: {exception!r}" + ) + + drive_service = self._cached_client.client.drive_service + for i in range(0, len(messages), DRIVE_BATCH_SIZE): + batch = drive_service.new_batch_http_request(callback=callback) + for m in messages[i : i + DRIVE_BATCH_SIZE]: + batch.add( + self._archive_request(m.handle, m.sender), request_id=m.handle + ) + batch_execute_with_retries(batch) + + def outbox_messages(self, peer: str, since_timestamp: float | None) -> list[bytes]: + """The messages in ``peer``'s outbox to me, newer than ``since_timestamp``.""" + return [ + self._cached_client.client.download_file(m["file_id"]) + for m in self.outbox_message_metadatas(peer, since_timestamp) + ] + + def outbox_message_metadatas( + self, peer: str, since_timestamp: float | None + ) -> list[dict]: + """Like ``outbox_messages``, without downloading: each entry's + ``file_id`` is the handle to ``fetch_message`` it with.""" + folder_id = self._cached_client.get_peer_datasite_outbox_id(peer) + if folder_id is None: + return [] + + file_metadatas = self._cached_client.get_file_metadatas_from_folder( + folder_id, since_timestamp=since_timestamp + ) + valid_fname_objs = valid_events_from_file_metadatas(file_metadatas) + name_to_id = {f["name"]: f["id"] for f in file_metadatas} + + result = [] + for fname_obj in sorted(valid_fname_objs, key=lambda x: x.timestamp): + if since_timestamp is None or fname_obj.timestamp > since_timestamp: + file_name = fname_obj.as_string() + if file_name in name_to_id: + result.append( + { + "file_id": name_to_id[file_name], + "file_name": file_name, + "timestamp": fname_obj.timestamp, + } + ) + return result + + def fetch_message(self, handle: str) -> bytes: + return self._cached_client.client.download_file(handle) + + def own_outbox_message_metadatas(self, peer: str) -> list[dict]: + """The messages in my outbox to ``peer``, without downloading. An + outbox whose folder is not found holds none.""" + folder_id = self._cached_client.get_own_datasite_outbox_id(peer) + if folder_id is None: + return [] + return self._cached_client.get_file_metadatas_from_folder( + folder_id, since_timestamp=None + ) + + def delete_messages(self, handles: list[str]) -> None: + self._cached_client.client.delete_multiple_files_by_ids(handles) + + # -- Drive ----------------------------------------------------------------- + + def _list_inbox(self, peer: str) -> list[dict]: + """The messages in the inbox from ``peer``, oldest first, each with + ``file_id`` and ``file_name``, without downloading them.""" + inbox_folder_id = self._cached_client.get_own_datasite_inbox_id(peer) + if inbox_folder_id is None: + raise ValueError(f"Inbox folder for {peer} not found") + file_metadatas = self._cached_client.get_file_metadatas_from_folder( + inbox_folder_id + ) + name_to_id = {f["name"]: f["id"] for f in file_metadatas} + return [ + {"file_id": name_to_id[name.as_string()], "file_name": name.as_string()} + for name in sorted( + _valid_message_names(file_metadatas), + key=lambda x: x.submitted_timestamp, + ) + ] + + def _archive_request(self, file_id: str, sender: str): + """The request that moves an inbox file to its archive folder. The inbox + folder is the file's only parent and both folder ids are cached, so the + move is one Drive call. ``sender`` names the archive folder, so the + caller reads it from the transport, never from the message body.""" + return self._cached_client.client.drive_service.files().update( + fileId=file_id, + addParents=self._cached_client.get_archive_folder_id(sender), + removeParents=self._cached_client.get_own_datasite_inbox_id(sender), + fields="id", + supportsAllDrives=True, + ) + + def watcher_send_raw_bytes_to_inbox( + self, recipient: str, filename: str, data: bytes + ) -> None: + inbox_id = self._cached_client.get_peer_datasite_inbox_id(recipient) + if inbox_id is None: + raise Exception(f"Inbox folder for {recipient}'s datasite not found") + + payload, _ = self._cached_client.client.create_file_payload(data) + file_metadata = {"name": filename, "parents": [inbox_id]} + + result = self._cached_client.client.create_syft_file( + file_metadata, payload, fields="id, parents" + ) + + file_id = result.get("id") + actual_parents = result.get("parents", []) + if inbox_id not in actual_parents: + print( + f"WARNING: Message file {file_id} was not placed in inbox folder " + f"{inbox_id}. Actual parents: {actual_parents}. " + f"Moving file to correct folder..." + ) + execute_with_retries( + self._cached_client.client.drive_service.files().update( + fileId=file_id, + addParents=inbox_id, + fields="id, parents", + ) + ) + + def owner_write_raw_bytes_to_outbox( + self, recipient: str, filename: str, data: bytes + ) -> None: + outbox_folder_id = self._cached_client.get_own_datasite_outbox_id(recipient) + if outbox_folder_id is None: + raise ValueError(f"Outbox folder for {recipient} not found") + + file_payload, _ = self._cached_client.client.create_file_payload(data) + file_metadata = {"name": filename, "parents": [outbox_folder_id]} + + result = self._cached_client.client.create_syft_file( + file_metadata, file_payload, fields="id, parents" + ) + + file_id = result.get("id") + actual_parents = result.get("parents", []) + if outbox_folder_id not in actual_parents: + print( + f"WARNING: Event message {file_id} was not placed in outbox folder " + f"{outbox_folder_id}. Actual parents: {actual_parents}. " + f"Moving file to correct folder..." + ) + execute_with_retries( + self._cached_client.client.drive_service.files().update( + fileId=file_id, + addParents=outbox_folder_id, + fields="id, parents", + ) + ) diff --git a/syft/sync/connections/drive/drive_peering.py b/syft/sync/connections/drive/drive_peering.py new file mode 100644 index 00000000000..7884de0f838 --- /dev/null +++ b/syft/sync/connections/drive/drive_peering.py @@ -0,0 +1,457 @@ +"""Finding and trusting peers on Google Drive. + +The peer list is /SyftBox/SYFT_peers.json, the version document is +/SyftBox/SYFT_version.json, and the key bundles sit in a folder of their own. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, Callable, Optional + +from pydantic import BaseModel +from syft_migration import MigrationError + +from syft.sync.connections.drive.gdrive_client import GOOGLE_FOLDER_MIME_TYPE +from syft.sync.connections.drive.drive_cached_client import ( + GDRIVE_P2P_FOLDER_DATASITE_PREFIX, + DriveCachedClient, + GdriveP2PFolder, +) +from syft.sync.connections.drive.gdrive_retry import ( + execute_with_retries, + raise_unavailable, +) +from syft.sync.transport.naming import PEERS_META_KEY, SYFT_PEERS_VERSION +from syft.sync.transport.protocols import Peering + +if TYPE_CHECKING: + from syft.sync.version.peer_info import PeerInfo + + +SYFT_PEERS_FILE = "SYFT_peers.json" + +SYFT_VERSION_FILE = "SYFT_version.json" + + +class GdriveEncryptionBundlesFolder(BaseModel): + """Folder for encryption bundles: syft_encryption_bundles#email""" + + email: str + + def as_string(self) -> str: + return f"syft_encryption_bundles#{self.email}" + + +class DrivePeering(Peering): + """Identity from the Google account; channels, the peer list, key bundles + and version files in /SyftBox.""" + + def __init__(self, cached_client: DriveCachedClient): + self._cached_client = cached_client + # Encryption bundles folder cache + self._encryption_bundles_folder_id: str | None = None + # SYFT_peers.json as loaded: peer email -> {"state", "public_encryption_bundle"}, + # plus "_meta" -> {"version"}. None until the first read. + self._peers_json_cache: dict[str, dict[str, Any]] | None = None + + def clone_for_thread(self) -> "DrivePeering": + """Peering over a new Drive client, for use from another thread, with + these caches.""" + clone = DrivePeering(self._cached_client.clone_for_thread()) + clone._encryption_bundles_folder_id = self._encryption_bundles_folder_id + clone._peers_json_cache = ( + dict(self._peers_json_cache) if self._peers_json_cache is not None else None + ) + return clone + + def reset_caches(self) -> None: + self._cached_client.reset_caches() + self._encryption_bundles_folder_id = None + self._peers_json_cache = None + + def whoami(self) -> str: + """Return the email of the Google account behind the Drive service.""" + about = execute_with_retries( + self._cached_client.client.drive_service.about().get( + fields="user(emailAddress)" + ) + ) + return about["user"]["emailAddress"] + + def open_channels(self, peer: str): + """Create inbox + outbox folders for peer's datasite on own drive, share with peer. + + Creates: + - syft_datasite_{peer}_inbox_{self} — I submit proposals to peer here + - syft_datasite_{peer}_outbox_{self} — peer pushes events to me here + """ + # Inbox: I write proposals, peer reads + inbox_id = self._cached_client.get_peer_datasite_inbox_id(peer) + if inbox_id is None: + inbox_id = self._cached_client.create_peer_datasite_folder(peer, "inbox") + self._cached_client.client.add_permission(inbox_id, peer, write=True) + + # Outbox: peer writes events, I read + outbox_id = self._cached_client.get_peer_datasite_outbox_id(peer) + if outbox_id is None: + outbox_id = self._cached_client.create_peer_datasite_folder(peer, "outbox") + self._cached_client.client.add_permission(outbox_id, peer, write=True) + + def _get_peers_file_id(self) -> str | None: + """Find SYFT_peers.json file in /SyftBox folder""" + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + query = f"name='{SYFT_PEERS_FILE}' and '{syftbox_folder_id}' in parents and trashed=false" + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) + ) + items = results.get("files", []) + return items[0]["id"] if items else None + + def _download_peers_json(self) -> dict[str, dict[str, Any]]: + """Fetch peers JSON from GDrive. Returns empty dict if not found. + + Maps each peer email to ``{"state": str, "public_encryption_bundle": + dict}``, and ``_meta`` to ``{"version": int}``. + + Raises: + TransportUnavailableError: when Drive cannot be reached. + ValueError: when the file is not valid JSON. + """ + with raise_unavailable("Could not read the peers file from Drive."): + file_id = self._get_peers_file_id() + if file_id is None: + return {} + file_data = self._cached_client.client.download_file(file_id) + try: + return json.loads(file_data.decode("utf-8")) + except ValueError as e: + raise ValueError(f"Could not read the peers file: {e}") from e + + def read_peer_registry( + self, force_download: bool = False + ) -> dict[str, dict[str, Any]]: + """Return peers JSON, using the in-memory cache when available.""" + if self._peers_json_cache is not None and not force_download: + return self._peers_json_cache + self._peers_json_cache = self._download_peers_json() + return self._peers_json_cache + + def _write_peers_json(self, peers_data: dict[str, dict[str, Any]]): + """Write peers JSON to GDrive. Creates or updates the file.""" + peers_data = { + **peers_data, + PEERS_META_KEY: {"version": SYFT_PEERS_VERSION}, + } + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + file_id = self._get_peers_file_id() + + # Convert to JSON bytes + json_data = json.dumps(peers_data, indent=2) + file_payload, _ = self._cached_client.client.create_file_payload(json_data) + + if file_id is None: + # Create new file + file_metadata = { + "name": SYFT_PEERS_FILE, + "parents": [syftbox_folder_id], + } + result = self._cached_client.client.create_syft_file( + file_metadata, file_payload + ) + self._peers_json_cache = peers_data + return result.get("id") + else: + # Update existing file + execute_with_retries( + self._cached_client.client.drive_service.files().update( + fileId=file_id, media_body=file_payload + ) + ) + self._peers_json_cache = peers_data + return file_id + + def update_peer_registry( + self, + peer: str, + state: str, + public_encryption_bundle: dict | None = None, + ): + """Update a single peer's state in the JSON file. + + Preserves existing fields when updating state. + """ + peers_data = self.read_peer_registry() + existing = peers_data.get(peer, {}) + existing["state"] = state + if public_encryption_bundle is not None: + existing["public_encryption_bundle"] = public_encryption_bundle + peers_data[peer] = existing + self._write_peers_json(peers_data) + + def incoming_requests(self) -> list[str]: + """Get list of pending peer requests. + + Scans for syft_datasite_#version#{self}_* folders NOT owned by self — those are + peers who created folders for our datasite. Filters out already-accepted + or rejected peers from SYFT_peers.json. + """ + folders = self._cached_client.client.list_all_files( + f"name contains '{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#' " + f"and name contains '#{self._cached_client.email}#' " + f"and trashed=false " + f"and mimeType = '{GOOGLE_FOLDER_MIME_TYPE}' " + f"and not 'me' in owners", + fields="id,name,owners", + ) + + all_folder_peers = set() + for f in folders: + try: + folder = GdriveP2PFolder.from_name(f["name"]) + except ValueError: + # The query matches a name prefix, so a folder with another shape + # can appear here. + continue + # Any account can make a folder with this name. The request is from + # the owner, so the name must name the owner. + owners = f.get("owners", []) + if ( + len(owners) != 1 + or owners[0].get("emailAddress", "").casefold() + != folder.peer_email.casefold() + ): + continue + if folder.datasite_email == self._cached_client.email: + all_folder_peers.add(folder.peer_email) + + peers_data = self.read_peer_registry() + pending_peers = [] + for peer_email in all_folder_peers: + undecided = peers_data.get(peer_email, {}).get("state") not in [ + "accepted", + "rejected", + ] + if undecided: + pending_peers.append(peer_email) + + return pending_peers + + def _get_own_peer_info_file_id(self) -> Optional[str]: + """Find the own peer info file, SYFT_version.json, in /SyftBox.""" + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + query = f"name='{SYFT_VERSION_FILE}' and '{syftbox_folder_id}' in parents and trashed=false" + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) + ) + items = results.get("files", []) + return items[0]["id"] if items else None + + def publish_peer_info(self, peer_info: "PeerInfo") -> None: + """Write the peer info file to /SyftBox. Creates or updates the file.""" + + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() + file_id = self._get_own_peer_info_file_id() + + # Convert to JSON string + json_data = peer_info.to_json() + file_payload, _ = self._cached_client.client.create_file_payload(json_data) + + if file_id is None: + # Create new file + file_metadata = { + "name": SYFT_VERSION_FILE, + "parents": [syftbox_folder_id], + } + self._cached_client.client.create_syft_file(file_metadata, file_payload) + else: + # Update existing file + execute_with_retries( + self._cached_client.client.drive_service.files().update( + fileId=file_id, media_body=file_payload + ) + ) + + def _get_peer_info_file_id_of(self, peer_email: str) -> Optional[str]: + """Find the peer info file, SYFT_version.json, of ``peer_email``.""" + query = ( + f"name='{SYFT_VERSION_FILE}' and trashed=false and '{peer_email}' in owners" + ) + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) + ) + items = results.get("files", []) + return items[0]["id"] if items else None + + def read_own_peer_info(self) -> Optional["PeerInfo"]: + """Read the own peer info file from /SyftBox. + + Returns None when there is no file. + + Raises: + TransportUnavailableError: when Drive cannot be reached. + ValueError: when the file is not a peer info file this client reads. + """ + return self._read_peer_info_file( + self._get_own_peer_info_file_id, "the own peer info file" + ) + + def read_peer_info(self, peer: str) -> Optional["PeerInfo"]: + """Read the peer info file from a peer's /SyftBox folder. + + Returns None when the peer has no file that this account may read. + + Raises: + TransportUnavailableError: when Drive cannot be reached. + ValueError: when the file is not a peer info file this client reads. + """ + return self._read_peer_info_file( + lambda: self._get_peer_info_file_id_of(peer), + f"the peer info file of {peer}", + ) + + def _read_peer_info_file( + self, find_file_id: Callable[[], str | None], file_label: str + ) -> Optional["PeerInfo"]: + """Find, download and parse a peer info file. + + Returns None when ``find_file_id`` finds no file. ``file_label`` names + the file in errors. + """ + from syft.sync.version.peer_info import PeerInfo + + with raise_unavailable(f"Could not read {file_label} from Drive."): + file_id = find_file_id() + if file_id is None: + return None + file_data = self._cached_client.client.download_file(file_id) + try: + return PeerInfo.from_json(file_data.decode("utf-8")) + except (ValueError, MigrationError) as e: + raise ValueError(f"Could not read {file_label}: {e}") from e + + def grant_peer_info_read(self, peer: str, peer_info: "PeerInfo") -> None: + """Share the peer info file with a peer so they can read it. + + A missing file is created from ``peer_info``, the info of this client. + """ + file_id = self._get_own_peer_info_file_id() + if file_id is None: + self.publish_peer_info(peer_info) + file_id = self._get_own_peer_info_file_id() + + if file_id: + self._cached_client.client.add_permission(file_id, peer, write=False) + + def _find_encryption_bundles_folder_id(self) -> str | None: + """The id of the encryption bundles folder of this user, or None when it does not exist.""" + if self._encryption_bundles_folder_id: + return self._encryption_bundles_folder_id + folder_name = GdriveEncryptionBundlesFolder( + email=self._cached_client.email + ).as_string() + folder_id = self._cached_client.client.find_folder_by_name( + folder_name, parent_id=self._cached_client.get_syftbox_folder_id() + ) + if folder_id: + self._encryption_bundles_folder_id = folder_id + return folder_id + + def _get_or_create_encryption_bundles_folder_id(self) -> str: + """Get or create the encryption bundles folder for this user.""" + folder_id = self._find_encryption_bundles_folder_id() + if not folder_id: + folder_name = GdriveEncryptionBundlesFolder( + email=self._cached_client.email + ).as_string() + folder_id = self._cached_client.client.create_folder( + folder_name, self._cached_client.get_syftbox_folder_id() + ) + self._encryption_bundles_folder_id = folder_id + return folder_id + + def _encryption_bundle_filename(self, owner_email: str, peer_email: str) -> str: + return f"encryption_bundle_{owner_email}_for_{peer_email}.json" + + def publish_key_bundle(self, peer: str, bundle_json: str) -> None: + """Write own encryption bundle for a peer to own bundles folder.""" + folder_id = self._get_or_create_encryption_bundles_folder_id() + filename = self._encryption_bundle_filename(self._cached_client.email, peer) + file_payload, _ = self._cached_client.client.create_file_payload(bundle_json) + + # Check if file already exists + query = f"name='{filename}' and '{folder_id}' in parents and trashed=false" + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) + ) + items = results.get("files", []) + + if items: + execute_with_retries( + self._cached_client.client.drive_service.files().update( + fileId=items[0]["id"], media_body=file_payload + ) + ) + else: + file_metadata = {"name": filename, "parents": [folder_id]} + self._cached_client.client.create_syft_file(file_metadata, file_payload) + + def grant_key_bundle_read(self, peer: str) -> None: + """Share the bundles folder with a peer so they can read bundles.""" + folder_id = self._get_or_create_encryption_bundles_folder_id() + self._cached_client.client.add_permission(folder_id, peer, write=False) + + def read_own_key_bundle(self, peer: str) -> str | None: + """Read the bundle this client published for ``peer``, or None. + + A read never creates the bundles folder: without it, nothing is published. + """ + folder_id = self._find_encryption_bundles_folder_id() + if folder_id is None: + return None + filename = self._encryption_bundle_filename(self._cached_client.email, peer) + query = f"name='{filename}' and '{folder_id}' in parents and trashed=false" + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) + ) + items = results.get("files", []) + if not items: + return None + return self._cached_client.client.download_file(items[0]["id"]).decode("utf-8") + + def read_peer_key_bundle(self, peer: str) -> str | None: + """Read encryption bundle that peer wrote for us. + + Searches for: encryption_bundle_{peer}_for_{self._cached_client.email}.json + in peer's bundles folder. + """ + filename = self._encryption_bundle_filename(peer, self._cached_client.email) + query = f"name='{filename}' and trashed=false and '{peer}' in owners" + results = execute_with_retries( + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) + ) + items = results.get("files", []) + if not items: + return None + try: + data = self._cached_client.client.download_file(items[0]["id"]) + except Exception as e: + print(f"Warning: could not download the bundle of {peer}: {e}") + return None + try: + return data.decode("utf-8") + except ValueError as e: + print(f"Warning: could not read the bundle of {peer}: {e}") + return None diff --git a/syft/sync/connections/drive/gdrive_client.py b/syft/sync/connections/drive/gdrive_client.py new file mode 100644 index 00000000000..46db30bd4aa --- /dev/null +++ b/syft/sync/connections/drive/gdrive_client.py @@ -0,0 +1,297 @@ +"""The Google Drive API client that the Drive transport jobs share. + +One client per account. It holds the Drive service, which is not safe to +use from more than one thread. +""" + +import io +import json +import logging +import pickle +from dataclasses import dataclass +from typing import Any, Callable, Optional + +from google.oauth2.credentials import Credentials as GoogleCredentials +from google_auth_httplib2 import AuthorizedHttp +from googleapiclient.discovery import build +from googleapiclient.errors import HttpError +from googleapiclient.http import MediaIoBaseDownload, MediaIoBaseUpload, build_http + +from syft.sync.connections.drive.gdrive_retry import ( + batch_execute_with_retries, + execute_with_retries, + next_chunk_with_retries, +) +from syft.sync.environments.environment import Environment + +GOOGLE_FOLDER_MIME_TYPE = "application/vnd.google-apps.folder" +# Every file and folder that syft creates carries this property. +SYFT_FILE_PROPERTY_KEY = "syft" +SYFT_FILE_PROPERTY_VALUE = "true" +# Google Drive batch API has a limit of 100 requests per batch +DRIVE_BATCH_SIZE = 100 + + +# Timeout for Google API requests (in seconds) +GOOGLE_API_TIMEOUT = 120 # 2 minutes + +SCOPES = ["https://www.googleapis.com/auth/drive"] + + +logging.getLogger("google_auth_httplib2").setLevel(logging.ERROR) + + +def build_drive_service( + credentials: GoogleCredentials, + timeout: int = GOOGLE_API_TIMEOUT, + environment: Environment | None = None, +): + """Build a Google Drive service with timeout-enabled authorized HTTP.""" + # Build the http via googleapiclient's own factory rather than constructing + # httplib2.Http() directly. build_http() applies Google-API-specific tweaks. + http = build_http() + http.timeout = timeout + if environment == Environment.COLAB: + import google.auth + from google.colab import auth as colab_auth + + colab_auth.authenticate_user() + creds, _ = google.auth.default() + authed_http = AuthorizedHttp(creds, http=http) + # Build service without explicit credentials in Colab + return build("drive", "v3", http=authed_http) + else: + authorized_http = AuthorizedHttp(credentials, http=http) + return build("drive", "v3", http=authorized_http) + + +def _has_reason(error: HttpError, reason: str) -> bool: + """True when Drive gave ``reason`` for ``error``.""" + details = error.error_details + return isinstance(details, list) and any(d.get("reason") == reason for d in details) + + +@dataclass +class GDriveClient: + # A googleapiclient Drive v3 service, or a MockDriveService in tests. + drive_service: Any + # Builds another service like ``drive_service``, for a copy in another thread. + new_service: Callable[[], Any] + + def clone_for_thread(self) -> "GDriveClient": + """A client over a new Drive service, for use from another thread.""" + return GDriveClient(self.new_service(), self.new_service) + + def download_file(self, file_id: str) -> bytes: + request = self.drive_service.files().get_media(fileId=file_id) + + file_buffer = io.BytesIO() + downloader = MediaIoBaseDownload( + file_buffer, request, chunksize=1024 * 1024 * 10 + ) + + done = False + while not done: + _, done = next_chunk_with_retries(downloader) + + message_data = file_buffer.getvalue() + return message_data + + def create_file_payload(self, data: Any) -> tuple[MediaIoBaseUpload, str]: + """Create a file payload for the GDrive""" + if isinstance(data, str): + file_data = data.encode("utf-8") + mime_type = "text/plain" + extension = ".txt" + elif isinstance(data, dict): + file_data = json.dumps(data, indent=2).encode("utf-8") + mime_type = "application/json" + extension = ".json" + elif isinstance(data, bytes): + file_data = data + mime_type = "application/octet-stream" + extension = ".bin" + else: + # Pickle for other data types + file_data = pickle.dumps(data) + mime_type = "application/octet-stream" + extension = ".pkl" + + media = MediaIoBaseUpload( + io.BytesIO(file_data), mimetype=mime_type, resumable=True + ) + + return media, extension + + def create_folder(self, folder_name: str, parent_id: str) -> str: + folder = self.create_syft_file( + { + "name": folder_name, + "mimeType": GOOGLE_FOLDER_MIME_TYPE, + "parents": [parent_id], + } + ) + return folder["id"] + + def create_syft_file( + self, + body: dict, + media_body: MediaIoBaseUpload | None = None, + fields: str = "id", + ) -> dict: + """Create a file or folder with the syft tag.""" + return execute_with_retries( + self.drive_service.files().create( + body={ + **body, + "properties": {SYFT_FILE_PROPERTY_KEY: SYFT_FILE_PROPERTY_VALUE}, + }, + media_body=media_body, + fields=fields, + ) + ) + + def delete_file_by_id(self, file_id: str) -> None: + execute_with_retries(self.drive_service.files().delete(fileId=file_id)) + + def delete_multiple_files_by_ids( + self, + file_ids: list[str], + ignore_permissions_errors: bool = True, + ignore_file_not_found: bool = True, + ): + def callback(request_id, response, exception): + if exception is None: + return + # The file is already deleted. + if ignore_file_not_found and exception.resp.status == 404: + return + # The file belongs to another account, or was removed by one. + if ignore_permissions_errors and _has_reason( + exception, "insufficientFilePermissions" + ): + return + raise exception + + for i in range(0, len(file_ids), DRIVE_BATCH_SIZE): + chunk = file_ids[i : i + DRIVE_BATCH_SIZE] + batch = self.drive_service.new_batch_http_request(callback=callback) + for file_id in chunk: + batch.add(self.drive_service.files().delete(fileId=file_id)) + batch_execute_with_retries(batch) + + def has_permission(self, file_id: str, email: str) -> bool: + """Check if user already has permission on the file.""" + perms = execute_with_retries( + self.drive_service.permissions().list( + fileId=file_id, fields="permissions(emailAddress)" + ) + ) + for p in perms.get("permissions", []): + if p.get("emailAddress", "").lower() == email.lower(): + return True + return False + + def add_permission(self, file_id: str, recipient: str, write=False): + """Add permission to the file if not already shared.""" + if self.has_permission(file_id, recipient): + return + + role = "writer" if write else "reader" + permission = { + "type": "user", + "role": role, + "emailAddress": recipient, + } + execute_with_retries( + self.drive_service.permissions().create( + fileId=file_id, body=permission, sendNotificationEmail=True + ) + ) + + def batch_add_permissions(self, file_id: str, users: list[str]) -> None: + """Add reader permissions for multiple users in a single batch request.""" + + def callback(request_id, response, exception): + # Drive answers a grant the user already has with that permission. + if exception is not None: + raise exception + + for i in range(0, len(users), DRIVE_BATCH_SIZE): + chunk = users[i : i + DRIVE_BATCH_SIZE] + batch = self.drive_service.new_batch_http_request(callback=callback) + for user_email in chunk: + permission = { + "type": "user", + "role": "reader", + "emailAddress": user_email, + } + batch.add( + self.drive_service.permissions().create( + fileId=file_id, + body=permission, + sendNotificationEmail=True, + ) + ) + batch_execute_with_retries(batch) + + def find_folder_by_name( + self, folder_name: str, parent_id: str = None, owner_email: str = None + ) -> Optional[str]: + """Find a folder by name, optionally within a specific parent""" + # parent_id = "1AQ3WLnVlLd6Zjo7p9Z_qGA1Djjf6-KIh" + owner_email_clause = f"and '{owner_email}' in owners" if owner_email else "" + parent_id_clause = f"and '{parent_id}' in parents" if parent_id else "" + query = f"name='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false {owner_email_clause} {parent_id_clause}" + + results = execute_with_retries( + self.drive_service.files().list(q=query, fields="files(id)", pageSize=1) + ) + items = results.get("files", []) + return items[0]["id"] if items else None + + def list_all_files(self, query: str, fields: str) -> list[dict]: + """Every file that matches ``query``, from all result pages. + + ``fields`` names the file fields to return, such as ``"id, name"``. + """ + files: list[dict] = [] + page_token = None + while True: + page = execute_with_retries( + self.drive_service.files().list( + q=query, + fields=f"files({fields}), nextPageToken", + pageToken=page_token, + ) + ) + files.extend(page["files"]) + page_token = page.get("nextPageToken") + if not page_token: + return files + + def find_folders( + self, + *, + name_contains: list[str] = (), + parent_id: str | None = None, + owner_email: str | None = None, + ) -> list[tuple[str, str]]: + """List folders matching the constraints. Returns (id, name) pairs. + + Thin wrapper over Drive's files.list -- handles query building and + pagination, knows nothing about versions. Pair with + _partition_by_version or _sorted_by_version when the caller cares about + the version in the folder name. + """ + clauses = [f"mimeType='{GOOGLE_FOLDER_MIME_TYPE}'", "trashed=false"] + for substr in name_contains: + clauses.append(f"name contains '{substr}'") + if parent_id: + clauses.append(f"'{parent_id}' in parents") + if owner_email: + clauses.append(f"'{owner_email}' in owners") + query = " and ".join(clauses) + + files = self.list_all_files(query, fields="id, name") + return [(f["id"], f["name"]) for f in files] diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py deleted file mode 100644 index a589b9b23a7..00000000000 --- a/syft/sync/connections/drive/gdrive_transport.py +++ /dev/null @@ -1,2448 +0,0 @@ -"""Google Drive Files transport layer implementation""" - -import io -import json -import logging -import pickle -import re -import warnings -from functools import lru_cache -from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional - -from google.oauth2.credentials import Credentials as GoogleCredentials -from google_auth_httplib2 import AuthorizedHttp -from googleapiclient.discovery import build -from googleapiclient.http import MediaIoBaseDownload, MediaIoBaseUpload, build_http -from pydantic import BaseModel, field_validator -from syft_migration import MigrationError - -from syft.sync.checkpoints.checkpoint import ( - CHECKPOINT_FILENAME_PREFIX, - INCREMENTAL_CHECKPOINT_PREFIX, - Checkpoint, - IncrementalCheckpoint, -) -from syft.sync.checkpoints.rolling_state import ( - ROLLING_STATE_FILENAME_PREFIX, - RollingState, -) -from syft.sync.connections.base_connection import ( - FileCollection, - SyftboxPlatformConnection, -) -from syft.sync.connections.drive.gdrive_retry import ( - batch_execute_with_retries, - execute_with_retries, - next_chunk_with_retries, - raise_unavailable, -) -from syft.sync.connections.drive.gdrive_utils import ( - gather_all_file_and_folder_ids_recursive, -) -from syft.sync.environments.environment import Environment -from syft.sync.events.file_change_event import ( - FileChangeEventsMessage, - FileChangeEventsMessageFileName, -) -from syft.sync.messages.proposed_filechange import ( - FileNameParseError, - MessageFileName, - ProposedFileChangesMessage, -) -from syft.sync.transport.naming import ( # noqa: F401 (re-exported) - PEERS_META_KEY, - SYFT_PEERS_VERSION, - timestamp_from_message_name, -) -from syft.sync.utils.syftbox_utils import check_env -from syft.sync.version.version_info import _parse_semver -from syft.version import SYFT_VERSION - -if TYPE_CHECKING: - from syft.sync.connections.drive.grdrive_config import ( - GdriveConnectionConfig, - ) - from syft.sync.version.peer_info import PeerInfo - -# Timeout for Google API requests (in seconds) -GOOGLE_API_TIMEOUT = 120 # 2 minutes -# Google Drive batch API has a limit of 100 requests per batch -DRIVE_BATCH_SIZE = 100 - -SYFTBOX_FOLDER = "SyftBox" -GOOGLE_FOLDER_MIME_TYPE = "application/vnd.google-apps.folder" -SCOPES = ["https://www.googleapis.com/auth/drive"] - - -logging.getLogger("google_auth_httplib2").setLevel(logging.ERROR) - - -def build_drive_service( - credentials: GoogleCredentials, - timeout: int = GOOGLE_API_TIMEOUT, - environment: Environment | None = None, -): - """Build a Google Drive service with timeout-enabled authorized HTTP.""" - # Build the http via googleapiclient's own factory rather than constructing - # httplib2.Http() directly. build_http() applies Google-API-specific tweaks. - http = build_http() - http.timeout = timeout - if environment == Environment.COLAB: - import google.auth - from google.colab import auth as colab_auth - - colab_auth.authenticate_user() - creds, _ = google.auth.default() - authed_http = AuthorizedHttp(creds, http=http) - # Build service without explicit credentials in Colab - return build("drive", "v3", http=authed_http) - else: - authorized_http = AuthorizedHttp(credentials, http=http) - return build("drive", "v3", http=authorized_http) - - -GDRIVE_P2P_FOLDER_DATASITE_PREFIX = "syft_datasite" -SYFT_PEERS_FILE = "SYFT_peers.json" - -SYFT_VERSION_FILE = "SYFT_version.json" -# Every file and folder that syft creates carries this property. -SYFT_FILE_PROPERTY_KEY = "syft" -SYFT_FILE_PROPERTY_VALUE = "true" - - -class GdriveArchiveFolder(BaseModel): - sender_email: str - recipient_email: str - - def as_string(self) -> str: - return f"syft_{self.sender_email}_to_{self.recipient_email}_archive" - - -# An email part of a folder name goes into Drive queries and local paths. It has -# one '@', a '.' in the domain, and no quote, slash, '#', whitespace or control -# character. It is not RFC 5322. -_EMAIL_PART = r"[^@\s'\"\\/#\x00-\x1f\x7f]+" -_FOLDER_EMAIL_RE = re.compile(rf"{_EMAIL_PART}@{_EMAIL_PART}\.{_EMAIL_PART}") - - -class GdriveP2PFolder(BaseModel): - """Folder for peer communication: syft_datasite#version#datasite_email#inbox|outbox#peer_email""" - - datasite_email: str - folder_type: str # "inbox" or "outbox" - peer_email: str - - @field_validator("datasite_email", "peer_email") - @classmethod - def _check_email(cls, v: str) -> str: - if not _FOLDER_EMAIL_RE.fullmatch(v): - raise ValueError(f"Invalid email in P2P folder name: {v!r}") - return v - - def as_string(self) -> str: - return f"{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#{SYFT_VERSION}#{self.datasite_email}#{self.folder_type}#{self.peer_email}" - - @classmethod - def from_name(cls, name: str) -> "GdriveP2PFolder": - parts = name.split("#") - if len(parts) != 5 or parts[0] != GDRIVE_P2P_FOLDER_DATASITE_PREFIX: - raise ValueError(f"Invalid P2P folder name: {name}") - return cls(datasite_email=parts[2], folder_type=parts[3], peer_email=parts[4]) - - -class GdriveEncryptionBundlesFolder(BaseModel): - """Folder for encryption bundles: syft_encryption_bundles#email""" - - email: str - - def as_string(self) -> str: - return f"syft_encryption_bundles#{self.email}" - - -class GdrivePersonalSyftboxFolder(BaseModel): - """Folder for personal SyftBox: {version}#{email} under /SyftBox/""" - - email: str - - def as_string(self) -> str: - return f"{SYFT_VERSION}#{self.email}" - - -def collection_name_query(prefix: str) -> str: - """A Drive query that finds a collection of ``prefix`` in any layout. - - It has no trailing '_', because a layout writes its variant in that - position. A client that predates multi-layout searches with the '_', and so - never lists a layout it cannot read. - """ - return f"name contains '{prefix}'" - - -@lru_cache(maxsize=None) -def _collection_name_re(prefix: str) -> "re.Pattern[str]": - """Matches '{prefix}{variant}_{tag}_{hash}'. - - The variant and the hash hold no '_'. The tag can hold one, so it takes - every character in between. - """ - return re.compile( - rf"^{re.escape(prefix)}" - r"(?P[^_]*)_(?P.+)_(?P[^_]+)$" - ) - - -class CollectionFolder(BaseModel): - """Naming value object for collection folders (domain-agnostic). - - Collections (datasets, private datasets, or any future collection type) are - named ``{prefix}{variant}_{tag}_{content_hash}`` on the wire. The variant - names the layout (see CollectionSyncSpec) and is "" for the original one, so - that name is unchanged byte for byte. This is the single source of truth for - building and parsing the name; it also computes the content hash from the - folder's file contents. - - A writer normally folds the variant into ``prefix`` (the spec builds the - qualified prefix once) and leaves ``variant`` empty; ``from_name`` splits - them apart again for a reader. - """ - - prefix: str - tag: str - content_hash: str - variant: str = "" - - @property - def folder_name(self) -> str: - """The on-wire folder name. Doubles as the cache key (same string).""" - return f"{self.prefix}{self.variant}_{self.tag}_{self.content_hash}" - - @classmethod - def from_name(cls, prefix: str, name: str) -> "CollectionFolder": - """Parse '{prefix}{variant}_{tag}_{hash}' into a CollectionFolder. - - Raises ValueError if the name does not carry the prefix, so callers can - skip non-matching folders. - """ - match = _collection_name_re(prefix).match(name) - if match is None: - raise ValueError(f"Invalid collection folder name: {name}") - return cls(prefix=prefix, **match.groupdict()) - - @staticmethod - def compute_hash(files: dict[str, bytes]) -> str: - """Compute a content hash from file contents.""" - from syft.sync.file_utils import compute_file_hashes - - return compute_file_hashes(files) - - -# Helpers for finding folders whose names embed SYFT_VERSION. Folder -# names use '#' or '-' as field separators with the version as one field, -# so we walk those fields looking for an X.Y.Z-shaped chunk -- no per-format -# parser needed. Patch-compat means same major.minor (matches the protocol -# definition in version_info.is_compatible_with). - - -def _looks_like_version(s: str) -> bool: - """True if s is shaped like 'X.Y.Z' with all numeric parts.""" - parts = s.split(".") - return len(parts) == 3 and all(p.isdigit() for p in parts) - - -def _extract_version_from_name(name: str) -> str | None: - """Return the first '#'/'-'-separated field that looks like a semver.""" - for chunk in name.replace("-", "#").split("#"): - if _looks_like_version(chunk): - return chunk - return None - - -# A folder id and name, with the version from the name. The version fields come -# first, so the default sort puts these in version order. -_VersionedFolder = tuple[int, int, int, str, str] - - -def _sorted_by_version(folders: list[tuple[str, str]]) -> list[tuple[str, str]]: - """Folders from the lowest version to the highest. - - A name with no readable version sorts first, so a versioned folder always - wins when the caller takes the last entry. - """ - - def key(entry: tuple[str, str]) -> tuple[int, int, int]: - version_str = _extract_version_from_name(entry[1]) - if version_str is None: - return (-1, -1, -1) - try: - return _parse_semver(version_str) - except ValueError: - return (-1, -1, -1) - - return sorted(folders, key=key) - - -def _partition_by_version( - folders: list[tuple[str, str]], - current_version: str | None = None, -) -> tuple[list[tuple[str, str]], list[tuple[str, str]], list[tuple[str, str]]]: - """Split folders into (compatible, older, newer) by the version in the name. - - Compatible means the same major and minor as the current version. The function - drops a folder that has no version in its name. Each list starts at the lowest - version. - """ - if current_version is None: - current_version = SYFT_VERSION - try: - current = _parse_semver(current_version) - except ValueError: - return [], [], [] - - compatible: list[_VersionedFolder] = [] - older: list[_VersionedFolder] = [] - newer: list[_VersionedFolder] = [] - for fid, name in folders: - version_str = _extract_version_from_name(name) - if version_str is None: - continue - try: - found = _parse_semver(version_str) - except ValueError: - continue - entry = (*found, fid, name) - if found[:2] == current[:2]: - compatible.append(entry) - elif found < current: - older.append(entry) - else: - newer.append(entry) - - def _ordered(entries: list[_VersionedFolder]) -> list[tuple[str, str]]: - return [(fid, name) for *_, fid, name in sorted(entries)] - - return _ordered(compatible), _ordered(older), _ordered(newer) - - -class GDriveConnection(SyftboxPlatformConnection): - """Google Drive Files API transport layer""" - - class Config: - arbitrary_types_allowed = True - - drive_service: Any = None - credentials: GoogleCredentials | None = None - verbose: bool = True - email: str - token_path: Path | None = None - _is_setup: bool = False - - # /SyftBox - # this is the toplevel folder with inboxes, outboxes and personal syftbox - _syftbox_folder_id: str | None = None - - # /SyftBox/myemail - # this is where we store the personal data - _personal_syftbox_folder_id: str | None = None - - # peer_email -> folder_id (folders I created for peer's datasite) - peer_datasite_inbox_cache: dict[str, str] = {} - peer_datasite_outbox_cache: dict[str, str] = {} - - # peer_email -> folder_id (folders peer created for my datasite) - own_datasite_inbox_cache: dict[str, str] = {} - own_datasite_outbox_cache: dict[str, str] = {} - - # sender email -> archive folder id - archive_folder_id_cache: dict[str, str] = {} - - # fname -> gdrive id - personal_syftbox_event_id_cache: dict[str, str] = {} - - # tag -> dataset collection folder id - collection_folder_id_cache: Dict[str, str] = {} - - # Rolling state caches for single-API-call optimization - _rolling_state_folder_id: str | None = None - _rolling_state_file_id: str | None = None - - # Encryption bundles folder cache - _encryption_bundles_folder_id: str | None = None - - # Cached SYFT_peers.json contents (None = not loaded yet). - _peers_json_cache: dict[str, dict[str, Any]] | None = None - - @classmethod - def from_config(cls, config: "GdriveConnectionConfig") -> "GDriveConnection": - return cls.from_token_path(config.email, config.token_path) - - @classmethod - def from_token_path( - cls, - email: str, - token_path: Path | None, - warm_syftbox_folder_id_cache: bool = True, - ) -> "GDriveConnection": - res = cls(email=email, token_path=token_path) - if token_path: - credentials = GoogleCredentials.from_authorized_user_file( - token_path, SCOPES - ) - else: - credentials = None - res.setup( - credentials=credentials, - warm_syftbox_folder_id_cache=warm_syftbox_folder_id_cache, - ) - return res - - @classmethod - def from_service( - cls, - email: str, - mock_service: Any, - warm_syftbox_folder_id_cache: bool = True, - ) -> "GDriveConnection": - """Create a GDriveConnection using a mock drive service for testing. - - Args: - email: Email of the user - mock_service: MockDriveService instance to use instead of real API - warm_syftbox_folder_id_cache: See ``setup``. - - Returns: - GDriveConnection configured with the mock service - """ - from syft.sync.connections.drive.mock_drive_service import ( - MockDriveService, - ) - - res = cls(email=email, token_path=None) - if isinstance(mock_service, MockDriveService): - mock_service = MockDriveService(mock_service._backing_store, email) - res.setup( - drive_service=mock_service, - warm_syftbox_folder_id_cache=warm_syftbox_folder_id_cache, - ) - return res - - def setup( - self, - credentials: GoogleCredentials | None = None, - drive_service: Any | None = None, - warm_syftbox_folder_id_cache: bool = True, - ): - """Setup Drive transport with OAuth2 credentials, Colab auth, or mock service. - - Args: - credentials: OAuth2 credentials for real API access - drive_service: Drive service instance (e.g., mock for testing) - warm_syftbox_folder_id_cache: If True, eagerly resolve the personal - SyftBox folder id (2-3 API calls). Pass False when caches are - seeded by the caller (e.g. ``copy()``). - """ - self.credentials = credentials - if drive_service is not None: - self.drive_service = drive_service - else: - self.drive_service = build_drive_service( - self.credentials, environment=self.environment - ) - - if warm_syftbox_folder_id_cache: - self.get_personal_syftbox_folder_id() - self._is_setup = True - - def clone_for_thread(self) -> "GDriveConnection": - """A connection safe to use from another thread: the Drive client is not.""" - return self.copy() - - def copy(self) -> "GDriveConnection": - # if is mock - from syft.sync.connections.drive.mock_drive_service import ( - MockDriveService, - ) - - if isinstance(self.drive_service, MockDriveService): - new_conn = GDriveConnection.from_service( - self.email, self.drive_service, warm_syftbox_folder_id_cache=False - ) - else: - new_conn = GDriveConnection.from_token_path( - self.email, self.token_path, warm_syftbox_folder_id_cache=False - ) - self._copy_caches_to(new_conn) - return new_conn - - def _copy_caches_to(self, other: "GDriveConnection") -> None: - """Seed all in-memory caches on `other` from `self`, so a freshly-built - copy doesn't re-discover folder ids via Drive search on first use.""" - other._syftbox_folder_id = self._syftbox_folder_id - other._personal_syftbox_folder_id = self._personal_syftbox_folder_id - other._rolling_state_folder_id = self._rolling_state_folder_id - other._rolling_state_file_id = self._rolling_state_file_id - other._encryption_bundles_folder_id = self._encryption_bundles_folder_id - other._peers_json_cache = ( - dict(self._peers_json_cache) if self._peers_json_cache is not None else None - ) - other.peer_datasite_inbox_cache = dict(self.peer_datasite_inbox_cache) - other.peer_datasite_outbox_cache = dict(self.peer_datasite_outbox_cache) - other.own_datasite_inbox_cache = dict(self.own_datasite_inbox_cache) - other.own_datasite_outbox_cache = dict(self.own_datasite_outbox_cache) - other.archive_folder_id_cache = dict(self.archive_folder_id_cache) - other.personal_syftbox_event_id_cache = dict( - self.personal_syftbox_event_id_cache - ) - other.collection_folder_id_cache = dict(self.collection_folder_id_cache) - - @property - def environment(self) -> Environment: - return check_env() - - def get_authenticated_email(self) -> str: - """Return the email of the Google account behind drive_service.""" - about = execute_with_retries( - self.drive_service.about().get(fields="user(emailAddress)") - ) - return about["user"]["emailAddress"] - - def create_personal_syftbox_folder(self) -> str: - """Creates /SyftBox/{version}#{email}""" - syftbox_folder_id = self.get_syftbox_folder_id() - folder_name = GdrivePersonalSyftboxFolder(email=self.email).as_string() - return self.create_folder(folder_name, syftbox_folder_id) - - def create_syftbox_folder(self) -> str: - """Creates /SyftBox""" - folder = self._create_syft_file( - {"name": SYFTBOX_FOLDER, "mimeType": GOOGLE_FOLDER_MIME_TYPE} - ) - return folder["id"] - - def create_archive_folder(self, sender_email: str) -> str: - archive_folder = GdriveArchiveFolder( - sender_email=sender_email, recipient_email=self.email - ) - archive_folder_name = archive_folder.as_string() - syftbox_folder_id = self.get_syftbox_folder_id() - return self.create_folder(archive_folder_name, syftbox_folder_id) - - def create_peer_datasite_folders(self, peer_email: str): - """Create inbox + outbox folders for peer's datasite on own drive, share with peer. - - Creates: - - syft_datasite_{peer}_inbox_{self} — I submit proposals to peer here - - syft_datasite_{peer}_outbox_{self} — peer pushes events to me here - """ - # Inbox: I write proposals, peer reads - inbox_id = self._get_peer_datasite_inbox_id(peer_email) - if inbox_id is None: - inbox_id = self._create_peer_datasite_folder(peer_email, "inbox") - self.add_permission(inbox_id, peer_email, write=True) - - # Outbox: peer writes events, I read - outbox_id = self._get_peer_datasite_outbox_id(peer_email) - if outbox_id is None: - outbox_id = self._create_peer_datasite_folder(peer_email, "outbox") - self.add_permission(outbox_id, peer_email, write=True) - - def add_peer(self, peer_email: str): - """Alias for create_peer_datasite_folders.""" - self.create_peer_datasite_folders(peer_email) - - def _get_peers_file_id(self) -> str | None: - """Find SYFT_peers.json file in /SyftBox folder""" - syftbox_folder_id = self.get_syftbox_folder_id() - query = f"name='{SYFT_PEERS_FILE}' and '{syftbox_folder_id}' in parents and trashed=false" - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - return items[0]["id"] if items else None - - def _download_peers_json(self) -> dict[str, dict[str, Any]]: - """Fetch peers JSON from GDrive. Returns empty dict if not found. - - Maps each peer email to ``{"state": str, "public_encryption_bundle": - dict}``, and ``_meta`` to ``{"version": int}``. - - Raises: - TransportUnavailableError: when Drive cannot be reached. - ValueError: when the file is not valid JSON. - """ - with raise_unavailable("Could not read the peers file from Drive."): - file_id = self._get_peers_file_id() - if file_id is None: - return {} - file_data = self.download_file(file_id) - try: - return json.loads(file_data.decode("utf-8")) - except ValueError as e: - raise ValueError(f"Could not read the peers file: {e}") from e - - def _get_peers_json( - self, force_download: bool = False - ) -> dict[str, dict[str, Any]]: - """Return peers JSON, using the in-memory cache when available.""" - if self._peers_json_cache is not None and not force_download: - return self._peers_json_cache - self._peers_json_cache = self._download_peers_json() - return self._peers_json_cache - - def _write_peers_json(self, peers_data: dict[str, dict[str, Any]]): - """Write peers JSON to GDrive. Creates or updates the file.""" - peers_data = { - **peers_data, - PEERS_META_KEY: {"version": SYFT_PEERS_VERSION}, - } - syftbox_folder_id = self.get_syftbox_folder_id() - file_id = self._get_peers_file_id() - - # Convert to JSON bytes - json_data = json.dumps(peers_data, indent=2) - file_payload, _ = self.create_file_payload(json_data) - - if file_id is None: - # Create new file - file_metadata = { - "name": SYFT_PEERS_FILE, - "parents": [syftbox_folder_id], - } - result = self._create_syft_file(file_metadata, file_payload) - self._peers_json_cache = peers_data - return result.get("id") - else: - # Update existing file - execute_with_retries( - self.drive_service.files().update( - fileId=file_id, media_body=file_payload - ) - ) - self._peers_json_cache = peers_data - return file_id - - def _update_peer_state( - self, - peer_email: str, - state: str, - public_encryption_bundle: dict | None = None, - ): - """Update a single peer's state in the JSON file. - - Preserves existing fields when updating state. - """ - peers_data = self._get_peers_json() - existing = peers_data.get(peer_email, {}) - existing["state"] = state - if public_encryption_bundle is not None: - existing["public_encryption_bundle"] = public_encryption_bundle - peers_data[peer_email] = existing - self._write_peers_json(peers_data) - - def get_peer_requests(self) -> list[str]: - """Get list of pending peer requests. - - Scans for syft_datasite_#version#{self}_* folders NOT owned by self — those are - peers who created folders for our datasite. Filters out already-accepted - or rejected peers from SYFT_peers.json. - """ - results = execute_with_retries( - self.drive_service.files().list( - q=f"name contains '{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#' " - f"and name contains '#{self.email}#' " - f"and trashed=false " - f"and mimeType = '{GOOGLE_FOLDER_MIME_TYPE}' " - f"and not 'me' in owners", - fields="files(id,name,owners)", - ) - ) - - all_folder_peers = set() - for f in results.get("files", []): - try: - folder = GdriveP2PFolder.from_name(f["name"]) - except ValueError: - # The query matches a name prefix, so a folder with another shape - # can appear here. - continue - # Any account can make a folder with this name. The request is from - # the owner, so the name must name the owner. - owners = f.get("owners", []) - if ( - len(owners) != 1 - or owners[0].get("emailAddress", "").casefold() - != folder.peer_email.casefold() - ): - continue - if folder.datasite_email == self.email: - all_folder_peers.add(folder.peer_email) - - peers_data = self._get_peers_json() - pending_peers = [] - for peer_email in all_folder_peers: - if peer_email not in peers_data: - pending_peers.append(peer_email) - elif peers_data[peer_email].get("state") not in ["accepted", "rejected"]: - pending_peers.append(peer_email) - - return pending_peers - - def watcher_download_raw_events_from_outbox( - self, peer_email: str, since_timestamp: float | None - ) -> list[bytes]: - folder_id = self._get_peer_datasite_outbox_id(peer_email) - if folder_id is None: - return [] - - file_metadatas = self.get_file_metadatas_from_folder( - folder_id, since_timestamp=since_timestamp - ) - valid_fname_objs = self._get_valid_events_from_file_metadatas(file_metadatas) - name_to_id = {f["name"]: f["id"] for f in file_metadatas} - - sorted_fname_objs = [ - x - for x in sorted(valid_fname_objs, key=lambda x: x.timestamp) - if since_timestamp is None or x.timestamp > since_timestamp - ] - - if len(sorted_fname_objs) == 0: - return [] - - res = [] - for fname_obj in sorted_fname_objs: - file_name = fname_obj.as_string() - if file_name in name_to_id: - res.append(self.download_file(name_to_id[file_name])) - return res - - def watcher_get_events_messages( - self, peer_email: str, since_timestamp: float | None - ) -> list[FileChangeEventsMessage]: - raw_list = self.watcher_download_raw_events_from_outbox( - peer_email, since_timestamp - ) - return [FileChangeEventsMessage.from_compressed_data(data) for data in raw_list] - - def watcher_get_outbox_file_metadatas( - self, peer_email: str, since_timestamp: float | None - ) -> list[dict]: - """Get file metadata from peer's outbox folder without downloading.""" - folder_id = self._get_peer_datasite_outbox_id(peer_email) - if folder_id is None: - return [] - - file_metadatas = self.get_file_metadatas_from_folder( - folder_id, since_timestamp=since_timestamp - ) - valid_fname_objs = self._get_valid_events_from_file_metadatas(file_metadatas) - name_to_id = {f["name"]: f["id"] for f in file_metadatas} - - result = [] - for fname_obj in sorted(valid_fname_objs, key=lambda x: x.timestamp): - if since_timestamp is None or fname_obj.timestamp > since_timestamp: - file_name = fname_obj.as_string() - if file_name in name_to_id: - result.append( - { - "file_id": name_to_id[file_name], - "file_name": file_name, - "timestamp": fname_obj.timestamp, - } - ) - return result - - def owner_write_raw_bytes_to_syftbox(self, filename: str, data: bytes) -> str: - """Write raw bytes to /SyftBox/myemail.""" - personal_syftbox_folder_id = self.get_personal_syftbox_folder_id() - file_metadata = { - "name": filename, - "parents": [personal_syftbox_folder_id], - } - file_payload, _ = self.create_file_payload(data) - - res = self._create_syft_file(file_metadata, file_payload) - gdrive_id = res.get("id") - self.personal_syftbox_event_id_cache[filename] = gdrive_id - return gdrive_id - - def owner_download_raw_bytes_by_id(self, file_id: str) -> bytes: - """Download raw bytes by file ID.""" - return self.download_file(file_id) - - def owner_get_all_accepted_event_file_ids( - self, since_timestamp: float | None = None - ) -> list[str]: - personal_syftbox_folder_id = self.get_personal_syftbox_folder_id() - file_metadatas = self.get_file_metadatas_from_folder( - personal_syftbox_folder_id, since_timestamp=since_timestamp - ) - valid_fname_objs = self._filter_valid_file_metadatas(file_metadatas) - oldest_first = sorted( - valid_fname_objs, - key=lambda f: ( - FileChangeEventsMessageFileName.from_string(f["name"]).timestamp - ), - ) - return [f["id"] for f in oldest_first] - - def owner_download_all_raw_events_from_syftbox(self) -> list[bytes]: - """Download all event files from /SyftBox/myemail as raw bytes.""" - personal_syftbox_folder_id = self.get_personal_syftbox_folder_id() - file_metadatas = self.get_file_metadatas_from_folder(personal_syftbox_folder_id) - valid_fname_objs = self._get_valid_events_from_file_metadatas(file_metadatas) - - result = [] - for fname_obj in sorted(valid_fname_objs, key=lambda x: x.timestamp): - gdrive_id = [ - f for f in file_metadatas if f["name"] == fname_obj.as_string() - ][0]["id"] - try: - file_data = self.download_file(gdrive_id) - except Exception as e: - print(f"Warning: could not download event {fname_obj.as_string()}: {e}") - continue - result.append(file_data) - return result - - def owner_write_raw_bytes_to_outbox( - self, recipient: str, filename: str, data: bytes - ) -> None: - outbox_folder_id = self._get_own_datasite_outbox_id(recipient) - if outbox_folder_id is None: - raise ValueError(f"Outbox folder for {recipient} not found") - - file_payload, _ = self.create_file_payload(data) - file_metadata = {"name": filename, "parents": [outbox_folder_id]} - - result = self._create_syft_file( - file_metadata, file_payload, fields="id, parents" - ) - - file_id = result.get("id") - actual_parents = result.get("parents", []) - if outbox_folder_id not in actual_parents: - print( - f"WARNING: Event message {file_id} was not placed in outbox folder " - f"{outbox_folder_id}. Actual parents: {actual_parents}. " - f"Moving file to correct folder..." - ) - execute_with_retries( - self.drive_service.files().update( - fileId=file_id, - addParents=outbox_folder_id, - fields="id, parents", - ) - ) - - def owner_write_event_messages_to_outbox( - self, recipient: str, events_message: FileChangeEventsMessage - ): - data = events_message.as_compressed_data() - fname = events_message.message_filepath.as_string() - self.owner_write_raw_bytes_to_outbox(recipient, fname, data) - - def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: - """Move the inbox file ``file_id`` from ``sender_email`` to its archive folder. - - ``sender_email`` names the archive folder, so the caller reads it from - the transport, never from the sender-supplied message body. - """ - execute_with_retries(self._archive_inbox_file_request(file_id, sender_email)) - - def owner_archive_inbox_files(self, files: list[tuple[str, str]]) -> None: - """Move each inbox file, given as ``(file_id, sender_email)``, to its - sender's archive folder, up to ``DRIVE_BATCH_SIZE`` moves per request. - - A move that fails is warned about and its file stays in the inbox. - """ - - def callback(request_id, response, exception): - if exception is not None: - warnings.warn( - f"Could not archive inbox file {request_id}: {exception!r}" - ) - - for i in range(0, len(files), DRIVE_BATCH_SIZE): - batch = self.drive_service.new_batch_http_request(callback=callback) - for file_id, sender_email in files[i : i + DRIVE_BATCH_SIZE]: - request = self._archive_inbox_file_request(file_id, sender_email) - batch.add(request, request_id=file_id) - batch_execute_with_retries(batch) - - def _archive_inbox_file_request(self, file_id: str, sender_email: str): - """The request that moves an inbox file to its archive folder. The inbox - folder is the file's only parent and both folder ids are cached, so the - move is one Drive call.""" - return self.drive_service.files().update( - fileId=file_id, - addParents=self.owner_get_archive_folder_id(sender_email), - removeParents=self._get_own_datasite_inbox_id(sender_email), - fields="id", - supportsAllDrives=True, - ) - - def _has_permission(self, file_id: str, email: str) -> bool: - """Check if user already has permission on the file.""" - perms = execute_with_retries( - self.drive_service.permissions().list( - fileId=file_id, fields="permissions(emailAddress)" - ) - ) - for p in perms.get("permissions", []): - if p.get("emailAddress", "").lower() == email.lower(): - return True - return False - - def add_permission(self, file_id: str, recipient: str, write=False): - """Add permission to the file if not already shared.""" - if self._has_permission(file_id, recipient): - return - - role = "writer" if write else "reader" - permission = { - "type": "user", - "role": role, - "emailAddress": recipient, - } - execute_with_retries( - self.drive_service.permissions().create( - fileId=file_id, body=permission, sendNotificationEmail=True - ) - ) - - def _create_peer_datasite_folder(self, peer_email: str, folder_type: str) -> str: - """Create a datasite folder under /SyftBox for peer's datasite.""" - if folder_type not in ("inbox", "outbox"): - raise ValueError( - f"Invalid folder_type: {folder_type}. Must be 'inbox' or 'outbox'." - ) - parent_id = self.get_syftbox_folder_id() - folder = GdriveP2PFolder( - datasite_email=peer_email, folder_type=folder_type, peer_email=self.email - ) - folder_name = folder.as_string() - folder_id = self.create_folder(folder_name, parent_id) - if folder_type == "inbox": - self.peer_datasite_inbox_cache[peer_email] = folder_id - else: - self.peer_datasite_outbox_cache[peer_email] = folder_id - return folder_id - - def get_personal_syftbox_folder_id(self) -> str: - """/SyftBox/{version}#{email}""" - if self._personal_syftbox_folder_id: - return self._personal_syftbox_folder_id - folders = self._find_folders( - name_contains=[f"#{self.email}"], - parent_id=self.get_syftbox_folder_id(), - owner_email=self.email, - ) - # The substring '#{email}' also matches p2p folder names that end in - # '#{peer}#{type}#{email}'. Personal folder shape is exactly - # '{version}#{email}', so require a single '#'. - folders = [(fid, name) for fid, name in folders if name.count("#") == 1] - folder_id = self._find_or_adopt_versioned_folder( - folders, - current_name=GdrivePersonalSyftboxFolder(email=self.email).as_string(), - ) - if folder_id: - self._personal_syftbox_folder_id = folder_id - return folder_id - return self.create_personal_syftbox_folder() - - def get_syftbox_folder_id(self) -> str: - """/SyftBox""" - # cached - if self._syftbox_folder_id: - return self._syftbox_folder_id - else: - syftbox_folder_id = self.get_syftbox_folder_id_from_drive() - if syftbox_folder_id: - self._syftbox_folder_id = syftbox_folder_id - return self._syftbox_folder_id - else: - return self.create_syftbox_folder() - - def get_archive_folder_id_from_drive(self, sender_email: str) -> str | None: - archive_folder = GdriveArchiveFolder( - sender_email=sender_email, recipient_email=self.email - ) - archive_folder_name = archive_folder.as_string() - query = f"name='{archive_folder_name}' and mimeType='application/vnd.google-apps.folder' and 'me' in owners and trashed=false" - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - return items[0]["id"] if items else None - - def owner_get_archive_folder_id(self, sender_email: str) -> str: - if sender_email not in self.archive_folder_id_cache: - archive_folder_id = self.get_archive_folder_id_from_drive( - sender_email - ) or self.create_archive_folder(sender_email) - self.archive_folder_id_cache[sender_email] = archive_folder_id - return self.archive_folder_id_cache[sender_email] - - @staticmethod - def _extract_timestamp_from_filename(filename: str) -> float | None: - """The send time in a message file name. See ``naming.py``.""" - return timestamp_from_message_name(filename) - - def get_file_metadatas_from_folder( - self, - folder_id: str, - since_timestamp: float | None = None, - page_size: int = 100, - ) -> list[dict]: - """ - Get file metadatas from folder with early termination. - - Args: - folder_id: Google Drive folder ID - since_timestamp: Optional timestamp. If provided, only files with a - newer timestamp are returned, and pagination stops at - a file from an earlier second. - page_size: Number of files to fetch per API call. Default 100. - - Returns: - List of file metadata dicts, sorted by name descending (newest first) - """ - query = f"'{folder_id}' in parents and trashed=false" - all_files = [] - page_token = None - - while True: - results = execute_with_retries( - self.drive_service.files().list( - q=query, - fields="files(id, name, size, mimeType, modifiedTime), nextPageToken", - pageSize=page_size, - pageToken=page_token, - orderBy="name desc", - ) - ) - - page_files = results.get("files", []) - - # Early termination: Check if this page contains old files - if since_timestamp is not None and page_files: - should_stop = False - - for file in page_files: - timestamp = self._extract_timestamp_from_filename(file["name"]) - - if timestamp is not None: - if timestamp > since_timestamp: - all_files.append(file) - elif int(timestamp) < int(since_timestamp): - # Drive sorts names as strings, so a 0.9 fraction of - # a second lists before 0.95. We stop when the - # second flips. - should_stop = True - break - else: - # No timestamp in filename, include the file - all_files.append(file) - - if should_stop: - # Don't fetch more pages - break - else: - # No early termination check, add all files - all_files.extend(page_files) - - # Check for next page - page_token = results.get("nextPageToken") - if not page_token: - break - - return all_files - - @staticmethod - def _filter_valid_file_metadatas( - file_metadatas: list[dict], - ) -> list[dict]: - res = [] - for file_metadata in file_metadatas: - fname = file_metadata["name"] - try: - FileChangeEventsMessageFileName.from_string(fname) - except ValueError: - # The folder holds other files, so a name that is not an event - # name is normal here. This method filters them out. - continue - res.append(file_metadata) - return res - - @staticmethod - def _get_valid_events_from_file_metadatas( - file_metadatas: list[dict], - ) -> list[FileChangeEventsMessageFileName]: - res = [] - for file_metadata in file_metadatas: - fname = file_metadata["name"] - try: - message_filename = FileChangeEventsMessageFileName.from_string(fname) - except ValueError: - print(f"Warning: invalid event file name: {fname}") - continue - res.append(message_filename) - return res - - @staticmethod - def _get_valid_messages_from_file_metadatas( - file_metadatas: list[dict], - ) -> list[MessageFileName]: - res = [] - for file_metadata in file_metadatas: - try: - message_filename = MessageFileName.from_string(file_metadata["name"]) - res.append(message_filename) - except FileNameParseError: - continue - return res - - def owner_get_inbox_file_metadatas(self, sender_email: str) -> list[dict]: - """The messages in the inbox from ``sender_email``, oldest first, each - with ``file_id`` and ``file_name``, without downloading them.""" - inbox_folder_id = self._get_own_datasite_inbox_id(sender_email) - if inbox_folder_id is None: - raise ValueError(f"Inbox folder for {sender_email} not found") - file_metadatas = self.get_file_metadatas_from_folder(inbox_folder_id) - valid_file_names = self._get_valid_messages_from_file_metadatas(file_metadatas) - name_to_id = {f["name"]: f["id"] for f in file_metadatas} - return [ - {"file_id": name_to_id[name.as_string()], "file_name": name.as_string()} - for name in sorted(valid_file_names, key=lambda x: x.submitted_timestamp) - ] - - def _find_p2p_folder_id( - self, - datasite_email: str, - folder_type: str, - peer_email: str, - owner_email: str, - ) -> str | None: - folders = self._find_folders( - name_contains=[ - f"{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#", - f"#{datasite_email}#{folder_type}#{peer_email}", - ], - owner_email=owner_email, - ) - - # Drive's `name contains` is a fuzzy token/prefix match, not a substring - # match, so it can return folders for other peers when email tokens - # collide - def _is_exact_match(name: str) -> bool: - try: - folder = GdriveP2PFolder.from_name(name) - except ValueError: - return False - return ( - folder.datasite_email == datasite_email - and folder.folder_type == folder_type - and folder.peer_email == peer_email - ) - - # Ignore the version in the name. Each peer builds this name from its own - # client version, so a filter here hides the folder that the peer uses. - # After an upgrade the client therefore finds the old folder and writes to - # it. It makes no second folder, which an older peer would never look for. - candidates = _sorted_by_version( - [(fid, name) for fid, name in folders if _is_exact_match(name)] - ) - if not candidates: - return None - if len(candidates) > 1: - print( - f"Warning: {len(candidates)} P2P folders for {datasite_email} " - f"{folder_type} {peer_email}; using {candidates[-1][1]}" - ) - return candidates[-1][0] - - def _get_peer_datasite_inbox_id(self, peer_email: str) -> str | None: - """Get folder: syft_datasite_{peer}_inbox_{self}, owned by self.""" - if peer_email in self.peer_datasite_inbox_cache: - return self.peer_datasite_inbox_cache[peer_email] - folder_id = self._find_p2p_folder_id( - datasite_email=peer_email, - folder_type="inbox", - peer_email=self.email, - owner_email=self.email, - ) - if folder_id is not None: - self.peer_datasite_inbox_cache[peer_email] = folder_id - return folder_id - - def _get_peer_datasite_outbox_id(self, peer_email: str) -> str | None: - """Get folder: syft_datasite_{peer}_outbox_{self}, owned by self.""" - if peer_email in self.peer_datasite_outbox_cache: - return self.peer_datasite_outbox_cache[peer_email] - folder_id = self._find_p2p_folder_id( - datasite_email=peer_email, - folder_type="outbox", - peer_email=self.email, - owner_email=self.email, - ) - if folder_id is not None: - self.peer_datasite_outbox_cache[peer_email] = folder_id - return folder_id - - def _get_own_datasite_inbox_id(self, peer_email: str) -> str | None: - """Get folder: syft_datasite_{self}_inbox_{peer}, owned by peer.""" - if peer_email in self.own_datasite_inbox_cache: - return self.own_datasite_inbox_cache[peer_email] - folder_id = self._find_p2p_folder_id( - datasite_email=self.email, - folder_type="inbox", - peer_email=peer_email, - owner_email=peer_email, - ) - if folder_id is not None: - self.own_datasite_inbox_cache[peer_email] = folder_id - return folder_id - - def _get_own_datasite_outbox_id(self, peer_email: str) -> str | None: - """Get folder: syft_datasite_{self}_outbox_{peer}, owned by peer.""" - if peer_email in self.own_datasite_outbox_cache: - return self.own_datasite_outbox_cache[peer_email] - folder_id = self._find_p2p_folder_id( - datasite_email=self.email, - folder_type="outbox", - peer_email=peer_email, - owner_email=peer_email, - ) - if folder_id is not None: - self.own_datasite_outbox_cache[peer_email] = folder_id - return folder_id - - def watcher_send_raw_bytes_to_inbox( - self, recipient: str, filename: str, data: bytes - ) -> None: - inbox_id = self._get_peer_datasite_inbox_id(recipient) - if inbox_id is None: - raise Exception(f"Inbox folder for {recipient}'s datasite not found") - - payload, _ = self.create_file_payload(data) - file_metadata = {"name": filename, "parents": [inbox_id]} - - result = self._create_syft_file(file_metadata, payload, fields="id, parents") - - file_id = result.get("id") - actual_parents = result.get("parents", []) - if inbox_id not in actual_parents: - print( - f"WARNING: Message file {file_id} was not placed in inbox folder " - f"{inbox_id}. Actual parents: {actual_parents}. " - f"Moving file to correct folder..." - ) - execute_with_retries( - self.drive_service.files().update( - fileId=file_id, - addParents=inbox_id, - fields="id, parents", - ) - ) - - def watcher_send_proposed_file_changes_message( - self, - recipient: str, - proposed_file_changes_message: ProposedFileChangesMessage, - ): - data = proposed_file_changes_message.as_compressed_data() - filename = proposed_file_changes_message.message_filename.as_string() - self.watcher_send_raw_bytes_to_inbox(recipient, filename, data) - - def reset_caches(self): - self._syftbox_folder_id = None - self._personal_syftbox_folder_id = None - self.peer_datasite_inbox_cache.clear() - self.peer_datasite_outbox_cache.clear() - self.own_datasite_inbox_cache.clear() - self.own_datasite_outbox_cache.clear() - self.archive_folder_id_cache.clear() - self.personal_syftbox_event_id_cache.clear() - self.collection_folder_id_cache.clear() - self._rolling_state_folder_id = None - self._rolling_state_file_id = None - self._encryption_bundles_folder_id = None - self._peers_json_cache = None - - def gather_all_file_and_folder_ids(self) -> list[str]: - syftbox_folder_id = self.get_syftbox_folder_id() - return gather_all_file_and_folder_ids_recursive( - self.drive_service, syftbox_folder_id - ) - - def delete_multiple_files_by_ids( - self, - file_ids: list[str], - ignore_permissions_errors: bool = True, - ignore_file_not_found: bool = True, - ): - def callback(request_id, response, exception): - if exception: - exception_str = str(exception) - # insufficientFilePermissions is a common error when deleting files that may already be removed - if ( - ignore_permissions_errors - and "insufficientFilePermissions" in exception_str - ): - return - # 404 errors occur when files are already deleted - if ignore_file_not_found and ( - "404" in exception_str or "notFound" in exception_str - ): - return - raise exception - if ( - response is not None - and not isinstance(response, str) - and response.get("status") - and int(response.get("status")) >= 400 - ): - raise Exception( - f"Failed to delete {request_id}: error status {response.get('status')}" - ) - - for i in range(0, len(file_ids), DRIVE_BATCH_SIZE): - chunk = file_ids[i : i + DRIVE_BATCH_SIZE] - batch = self.drive_service.new_batch_http_request(callback=callback) - for file_id in chunk: - batch.add(self.drive_service.files().delete(fileId=file_id)) - batch_execute_with_retries(batch) - - def delete_file_by_id(self, file_id: str, raise_on_error: bool = False): - try: - execute_with_retries(self.drive_service.files().delete(fileId=file_id)) - except Exception as e: - if raise_on_error: - raise e - print(f"Warning: could not delete file {file_id}: {e}") - - def find_orphaned_syft_files(self) -> list[str]: - """Find syft files of this account in its My Drive root, except /SyftBox. - - Drive moves a file to the root of its owner when the file loses its - folder. Only /SyftBox belongs in root. - - Returns list of orphaned file IDs. - """ - query = ( - "'root' in parents and 'me' in owners and trashed=false and " - f"name != '{SYFTBOX_FOLDER}' and " - f"properties has {{ key='{SYFT_FILE_PROPERTY_KEY}' " - f"and value='{SYFT_FILE_PROPERTY_VALUE}' }}" - ) - file_ids = [] - page_token = None - while True: - results = execute_with_retries( - self.drive_service.files().list( - q=query, fields="files(id), nextPageToken", pageToken=page_token - ) - ) - file_ids.extend(item["id"] for item in results["files"]) - page_token = results.get("nextPageToken") - if not page_token: - return file_ids - - def create_file_payload(self, data: Any) -> tuple[MediaIoBaseUpload, str]: - """Create a file payload for the GDrive""" - if isinstance(data, str): - file_data = data.encode("utf-8") - mime_type = "text/plain" - extension = ".txt" - elif isinstance(data, dict): - file_data = json.dumps(data, indent=2).encode("utf-8") - mime_type = "application/json" - extension = ".json" - elif isinstance(data, bytes): - file_data = data - mime_type = "application/octet-stream" - extension = ".bin" - else: - # Pickle for other data types - file_data = pickle.dumps(data) - mime_type = "application/octet-stream" - extension = ".pkl" - - media = MediaIoBaseUpload( - io.BytesIO(file_data), mimetype=mime_type, resumable=True - ) - - return media, extension - - def _find_folder_by_name( - self, folder_name: str, parent_id: str = None, owner_email: str = None - ) -> Optional[str]: - """Find a folder by name, optionally within a specific parent""" - # parent_id = "1AQ3WLnVlLd6Zjo7p9Z_qGA1Djjf6-KIh" - owner_email_clause = f"and '{owner_email}' in owners" if owner_email else "" - parent_id_clause = f"and '{parent_id}' in parents" if parent_id else "" - query = f"name='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false {owner_email_clause} {parent_id_clause}" - - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id)", pageSize=1) - ) - items = results.get("files", []) - return items[0]["id"] if items else None - - def _find_folders( - self, - *, - name_contains: list[str] = (), - parent_id: str | None = None, - owner_email: str | None = None, - ) -> list[tuple[str, str]]: - """List folders matching the constraints. Returns (id, name) pairs. - - Thin wrapper over Drive's files.list -- handles query building and - pagination, knows nothing about versions. Pair with - _partition_by_version or _sorted_by_version when the caller cares about - the version in the folder name. - """ - clauses = [f"mimeType='{GOOGLE_FOLDER_MIME_TYPE}'", "trashed=false"] - for substr in name_contains: - clauses.append(f"name contains '{substr}'") - if parent_id: - clauses.append(f"'{parent_id}' in parents") - if owner_email: - clauses.append(f"'{owner_email}' in owners") - query = " and ".join(clauses) - - out: list[tuple[str, str]] = [] - page_token = None - while True: - page = execute_with_retries( - self.drive_service.files().list( - q=query, - fields="files(id, name), nextPageToken", - pageToken=page_token, - ) - ) - out.extend((f["id"], f["name"]) for f in page.get("files", [])) - page_token = page.get("nextPageToken") - if not page_token: - break - return out - - def _expect_one(self, folders: list[tuple[str, str]]) -> str | None: - """Return the single folder id, None if empty, raise if more than one. - - Auto-picking from multiple matches risks orphaning data, so the user - must reconcile on Drive themselves. - """ - if not folders: - return None - if len(folders) == 1: - return folders[0][0] - names = [n for _, n in folders] - raise RuntimeError( - f"Found {len(folders)} compatible folders on Drive: {names}. " - f"Exactly one is expected. This usually indicates a stale folder " - f"left behind by a prior client version. Please delete the stale " - f"folder(s) on Drive (keeping the one with your data) and retry." - ) - - def _find_or_adopt_versioned_folder( - self, - folders: list[tuple[str, str]], - current_name: str, - current_version: str | None = None, - ) -> str | None: - """Return the id of a PRIVATE folder for this client version, or None. - - A private folder name holds the client version, so a minor upgrade looks - for a name that does not exist yet. This method renames the folder of the - highest earlier version to `current_name` and keeps the data. A new folder - would leave the data of the user on Drive and out of reach. - - Renames the folder, so the caller must own it and no peer may look it up by - name. A P2P folder fails both conditions: use `_expect_one` for those. - - Raises RuntimeError if only a folder from a later version exists, or if - more than one compatible folder exists. - """ - compatible, older, newer = _partition_by_version(folders, current_version) - if compatible: - return self._expect_one(compatible) - if newer: - names = [n for _, n in newer] - latest = _extract_version_from_name(names[-1]) - raise RuntimeError( - f"Found a folder from a later client version on Drive: {names}. " - f"This client is {current_version or SYFT_VERSION} and " - f"cannot read that data. Install syft {latest} or later." - ) - if not older: - return None - - folder_id, name = older[-1] - execute_with_retries( - self.drive_service.files().update( - fileId=folder_id, body={"name": current_name} - ) - ) - print(f"Adopted the folder of an earlier version: {name} -> {current_name}") - if len(older) > 1: - stale = [n for _, n in older[:-1]] - print( - f"Warning: {len(stale)} folder(s) of earlier versions stay on " - f"Drive: {stale}" - ) - return folder_id - - def download_file(self, file_id: str) -> bytes: - request = self.drive_service.files().get_media(fileId=file_id) - - file_buffer = io.BytesIO() - downloader = MediaIoBaseDownload( - file_buffer, request, chunksize=1024 * 1024 * 10 - ) - - done = False - while not done: - _, done = next_chunk_with_retries(downloader) - - message_data = file_buffer.getvalue() - return message_data - - def create_folder(self, folder_name: str, parent_id: str) -> str: - folder = self._create_syft_file( - { - "name": folder_name, - "mimeType": GOOGLE_FOLDER_MIME_TYPE, - "parents": [parent_id], - } - ) - return folder["id"] - - def _create_syft_file( - self, - body: dict, - media_body: MediaIoBaseUpload | None = None, - fields: str = "id", - ) -> dict: - """Create a file or folder with the syft tag.""" - return execute_with_retries( - self.drive_service.files().create( - body={ - **body, - "properties": {SYFT_FILE_PROPERTY_KEY: SYFT_FILE_PROPERTY_VALUE}, - }, - media_body=media_body, - fields=fields, - ) - ) - - def get_syftbox_folder_id_from_drive(self) -> str | None: - query = f"name='{SYFTBOX_FOLDER}' and mimeType='application/vnd.google-apps.folder' and 'me' in owners and trashed=false" - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id, name)") - ) - items = results.get("files", []) - return items[0]["id"] if items else None - - # ========================================================================= - # GENERIC COLLECTION METHODS (parameterized on `prefix`) - # - # Folder names are built and parsed via CollectionFolder, the single source - # of truth for the f"{prefix}_{tag}_{content_hash}" scheme. On-wire behavior - # is byte-identical to the previous inline construction. - # ========================================================================= - - def owner_create_collection_folder( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> str: - """Create /SyftBox/{prefix}_{tag}_{hash} folder.""" - folder_name = CollectionFolder( - prefix=prefix, tag=tag, content_hash=content_hash - ).folder_name - - # Check cache - if folder_name in self.collection_folder_id_cache: - return self.collection_folder_id_cache[folder_name] - - syftbox_folder_id = self.get_syftbox_folder_id() - - # Check if exists - folder_id = self._find_folder_by_name(folder_name, parent_id=syftbox_folder_id) - if folder_id: - self.collection_folder_id_cache[folder_name] = folder_id - return folder_id - - # Create new folder - folder_id = self.create_folder(folder_name, syftbox_folder_id) - self.collection_folder_id_cache[folder_name] = folder_id - return folder_id - - def owner_tag_collection_as_any( - self, prefix: str, tag: str, content_hash: str - ) -> None: - """Mark collection as shared with 'any' via appProperties.""" - folder_id = self._get_collection_folder_id(prefix, tag, content_hash) - execute_with_retries( - self.drive_service.files().update( - fileId=folder_id, - body={"appProperties": {"syft_shared_with_any": "true"}}, - ) - ) - - def owner_share_collection( - self, prefix: str, tag: str, content_hash: str, users: list[str] - ) -> None: - """Share collection folder with specific users via batch API.""" - if not users: - return - folder_id = self._get_collection_folder_id(prefix, tag, content_hash) - self._batch_add_permissions(folder_id, users) - - def _batch_add_permissions(self, file_id: str, users: list[str]) -> None: - """Add reader permissions for multiple users in a single batch request.""" - - def callback(request_id, response, exception): - # Ignore "already shared" errors - if exception and "alreadyShared" not in str(exception): - raise exception - - for i in range(0, len(users), DRIVE_BATCH_SIZE): - chunk = users[i : i + DRIVE_BATCH_SIZE] - batch = self.drive_service.new_batch_http_request(callback=callback) - for user_email in chunk: - permission = { - "type": "user", - "role": "reader", - "emailAddress": user_email, - } - batch.add( - self.drive_service.permissions().create( - fileId=file_id, - body=permission, - sendNotificationEmail=True, - ) - ) - batch_execute_with_retries(batch) - - def owner_upload_collection_files( - self, prefix: str, tag: str, content_hash: str, files: dict[str, bytes] - ) -> None: - """Upload files to a collection folder.""" - folder_id = self._get_collection_folder_id(prefix, tag, content_hash) - - for file_path, content in files.items(): - file_payload, _ = self.create_file_payload(content) - file_name = Path(file_path).name - - file_metadata = {"name": file_name, "parents": [folder_id]} - self._create_syft_file(file_metadata, file_payload) - - def owner_list_collections(self, prefix: str) -> list[str]: - """List collections created by DO (owned by me). Returns list of tags. - - A tag appears once however many layouts hold it. - """ - syftbox_folder_id = self.get_syftbox_folder_id() - query = ( - f"{collection_name_query(prefix)} and '{syftbox_folder_id}' in parents " - f"and 'me' in owners and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" - ) - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(name)") - ) - - result: list[str] = [] - for folder in results.get("files", []): - try: - cf = CollectionFolder.from_name(prefix, folder["name"]) - except ValueError: - continue - if cf.tag not in result: - result.append(cf.tag) - return result - - def owner_list_all_collections_with_permissions( - self, - prefix: str, - ) -> list[FileCollection]: - """List all DO's collections with permissions info, one per layout.""" - syftbox_folder_id = self.get_syftbox_folder_id() - query = ( - f"{collection_name_query(prefix)} and '{syftbox_folder_id}' in parents " - f"and 'me' in owners and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" - ) - results = execute_with_retries( - self.drive_service.files().list( - q=query, fields="files(id,name,appProperties)" - ) - ) - - collections = [] - for folder in results.get("files", []): - try: - cf = CollectionFolder.from_name(prefix, folder["name"]) - except ValueError: - continue - has_anyone = ( - folder.get("appProperties", {}).get("syft_shared_with_any") == "true" - ) - collections.append( - FileCollection( - folder_id=folder["id"], - tag=cf.tag, - content_hash=cf.content_hash, - has_any_permission=has_anyone, - variant=cf.variant, - ) - ) - - return collections - - def owner_delete_collection(self, prefix: str, tag: str) -> None: - """Delete every layout of ``tag`` published under ``prefix``.""" - collections = self.owner_list_all_collections_with_permissions(prefix) - for c in collections: - if c.tag == tag: - self.delete_file_by_id(c.folder_id) - folder_name = CollectionFolder( - prefix=prefix, - tag=c.tag, - content_hash=c.content_hash, - variant=c.variant, - ).folder_name - self.collection_folder_id_cache.pop(folder_name, None) - - def watcher_list_collections(self, prefix: str) -> list[dict]: - """List collections shared with DS (not owned by me), one per layout. - - Returns list of dicts with keys: owner_email, tag, content_hash, variant - """ - query = ( - f"{collection_name_query(prefix)} and not 'me' in owners " - f"and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" - ) - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(name, owners)") - ) - - result = [] - for folder in results.get("files", []): - try: - cf = CollectionFolder.from_name(prefix, folder["name"]) - except ValueError: - # Skip folders that don't match the expected format - continue - owner_email = folder.get("owners", [{}])[0].get("emailAddress", "unknown") - result.append( - { - "owner_email": owner_email, - "tag": cf.tag, - "content_hash": cf.content_hash, - "variant": cf.variant, - } - ) - return result - - def watcher_download_collection( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> dict[str, bytes]: - """Download all files from a collection.""" - folder_name = CollectionFolder( - prefix=prefix, tag=tag, content_hash=content_hash - ).folder_name - # Try to find folder by name (could be owned by someone else) - folder_id = self._find_folder_by_name(folder_name, owner_email=owner_email) - - if not folder_id: - raise ValueError(f"Collection {tag} with hash {content_hash} not found") - - file_metadatas = self.get_file_metadatas_from_folder(folder_id) - files = {} - for file_meta in file_metadatas: - file_id = file_meta["id"] - file_name = file_meta["name"] - files[file_name] = self.download_file(file_id) - - return files - - def watcher_get_collection_file_metadatas( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> List[Dict]: - """Get file metadata from a collection without downloading.""" - folder_name = CollectionFolder( - prefix=prefix, tag=tag, content_hash=content_hash - ).folder_name - folder_id = self._find_folder_by_name(folder_name, owner_email=owner_email) - - if not folder_id: - raise ValueError(f"Collection {tag} with hash {content_hash} not found") - - file_metadatas = self.get_file_metadatas_from_folder(folder_id) - return [{"file_id": f["id"], "file_name": f["name"]} for f in file_metadatas] - - def watcher_download_collection_file(self, file_id: str) -> bytes: - """Download a single file from a collection.""" - return self.download_file(file_id) - - def _get_collection_folder_id( - self, prefix: str, tag: str, content_hash: str - ) -> str: - """Get folder ID for a collection, with caching.""" - folder_name = CollectionFolder( - prefix=prefix, tag=tag, content_hash=content_hash - ).folder_name - if folder_name in self.collection_folder_id_cache: - return self.collection_folder_id_cache[folder_name] - - syftbox_folder_id = self.get_syftbox_folder_id() - folder_id = self._find_folder_by_name(folder_name, parent_id=syftbox_folder_id) - - if not folder_id: - raise ValueError( - f"Collection folder {tag} with hash {content_hash} not found" - ) - - self.collection_folder_id_cache[folder_name] = folder_id - return folder_id - - def _get_own_peer_info_file_id(self) -> Optional[str]: - """Find the own peer info file, SYFT_version.json, in /SyftBox.""" - syftbox_folder_id = self.get_syftbox_folder_id() - query = f"name='{SYFT_VERSION_FILE}' and '{syftbox_folder_id}' in parents and trashed=false" - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - return items[0]["id"] if items else None - - def write_peer_info_file(self, peer_info: "PeerInfo") -> None: - """Write the peer info file to /SyftBox. Creates or updates the file.""" - - syftbox_folder_id = self.get_syftbox_folder_id() - file_id = self._get_own_peer_info_file_id() - - # Convert to JSON string - json_data = peer_info.to_json() - file_payload, _ = self.create_file_payload(json_data) - - if file_id is None: - # Create new file - file_metadata = { - "name": SYFT_VERSION_FILE, - "parents": [syftbox_folder_id], - } - self._create_syft_file(file_metadata, file_payload) - else: - # Update existing file - execute_with_retries( - self.drive_service.files().update( - fileId=file_id, media_body=file_payload - ) - ) - - def _get_peer_info_file_id_of(self, peer_email: str) -> Optional[str]: - """Find the peer info file, SYFT_version.json, of ``peer_email``.""" - query = ( - f"name='{SYFT_VERSION_FILE}' and trashed=false and '{peer_email}' in owners" - ) - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - return items[0]["id"] if items else None - - def read_own_peer_info_file(self) -> Optional["PeerInfo"]: - """Read the own peer info file from /SyftBox. - - Returns None when there is no file. - - Raises: - TransportUnavailableError: when Drive cannot be reached. - ValueError: when the file is not a peer info file this client reads. - """ - return self._read_peer_info_file( - self._get_own_peer_info_file_id, "the own peer info file" - ) - - def read_peer_info_file(self, peer_email: str) -> Optional["PeerInfo"]: - """Read the peer info file from a peer's /SyftBox folder. - - Returns None when the peer has no file that this account may read. - - Raises: - TransportUnavailableError: when Drive cannot be reached. - ValueError: when the file is not a peer info file this client reads. - """ - return self._read_peer_info_file( - lambda: self._get_peer_info_file_id_of(peer_email), - f"the peer info file of {peer_email}", - ) - - def _read_peer_info_file( - self, find_file_id: Callable[[], str | None], file_label: str - ) -> Optional["PeerInfo"]: - """Find, download and parse a peer info file. - - Returns None when ``find_file_id`` finds no file. ``file_label`` names - the file in errors. - """ - from syft.sync.version.peer_info import PeerInfo - - with raise_unavailable(f"Could not read {file_label} from Drive."): - file_id = find_file_id() - if file_id is None: - return None - file_data = self.download_file(file_id) - try: - return PeerInfo.from_json(file_data.decode("utf-8")) - except (ValueError, MigrationError) as e: - raise ValueError(f"Could not read {file_label}: {e}") from e - - def share_peer_info_file_with_peer( - self, peer_email: str, peer_info: "PeerInfo" - ) -> None: - """Share the peer info file with a peer so they can read it. - - A missing file is created from ``peer_info``, the info of this client. - """ - file_id = self._get_own_peer_info_file_id() - if file_id is None: - self.write_peer_info_file(peer_info) - file_id = self._get_own_peer_info_file_id() - - if file_id: - self.add_permission(file_id, peer_email, write=False) - - # ========================================================================= - # CHECKPOINT METHODS - # ========================================================================= - - def _get_checkpoints_folder_name(self) -> str: - """Get the checkpoints folder name: {email}-{version}-checkpoints""" - return f"{self.email}-{SYFT_VERSION}-checkpoints" - - def _get_checkpoints_folder_id(self) -> str | None: - """Find the checkpoints folder ID from Google Drive.""" - folders = self._find_folders( - name_contains=[f"{self.email}-", "-checkpoints"], - parent_id=self.get_syftbox_folder_id(), - ) - return self._find_or_adopt_versioned_folder( - folders, current_name=self._get_checkpoints_folder_name() - ) - - def _get_or_create_checkpoints_folder_id(self) -> str: - """Get or create the checkpoints folder.""" - folder_id = self._get_checkpoints_folder_id() - if folder_id is not None: - return folder_id - # Create the folder - folder_name = self._get_checkpoints_folder_name() - syftbox_folder_id = self.get_syftbox_folder_id() - return self.create_folder(folder_name, syftbox_folder_id) - - def upload_raw_checkpoint(self, filename: str, data: bytes) -> str: - """Upload raw checkpoint bytes to Google Drive. - - Uploads the new checkpoint first, then deletes old ones. - """ - folder_id = self._get_or_create_checkpoints_folder_id() - payload, _ = self.create_file_payload(data) - - file_metadata = { - "name": filename, - "parents": [folder_id], - } - - result = self._create_syft_file(file_metadata, payload) - - # Only delete old checkpoints after successful upload - self.delete_all_checkpoints(exclude_file_id=result.get("id")) - - return result.get("id") - - def download_raw_latest_checkpoint(self) -> bytes | None: - """Download the latest checkpoint as raw bytes, or None.""" - folder_id = self._get_checkpoints_folder_id() - if folder_id is None: - return None - - query = ( - f"'{folder_id}' in parents and trashed=false " - f"and name contains '{CHECKPOINT_FILENAME_PREFIX}'" - ) - results = ( - self.drive_service.files().list(q=query, fields="files(id, name)").execute() - ) - items = results.get("files", []) - - if not items: - return None - - latest_file = None - latest_timestamp = -1.0 - for item in items: - timestamp = Checkpoint.filename_to_timestamp(item["name"]) - if timestamp is not None and timestamp > latest_timestamp: - latest_timestamp = timestamp - latest_file = item - - if latest_file is None: - return None - - try: - return self.download_file(latest_file["id"]) - except Exception as e: - print(f"Warning: Failed to load checkpoint: {e}") - return None - - def delete_all_checkpoints(self, exclude_file_id: str | None = None): - """Delete all existing full checkpoints (not incremental ones). - - Args: - exclude_file_id: If provided, skip deleting this file ID - (used to preserve a newly uploaded checkpoint). - """ - folder_id = self._get_checkpoints_folder_id() - if folder_id is None: - return - - # List only full checkpoint files (start with "checkpoint_" not "incremental_checkpoint_") - query = ( - f"'{folder_id}' in parents and trashed=false " - f"and name contains '{CHECKPOINT_FILENAME_PREFIX}'" - ) - results = ( - self.drive_service.files().list(q=query, fields="files(id, name)").execute() - ) - items = results.get("files", []) - - # Delete only full checkpoints (not incremental ones) - for item in items: - if item["name"].startswith(INCREMENTAL_CHECKPOINT_PREFIX): - continue # Skip incremental checkpoints - if item["id"] == exclude_file_id: - continue # Skip the newly uploaded checkpoint - try: - self.drive_service.files().delete(fileId=item["id"]).execute() - except Exception as e: - print(f"Warning: Failed to delete checkpoint {item['name']}: {e}") - - # ========================================================================= - # INCREMENTAL CHECKPOINT METHODS - # ========================================================================= - - def upload_raw_incremental_checkpoint(self, filename: str, data: bytes) -> str: - """Upload raw incremental checkpoint bytes to Google Drive.""" - folder_id = self._get_or_create_checkpoints_folder_id() - payload, _ = self.create_file_payload(data) - - file_metadata = { - "name": filename, - "parents": [folder_id], - } - - result = self._create_syft_file(file_metadata, payload) - return result.get("id") - - def download_all_raw_incremental_checkpoints(self) -> list[bytes]: - """Download all incremental checkpoints as raw bytes.""" - folder_id = self._get_checkpoints_folder_id() - if folder_id is None: - return [] - - query = ( - f"'{folder_id}' in parents and trashed=false " - f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" - ) - results = ( - self.drive_service.files().list(q=query, fields="files(id, name)").execute() - ) - items = results.get("files", []) - - if not items: - return [] - - result = [] - for item in items: - try: - file_data = self.download_file(item["id"]) - result.append(file_data) - except Exception as e: - print( - f"Warning: Failed to load incremental checkpoint {item['name']}: {e}" - ) - continue - return result - - def get_incremental_checkpoint_count(self) -> int: - """Get the number of incremental checkpoints on Google Drive.""" - folder_id = self._get_checkpoints_folder_id() - if folder_id is None: - return 0 - - query = ( - f"'{folder_id}' in parents and trashed=false " - f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" - ) - results = self.drive_service.files().list(q=query, fields="files(id)").execute() - return len(results.get("files", [])) - - def get_next_incremental_sequence_number(self) -> int: - """Get the next sequence number for incremental checkpoints.""" - folder_id = self._get_checkpoints_folder_id() - if folder_id is None: - return 1 - - query = ( - f"'{folder_id}' in parents and trashed=false " - f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" - ) - results = ( - self.drive_service.files().list(q=query, fields="files(name)").execute() - ) - items = results.get("files", []) - - if not items: - return 1 - - # Find highest sequence number - max_seq = 0 - for item in items: - seq = IncrementalCheckpoint.filename_to_sequence_number(item["name"]) - if seq is not None and seq > max_seq: - max_seq = seq - - return max_seq + 1 - - def delete_all_incremental_checkpoints(self) -> None: - """Delete all incremental checkpoints (called after compacting).""" - folder_id = self._get_checkpoints_folder_id() - if folder_id is None: - return - - query = ( - f"'{folder_id}' in parents and trashed=false " - f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" - ) - results = ( - self.drive_service.files().list(q=query, fields="files(id, name)").execute() - ) - items = results.get("files", []) - - for item in items: - try: - self.drive_service.files().delete(fileId=item["id"]).execute() - except Exception as e: - print( - f"Warning: Failed to delete incremental checkpoint {item['name']}: {e}" - ) - - def get_events_count_since_checkpoint( - self, checkpoint_timestamp: float | None - ) -> int: - """ - Count events created after the checkpoint timestamp. - - Args: - checkpoint_timestamp: The timestamp of the checkpoint, or None for all events. - - Returns: - Number of events since the checkpoint. - """ - personal_folder_id = self.get_personal_syftbox_folder_id() - file_metadatas = self.get_file_metadatas_from_folder(personal_folder_id) - - if checkpoint_timestamp is None: - # Count all events - return len( - [ - f - for f in file_metadatas - if f["name"].startswith("syfteventsmessagev3_") - ] - ) - - # Count events after checkpoint - count = 0 - for f in file_metadatas: - if not f["name"].startswith("syfteventsmessagev3_"): - continue - event_timestamp = self._extract_timestamp_from_filename(f["name"]) - if event_timestamp is not None and event_timestamp > checkpoint_timestamp: - count += 1 - return count - - def download_raw_events_since_timestamp( - self, since_timestamp: float - ) -> list[bytes]: - """Download event files created after a timestamp as raw bytes.""" - personal_folder_id = self.get_personal_syftbox_folder_id() - file_metadatas = self.get_file_metadatas_from_folder( - personal_folder_id, since_timestamp=since_timestamp - ) - - valid_fname_objs = self._get_valid_events_from_file_metadatas(file_metadatas) - - result = [] - for fname_obj in sorted(valid_fname_objs, key=lambda x: x.timestamp): - if fname_obj.timestamp <= since_timestamp: - continue - - gdrive_id = [ - f for f in file_metadatas if f["name"] == fname_obj.as_string() - ][0]["id"] - - try: - file_data = self.download_file(gdrive_id) - result.append(file_data) - except Exception as e: - print(f"Warning: Failed to download event: {e}") - continue - - return result - - # ========================================================================= - # ROLLING STATE METHODS - # ========================================================================= - - def _get_rolling_state_folder_name(self) -> str: - """Get the rolling state folder name: {email}-{version}-rolling-state""" - return f"{self.email}-{SYFT_VERSION}-rolling-state" - - def _get_rolling_state_folder_id(self, use_cache: bool = True) -> str | None: - """ - Find the rolling state folder ID from Google Drive. - - Args: - use_cache: If True, return cached value if available. - """ - if use_cache and self._rolling_state_folder_id is not None: - return self._rolling_state_folder_id - - folders = self._find_folders( - name_contains=[f"{self.email}-", "-rolling-state"], - parent_id=self.get_syftbox_folder_id(), - ) - folder_id = self._find_or_adopt_versioned_folder( - folders, current_name=self._get_rolling_state_folder_name() - ) - if folder_id is not None: - self._rolling_state_folder_id = folder_id - return folder_id - - def _get_or_create_rolling_state_folder_id(self) -> str: - """Get or create the rolling state folder.""" - folder_id = self._get_rolling_state_folder_id() - if folder_id is not None: - return folder_id - # Create the folder - folder_name = self._get_rolling_state_folder_name() - syftbox_folder_id = self.get_syftbox_folder_id() - folder_id = self.create_folder(folder_name, syftbox_folder_id) - self._rolling_state_folder_id = folder_id - return folder_id - - def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: - """Upload raw rolling state bytes to Google Drive. - - Optimized to use update() if file ID is cached. - """ - payload, _ = self.create_file_payload(data) - - # Try to update existing file if we have a cached ID - if self._rolling_state_file_id is not None: - try: - self.drive_service.files().update( - fileId=self._rolling_state_file_id, - media_body=payload, - ).execute() - return self._rolling_state_file_id - except Exception as e: - # The cached file is gone or unreachable. Clear the cache and - # write a new file below. - print( - f"Warning: could not update rolling state " - f"{self._rolling_state_file_id}, writing a new file: {e}" - ) - self._rolling_state_file_id = None - - folder_id = self._get_or_create_rolling_state_folder_id() - - file_metadata = { - "name": filename, - "parents": [folder_id], - } - - result = self._create_syft_file(file_metadata, payload) - file_id = result.get("id") - self._rolling_state_file_id = file_id - return file_id - - def download_raw_rolling_state(self) -> bytes | None: - """Download the latest rolling state as raw bytes, or None. - - Also populates the folder and file ID caches for subsequent uploads. - """ - folder_id = self._get_rolling_state_folder_id() - if folder_id is None: - return None - - query = ( - f"'{folder_id}' in parents and trashed=false " - f"and name contains '{ROLLING_STATE_FILENAME_PREFIX}'" - ) - results = ( - self.drive_service.files().list(q=query, fields="files(id, name)").execute() - ) - items = results.get("files", []) - - if not items: - return None - - latest_file = None - latest_timestamp = -1.0 - for item in items: - timestamp = RollingState.filename_to_timestamp(item["name"]) - if timestamp is not None and timestamp > latest_timestamp: - latest_timestamp = timestamp - latest_file = item - - if latest_file is None: - return None - - self._rolling_state_file_id = latest_file["id"] - - try: - return self.download_file(latest_file["id"]) - except Exception as e: - print(f"Warning: Failed to load rolling state: {e}") - self._rolling_state_file_id = None - return None - - def delete_rolling_state(self) -> None: - """Delete all existing rolling state files and clear cache.""" - # Clear the file ID cache - self._rolling_state_file_id = None - - folder_id = self._get_rolling_state_folder_id() - if folder_id is None: - return - - # List all rolling state files - query = ( - f"'{folder_id}' in parents and trashed=false " - f"and name contains '{ROLLING_STATE_FILENAME_PREFIX}'" - ) - results = ( - self.drive_service.files().list(q=query, fields="files(id, name)").execute() - ) - items = results.get("files", []) - - # Delete each rolling state file - for item in items: - try: - self.drive_service.files().delete(fileId=item["id"]).execute() - except Exception as e: - print(f"Warning: Failed to delete rolling state {item['name']}: {e}") - - # ========================================================================= - # ENCRYPTION BUNDLE METHODS - # ========================================================================= - - def _find_encryption_bundles_folder_id(self) -> str | None: - """The id of the encryption bundles folder of this user, or None when it does not exist.""" - if self._encryption_bundles_folder_id: - return self._encryption_bundles_folder_id - folder_name = GdriveEncryptionBundlesFolder(email=self.email).as_string() - folder_id = self._find_folder_by_name( - folder_name, parent_id=self.get_syftbox_folder_id() - ) - if folder_id: - self._encryption_bundles_folder_id = folder_id - return folder_id - - def _get_or_create_encryption_bundles_folder_id(self) -> str: - """Get or create the encryption bundles folder for this user.""" - folder_id = self._find_encryption_bundles_folder_id() - if not folder_id: - folder_name = GdriveEncryptionBundlesFolder(email=self.email).as_string() - folder_id = self.create_folder(folder_name, self.get_syftbox_folder_id()) - self._encryption_bundles_folder_id = folder_id - return folder_id - - def _encryption_bundle_filename(self, owner_email: str, peer_email: str) -> str: - return f"encryption_bundle_{owner_email}_for_{peer_email}.json" - - def write_encryption_bundle(self, peer_email: str, bundle_json: str) -> None: - """Write own encryption bundle for a peer to own bundles folder.""" - folder_id = self._get_or_create_encryption_bundles_folder_id() - filename = self._encryption_bundle_filename(self.email, peer_email) - file_payload, _ = self.create_file_payload(bundle_json) - - # Check if file already exists - query = f"name='{filename}' and '{folder_id}' in parents and trashed=false" - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - - if items: - execute_with_retries( - self.drive_service.files().update( - fileId=items[0]["id"], media_body=file_payload - ) - ) - else: - file_metadata = {"name": filename, "parents": [folder_id]} - self._create_syft_file(file_metadata, file_payload) - - def share_encryption_bundles_folder(self, peer_email: str) -> None: - """Share the bundles folder with a peer so they can read bundles.""" - folder_id = self._get_or_create_encryption_bundles_folder_id() - self.add_permission(folder_id, peer_email, write=False) - - def read_own_encryption_bundle(self, peer_email: str) -> str | None: - """Read the bundle this client published for ``peer_email``, or None. - - A read never creates the bundles folder: without it, nothing is published. - """ - folder_id = self._find_encryption_bundles_folder_id() - if folder_id is None: - return None - filename = self._encryption_bundle_filename(self.email, peer_email) - query = f"name='{filename}' and '{folder_id}' in parents and trashed=false" - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - if not items: - return None - return self.download_file(items[0]["id"]).decode("utf-8") - - def read_peer_encryption_bundle(self, peer_email: str) -> str | None: - """Read encryption bundle that peer wrote for us. - - Searches for: encryption_bundle_{peer_email}_for_{self.email}.json - in peer's bundles folder. - """ - filename = self._encryption_bundle_filename(peer_email, self.email) - query = f"name='{filename}' and trashed=false and '{peer_email}' in owners" - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - if not items: - return None - try: - data = self.download_file(items[0]["id"]) - except Exception as e: - print(f"Warning: could not download the bundle of {peer_email}: {e}") - return None - try: - return data.decode("utf-8") - except ValueError as e: - print(f"Warning: could not read the bundle of {peer_email}: {e}") - return None diff --git a/syft/sync/connections/drive/grdrive_config.py b/syft/sync/connections/drive/grdrive_config.py index 95131c20ab8..6848af1bf95 100644 --- a/syft/sync/connections/drive/grdrive_config.py +++ b/syft/sync/connections/drive/grdrive_config.py @@ -1,11 +1,8 @@ from pathlib import Path -from typing import ClassVar, Type from syft.sync.connections.base_connection import ConnectionConfig -from syft.sync.connections.drive.gdrive_transport import GDriveConnection class GdriveConnectionConfig(ConnectionConfig): - connection_type: ClassVar[Type["GDriveConnection"]] = GDriveConnection email: str token_path: Path | None = None diff --git a/syft/sync/connections/drive/mock_drive_service.py b/syft/sync/connections/drive/mock_drive_service.py index 58c60fa0876..51c9a1c6eca 100644 --- a/syft/sync/connections/drive/mock_drive_service.py +++ b/syft/sync/connections/drive/mock_drive_service.py @@ -1,7 +1,7 @@ -"""Mock Google Drive service for testing GDriveConnection without real API calls. +"""Mock Google Drive service for testing the Drive code without real API calls. This module provides a mock implementation of the Google Drive API that can be -injected into GDriveConnection, allowing testing of actual GDrive code paths +injected into DriveCachedClient, allowing testing of actual GDrive code paths without making real API calls. Usage: @@ -17,16 +17,19 @@ user1_service = MockDriveService(shared_backing_store, "user1@example.com") user2_service = MockDriveService(shared_backing_store, "user2@example.com") - # Inject into GDriveConnection - connection = GDriveConnection.from_mock_service("user1@example.com", user1_service) + # Inject into DriveCachedClient + cached_client = DriveCachedClient.from_service("user1@example.com", user1_service) """ +import json import re import uuid from datetime import datetime from typing import Any, Callable, Dict, List, Optional, Tuple +import httplib2 +from googleapiclient.errors import HttpError from pydantic import BaseModel, Field -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient GOOGLE_FOLDER_MIME_TYPE = "application/vnd.google-apps.folder" @@ -592,6 +595,13 @@ def execute(self) -> Dict[str, Any]: return file.to_metadata(self._fields) +def http_error(status: int, reason: str, message: str) -> HttpError: + """An HttpError with the JSON body that Drive sends.""" + errors = [{"reason": reason, "message": message}] + body = {"error": {"code": status, "message": message, "errors": errors}} + return HttpError(httplib2.Response({"status": status}), json.dumps(body).encode()) + + class MockDeleteRequest: """Mock request for files().delete().""" @@ -609,17 +619,17 @@ def execute(self) -> None: """Execute the delete request.""" file = self._backing_store.get_file(self._file_id) if file is None: - raise Exception( - f"" - ) + raise http_error(404, "notFound", f"File not found: {self._file_id}") # Check if user has permission to delete is_owner = any( owner.get("emailAddress") == self._current_user for owner in file.owners ) if not is_owner: - raise Exception( - f"" + raise http_error( + 403, + "insufficientFilePermissions", + f"The user may not delete {self._file_id}", ) self._delete_tree(self._file_id) @@ -979,8 +989,8 @@ def new_batch_http_request(self, callback=None) -> MockBatchHttpRequest: def pair_with_mock_service( email1: str, email2: str -) -> Tuple[GDriveConnection, GDriveConnection]: - """Pair two GDriveConnections using mock services. +) -> Tuple[DriveCachedClient, DriveCachedClient]: + """Pair two DriveCachedClients using mock services. Args: email1: Email for the first user @@ -989,6 +999,6 @@ def pair_with_mock_service( backing_store = MockDriveBackingStore() service1 = MockDriveService(backing_store, email1) service2 = MockDriveService(backing_store, email2) - return GDriveConnection.from_service( + return DriveCachedClient.from_service( email1, service1 - ), GDriveConnection.from_service(email2, service2) + ), DriveCachedClient.from_service(email2, service2) diff --git a/syft/sync/login_utils.py b/syft/sync/login_utils.py index ee757e0bea5..e3f5d028e81 100644 --- a/syft/sync/login_utils.py +++ b/syft/sync/login_utils.py @@ -19,10 +19,13 @@ def _read_remote_version( token_path: Optional[Path], ) -> Optional[VersionInfo]: """Read version file from GDrive without a full SyftboxManager.""" - from syft.sync.connections.drive.gdrive_transport import GDriveConnection + from syft.sync.connections.drive.drive_cached_client import DriveCachedClient + from syft.sync.connections.drive.drive_peering import DrivePeering - conn = GDriveConnection.from_token_path(email=email, token_path=token_path) - peer_info = conn.read_own_peer_info_file() + cached_client = DriveCachedClient.from_token_path( + email=email, token_path=token_path + ) + peer_info = DrivePeering(cached_client).read_own_peer_info() return peer_info.version_info if peer_info else None diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index a4c22fec7d9..60f4847eeb4 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -13,15 +13,15 @@ from syft.sync.callback_mixin import BaseModelCallbackMixin from syft.sync.connections.connection_router import ConnectionRouter from syft.sync.connections.drive import mock_drive_service +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient +from syft.sync.connections.drive.drive_collections import DriveCollections from syft.sync.connections.drive.drive_connections import ( - DriveCollections, - DriveOwnerStore, - DriveP2PConnection, - DrivePeering, DriveWatcherStore, - gdrive_connection_of, + drive_cached_client_of, ) -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore +from syft.sync.connections.drive.drive_p2p import DriveP2PConnection +from syft.sync.connections.drive.drive_peering import DrivePeering from syft.sync.connections.drive.grdrive_config import GdriveConnectionConfig from syft.sync.events.file_change_event import ( FileChangeEvent, @@ -603,18 +603,18 @@ def _transport_from_config(config: SyftboxManagerConfig) -> Transport: "Only one Google Drive connection for every transport component " "is built so far" ) - drive = GDriveConnection.from_config(drive_config) - return SyftboxManager._drive_transport(drive) + cached_client = DriveCachedClient.from_config(drive_config) + return SyftboxManager._drive_transport(cached_client) @staticmethod - def _drive_transport(drive: GDriveConnection) -> Transport: - """Every transport component over one Google Drive connection.""" + def _drive_transport(cached_client: DriveCachedClient) -> Transport: + """Every transport component over one Google Drive account.""" return Transport( - p2p=DriveP2PConnection(drive), - owner_store=DriveOwnerStore(drive), - watcher_store=DriveWatcherStore(drive), - collections=DriveCollections(drive), - peering=DrivePeering(drive), + p2p=DriveP2PConnection(cached_client), + owner_store=DriveOwnerStore(cached_client), + watcher_store=DriveWatcherStore(), + collections=DriveCollections(cached_client), + peering=DrivePeering(cached_client), platform=GdriveFilesPlatform(), ) @@ -782,7 +782,7 @@ def pair_with_mock_drive_service_connection( ): """Create a pair of managers using mock Google Drive services for testing. - This creates managers that use the actual GDriveConnection code but with + This creates managers that use the actual Google Drive code but with mock services instead of real Google Drive API calls. This allows testing the full GDrive code path without network calls. @@ -829,14 +829,14 @@ def pair_with_mock_drive_service_connection( encryption_policy=ds_encryption_policy or default_policy, ) - do_connection, ds_connection = mock_drive_service.pair_with_mock_service( + do_cached_client, ds_cached_client = mock_drive_service.pair_with_mock_service( do_config.email, ds_config.email ) do_manager = cls.from_config( - do_config, SyftboxManager._drive_transport(do_connection) + do_config, SyftboxManager._drive_transport(do_cached_client) ) ds_manager = cls.from_config( - ds_config, SyftboxManager._drive_transport(ds_connection) + ds_config, SyftboxManager._drive_transport(ds_cached_client) ) # Set up callbacks for DS -> DO communication @@ -1325,15 +1325,13 @@ def _copy(self): from copy import deepcopy new_config = deepcopy(self.config) - connection = gdrive_connection_of(self._connection_router.transport) - if not isinstance(connection, GDriveConnection): - raise ValueError("Only GDriveConnections can be copied") - if isinstance(connection.drive_service, mock_drive_service.MockDriveService): + drive_service = drive_cached_client_of( + self._connection_router.transport + ).client.drive_service + if isinstance(drive_service, mock_drive_service.MockDriveService): # A new connection to the same backing store - new_connection = GDriveConnection.from_service( - self.email, connection.drive_service - ) + cached_client = DriveCachedClient.from_service(self.email, drive_service) return SyftboxManager.from_config( - new_config, SyftboxManager._drive_transport(new_connection) + new_config, SyftboxManager._drive_transport(cached_client) ) return SyftboxManager.from_config(new_config) diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 99771f32fd5..43a38558f81 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -116,6 +116,19 @@ def _last_persisted_event_timestamp( return since +def _incrementals_after( + full_checkpoint: Checkpoint | None, incrementals: List[IncrementalCheckpoint] +) -> List[IncrementalCheckpoint]: + """The incrementals that ``full_checkpoint`` does not hold yet. + + A full checkpoint is built after the incrementals it holds, so an older + incremental is already in it. A failed delete can leave one on the backend. + """ + if full_checkpoint is None: + return incrementals + return [cp for cp in incrementals if cp.timestamp > full_checkpoint.timestamp] + + def _later_timestamp(current: float | None, candidate: float | None) -> float | None: """The later of two event timestamps, either of which may be missing.""" if candidate is None: @@ -389,7 +402,9 @@ def fetch_owner_state( """ router = self.connection_router full_checkpoint = router.get_latest_checkpoint() - incrementals = router.get_all_incremental_checkpoints() + incrementals = _incrementals_after( + full_checkpoint, router.get_all_incremental_checkpoints() + ) rolling_state = router.get_rolling_state() if rolling_state is not None and rolling_state.event_count == 0: rolling_state = None @@ -1237,10 +1252,11 @@ def compact_checkpoints(self) -> Checkpoint: existing_checkpoint = self.connection_router.get_latest_checkpoint() # Download all incremental checkpoints - incremental_cps = self.connection_router.get_all_incremental_checkpoints() + all_incremental_cps = self.connection_router.get_all_incremental_checkpoints() - if not incremental_cps: + if not all_incremental_cps: raise ValueError("No incremental checkpoints to compact") + incremental_cps = _incrementals_after(existing_checkpoint, all_incremental_cps) print(f"Found {len(incremental_cps)} incremental checkpoints to compact") diff --git a/syft/version.py b/syft/version.py index 9309193be21..b8a788bd066 100644 --- a/syft/version.py +++ b/syft/version.py @@ -7,7 +7,8 @@ # Current client version — the single source of truth. # pyproject.toml and syft/__init__.py read from here. # Must stay a plain X.Y.Z: peers compare it with _parse_semver and Drive folder -# names embed it (see sync/version/version_info.py and gdrive_transport.py). +# names embed it (see sync/version/version_info.py, and drive_cached_client.py and +# drive_owner_store.py in sync/connections/drive/). SYFT_VERSION = "0.10.1" # Minimum client version we support communicating with. diff --git a/tests/integration/with_unit_coverage/test_gdrive_pagination.py b/tests/integration/with_unit_coverage/test_gdrive_pagination.py index 1f2a9ae1046..a4f2b6cfb4f 100644 --- a/tests/integration/with_unit_coverage/test_gdrive_pagination.py +++ b/tests/integration/with_unit_coverage/test_gdrive_pagination.py @@ -17,7 +17,7 @@ @pytest.mark.usefixtures("setup_delete_syftboxes") def test_pagination_and_early_termination(): """Test pagination with small page size and early termination with since_timestamp""" - from syft.sync.connections.drive.gdrive_transport import GDriveConnection + from syft.sync.connections.drive.drive_cached_client import DriveCachedClient manager_ds, manager_do = SyftboxManager._pair_with_google_drive_testing_connection( do_email=EMAIL_DO, @@ -36,13 +36,13 @@ def test_pagination_and_early_termination(): sleep(1) # Test pagination with small page size - original_method = GDriveConnection.get_file_metadatas_from_folder + original_method = DriveCachedClient.get_file_metadatas_from_folder def get_files_page_size_2(self, folder_id, since_timestamp=None, page_size=100): return original_method(self, folder_id, since_timestamp, page_size=2) with patch.object( - GDriveConnection, + DriveCachedClient, "get_file_metadatas_from_folder", get_files_page_size_2, ): diff --git a/tests/integration/with_unit_coverage/test_sync_manager.py b/tests/integration/with_unit_coverage/test_sync_manager.py index 95a8002caa9..9bf2baccd0e 100644 --- a/tests/integration/with_unit_coverage/test_sync_manager.py +++ b/tests/integration/with_unit_coverage/test_sync_manager.py @@ -30,7 +30,6 @@ create_test_pair, is_mock_mode, ) -from syft.sync.connections.drive.drive_connections import gdrive_connection_of def path_for_job(do_email: str, ds_email: str, filename: str = "test.job") -> str: @@ -553,7 +552,7 @@ def test_version_upgrade_breaks_communication(): ) # Phase 2: Simulate DS "upgrading" to a new incompatible version - ds_connection = gdrive_connection_of(ds_manager._connection_router.transport) + ds_peering = ds_manager._connection_router.transport.peering current = VersionInfo.current() new_version = VersionInfo( @@ -566,7 +565,7 @@ def test_version_upgrade_breaks_communication(): # Write the new version to GDrive (simulating DS upgrading their client) own_info = ds_manager.peer_manager.get_own_peer_info() - ds_connection.write_peer_info_file( + ds_peering.publish_peer_info( own_info.model_copy(update={"version_info": new_version}) ) @@ -588,4 +587,4 @@ def test_version_upgrade_breaks_communication(): ) # Cleanup: Restore DS's version file to the original version - ds_connection.write_peer_info_file(own_info) + ds_peering.publish_peer_info(own_info) diff --git a/tests/migrations/p2p/test_dataset_multicopy_delivery.py b/tests/migrations/p2p/test_dataset_multicopy_delivery.py index c3f174c0ef9..ede59e1425b 100644 --- a/tests/migrations/p2p/test_dataset_multicopy_delivery.py +++ b/tests/migrations/p2p/test_dataset_multicopy_delivery.py @@ -11,10 +11,6 @@ """ import pytest -from syft.sync.connections.drive.gdrive_transport import ( - CollectionFolder, - collection_name_query, -) from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX from syft_migration import ProtocolSchema from syft_rds import SyftRDSClient @@ -24,6 +20,10 @@ dataset_variant, ) +from syft.sync.connections.drive.drive_collections import ( + CollectionFolder, + collection_name_query, +) from tests.unit.utils import create_tmp_dataset_files DATASET_COLLECTION_NAME_QUERY = collection_name_query(DATASET_COLLECTION_PREFIX) diff --git a/tests/migrations/p2p/test_message_protocol_downgrade.py b/tests/migrations/p2p/test_message_protocol_downgrade.py index 7f3766544ee..ce810d0bcfe 100644 --- a/tests/migrations/p2p/test_message_protocol_downgrade.py +++ b/tests/migrations/p2p/test_message_protocol_downgrade.py @@ -29,7 +29,6 @@ from syft_migration import MigrationError, ProtocolSchema from tests.unit.utils import get_mock_events_messages, mock_message -from syft.sync.connections.drive.drive_connections import gdrive_connection_of def _client_schema( @@ -129,9 +128,8 @@ def _raw_proposal_version(do_manager, ds_email: str) -> str: def _raw_outbox_versions(ds_manager, do_email: str) -> list[str]: """The version fields of the DO's outbox blobs for us, unparsed.""" - raw_list = gdrive_connection_of( - ds_manager._connection_router.transport - ).watcher_download_raw_events_from_outbox(do_email, None) + p2p = ds_manager._connection_router.transport.p2p + raw_list = p2p.outbox_messages(do_email, None) return [_wire_version(ds_manager, do_email, raw) for raw in raw_list] diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index 6f54a66cf50..ede30eddc23 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -11,16 +11,25 @@ from uuid import uuid4 import pytest +from googleapiclient.errors import HttpError -from syft.sync.sync import datasite_owner_syncer +from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore from syft.sync.checkpoints.checkpoint import ( + CHECKPOINT_FILENAME_PREFIX, + INCREMENTAL_CHECKPOINT_PREFIX, Checkpoint, IncrementalCheckpoint, ) from syft.sync.checkpoints.rolling_state import RollingState +from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.mock_drive_service import ( + http_error, + pair_with_mock_service, +) from syft.sync.events.file_change_event import FileChangeEvent from syft.sync.syftbox_manager import SyftboxManager -from tests.unit.utils import get_mock_event +from syft.sync.sync import datasite_owner_syncer +from tests.unit.utils import fail_next_listing, get_mock_event def test_checkpoint_create_and_restore(): @@ -740,3 +749,83 @@ def test_owner_state_is_written_by_default(): assert do_manager._get_all_accepted_events_do() != [] assert router.get_rolling_state() is not None assert do_manager.datasite_owner_syncer._rolling_state is not None + + +def test_latest_checkpoint_read_retries_drive_error(monkeypatch): + do_cached_client, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = DriveOwnerStore(do_cached_client) + store.put_checkpoint("checkpoint_1.0.tar.gz", b"full") + + fail_next_listing(monkeypatch, CHECKPOINT_FILENAME_PREFIX) + assert store.latest_checkpoint() == b"full" + + +def test_incremental_count_retries_drive_error(monkeypatch): + do_cached_client, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = DriveOwnerStore(do_cached_client) + store.put_incremental("incremental_checkpoint_1_1.0.tar.gz", b"inc") + + fail_next_listing(monkeypatch, INCREMENTAL_CHECKPOINT_PREFIX) + assert store.incremental_count() == 1 + + +def test_latest_checkpoint_download_error_raises(monkeypatch): + do_cached_client, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = DriveOwnerStore(do_cached_client) + store.put_checkpoint("checkpoint_1.0.tar.gz", b"full") + + def forbidden(self, file_id: str) -> bytes: + raise http_error(403, "cannotDownloadFile", "The file cannot be downloaded") + + monkeypatch.setattr(GDriveClient, "download_file", forbidden) + with pytest.raises(HttpError): + store.latest_checkpoint() + + +def _file_event(do_email: str, path: str, content: str) -> FileChangeEvent: + return FileChangeEvent( + id=uuid4(), + path_in_datasite=Path(path), + datasite_email=do_email, + content=content, + old_hash=None, + new_hash=f"hash_{content}", + is_deleted=False, + submitted_timestamp=time.time(), + timestamp=time.time(), + ) + + +def test_incremental_left_after_compaction_does_not_revert_a_file(): + """A delete that fails after compaction leaves a merged incremental on the + backend. Restore and the next compaction must not apply it again.""" + _, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=True, + ) + do_email = do_manager.email + router = do_manager._connection_router + syncer = do_manager.datasite_owner_syncer + older = IncrementalCheckpoint( + email=do_email, + sequence_number=1, + events=[_file_event(do_email, "file1.txt", "v1")], + ) + newer = IncrementalCheckpoint( + email=do_email, + sequence_number=2, + events=[_file_event(do_email, "file1.txt", "v2")], + ) + router.upload_incremental_checkpoint(older) + router.upload_incremental_checkpoint(newer) + syncer.compact_checkpoints() + # The delete of the older incremental failed: it is still on the backend. + router.upload_incremental_checkpoint(older) + + syncer.event_cache.clear_cache() + syncer.initial_sync_done = False + do_manager.sync(auto_checkpoint=False) + assert syncer.event_cache.file_connection.read_file("file1.txt") == "v2" + + compacted = syncer.compact_checkpoints() + [file1] = [f for f in compacted.files if f.path == "file1.txt"] + assert file1.content == "v2" diff --git a/tests/unit/test_dataset_collection_listing.py b/tests/unit/test_dataset_collection_listing.py index 5c74a2448f5..5aa34119d11 100644 --- a/tests/unit/test_dataset_collection_listing.py +++ b/tests/unit/test_dataset_collection_listing.py @@ -1,4 +1,4 @@ -"""owner_list_all_collections_with_permissions skips only bad names. +"""Collections.list_mine_with_permissions on Drive skips only bad names. The Drive query matches a name prefix, so another tool can return a folder that this client cannot parse. The listing skips that folder. Every other failure is a @@ -8,25 +8,28 @@ from unittest.mock import Mock import pytest - -from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX +from syft.sync.connections.drive.drive_collections import DriveCollections +from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient + VALID = f"{DATASET_COLLECTION_PREFIX}_mytag_abc123" UNPARSEABLE = DATASET_COLLECTION_PREFIX -def _conn(files): - conn = GDriveConnection(email="alice@example.com", verbose=False) - conn.drive_service = Mock() - conn._syftbox_folder_id = "syftbox-id" - conn.drive_service.files().list().execute.return_value = {"files": files} - return conn +def _collections(files): + cached_client = DriveCachedClient(GDriveClient(Mock(), Mock), "alice@example.com") + cached_client._syftbox_folder_id = "syftbox-id" + cached_client.client.drive_service.files().list().execute.return_value = { + "files": files + } + return DriveCollections(cached_client) def test_a_valid_collection_is_returned(): - conn = _conn([{"id": "f1", "name": VALID, "appProperties": {}}]) - got = conn.owner_list_all_collections_with_permissions(DATASET_COLLECTION_PREFIX) + collections = _collections([{"id": "f1", "name": VALID, "appProperties": {}}]) + got = collections.list_mine_with_permissions(DATASET_COLLECTION_PREFIX) assert [(c.folder_id, c.tag, c.content_hash) for c in got] == [ ("f1", "mytag", "abc123") ] @@ -34,7 +37,7 @@ def test_a_valid_collection_is_returned(): def test_the_any_permission_flag_comes_from_app_properties(): - conn = _conn( + collections = _collections( [ { "id": "f1", @@ -43,24 +46,24 @@ def test_the_any_permission_flag_comes_from_app_properties(): } ] ) - got = conn.owner_list_all_collections_with_permissions(DATASET_COLLECTION_PREFIX) + got = collections.list_mine_with_permissions(DATASET_COLLECTION_PREFIX) assert got[0].has_any_permission is True def test_a_name_the_client_cannot_parse_is_skipped(): - conn = _conn( + collections = _collections( [ {"id": "bad", "name": UNPARSEABLE, "appProperties": {}}, {"id": "f1", "name": VALID, "appProperties": {}}, ] ) - got = conn.owner_list_all_collections_with_permissions(DATASET_COLLECTION_PREFIX) + got = collections.list_mine_with_permissions(DATASET_COLLECTION_PREFIX) assert [c.folder_id for c in got] == ["f1"] def test_a_missing_name_field_raises(): # A blanket except turned this defect into a collection that disappears # without a message. - conn = _conn([{"id": "f1", "appProperties": {}}]) + collections = _collections([{"id": "f1", "appProperties": {}}]) with pytest.raises(KeyError): - conn.owner_list_all_collections_with_permissions(DATASET_COLLECTION_PREFIX) + collections.list_mine_with_permissions(DATASET_COLLECTION_PREFIX) diff --git a/tests/unit/test_delete_syftbox.py b/tests/unit/test_delete_syftbox.py index cb1fd97436b..f61bd8d1ccc 100644 --- a/tests/unit/test_delete_syftbox.py +++ b/tests/unit/test_delete_syftbox.py @@ -6,7 +6,7 @@ from syft.sync.login_utils import handle_potential_version_mismatches_on_login from syft.sync.syftbox_manager import SyftboxManager from syft.sync.version.version_info import VersionInfo -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of EMAIL = "test@example.com" @@ -128,11 +128,13 @@ def test_a_terminal_still_asks(self, mock_stdin): assert _prompt_mismatch(_old_version_info(), _old_version_info()) == "2" -def _query_files(connection, name_contains): +def _query_files(cached_client, name_contains): """Query mock drive for files/folders whose name contains a substring.""" q = f"name contains '{name_contains}' and trashed=false" results = ( - connection.drive_service.files().list(q=q, fields="files(id, name)").execute() + cached_client.client.drive_service.files() + .list(q=q, fields="files(id, name)") + .execute() ) return results.get("files", []) @@ -154,7 +156,7 @@ class TestOrphanedMessageFiles: def test_proposal_in_do_archive_not_found(self): ds_manager, _, archived = _archive_one_proposal() - found = _connection(ds_manager).find_orphaned_syft_files() + found = _owner_store(ds_manager).find_orphaned_syft_files() assert archived assert not set(archived) & set(found) @@ -163,45 +165,51 @@ def test_proposal_moved_to_root_by_do_delete_found(self): ds_manager, do_manager, archived = _archive_one_proposal() do_manager.delete_syftbox(broadcast_delete_events=False) - found = _connection(ds_manager).find_orphaned_syft_files() + found = _owner_store(ds_manager).find_orphaned_syft_files() assert archived assert set(archived) <= set(found) def test_proposal_moved_to_root_when_do_deletes_syftbox_folder_only(self): ds_manager, do_manager, archived = _archive_one_proposal() - do_conn = _connection(do_manager) - archive_id = do_conn.owner_get_archive_folder_id(ds_manager.email) - do_conn.drive_service.files().delete( - fileId=do_conn.get_syftbox_folder_id() + do_cached_client = _cached_client(do_manager) + archive_id = do_cached_client.get_archive_folder_id(ds_manager.email) + do_cached_client.client.drive_service.files().delete( + fileId=do_cached_client.get_syftbox_folder_id() ).execute() - found = _connection(ds_manager).find_orphaned_syft_files() + found = _owner_store(ds_manager).find_orphaned_syft_files() - assert archive_id not in do_conn.drive_service._backing_store.files + assert ( + archive_id not in do_cached_client.client.drive_service._backing_store.files + ) assert archived assert set(archived) <= set(found) def test_untagged_file_in_root_not_found(self): ds_manager, _, _ = _archive_one_proposal() - drive = _connection(ds_manager).drive_service + drive = _cached_client(ds_manager).client.drive_service user_file = drive.files().create(body={"name": "checkpoint_final.pt"}).execute() - found = _connection(ds_manager).find_orphaned_syft_files() + found = _owner_store(ds_manager).find_orphaned_syft_files() assert user_file["id"] not in found def test_syftbox_folder_not_found(self): ds_manager, _, _ = _archive_one_proposal() - conn = _connection(ds_manager) + ds_cached_client = _cached_client(ds_manager) + + found = _owner_store(ds_manager).find_orphaned_syft_files() + + assert ds_cached_client.get_syftbox_folder_id() not in found - found = conn.find_orphaned_syft_files() - assert conn.get_syftbox_folder_id() not in found +def _cached_client(manager: SyftboxManager): + return drive_cached_client_of(manager._connection_router.transport) -def _connection(manager: SyftboxManager): - return gdrive_connection_of(manager._connection_router.transport) +def _owner_store(manager: SyftboxManager): + return manager._connection_router.transport.owner_store def _archive_one_proposal(): @@ -212,7 +220,9 @@ def _archive_one_proposal(): ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection() ds_manager._send_file_change(f"{do_manager.email}/shared/x.txt", "hi") do_manager.sync() - do_conn = _connection(do_manager) - archive_id = do_conn.owner_get_archive_folder_id(ds_manager.email) - archived = [m["id"] for m in do_conn.get_file_metadatas_from_folder(archive_id)] + do_cached_client = _cached_client(do_manager) + archive_id = do_cached_client.get_archive_folder_id(ds_manager.email) + archived = [ + m["id"] for m in do_cached_client.get_file_metadatas_from_folder(archive_id) + ] return ds_manager, do_manager, archived diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 53f095cf690..2ae768ca054 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -1,7 +1,7 @@ """The transport jobs on Google Drive, and the router that calls them. -The Drive classes wrap today's GDriveConnection methods, so these tests pin -what each job promises, on the mock Drive service the rest of the suite uses. +These tests pin what each job promises, on the mock Drive service the rest of +the suite uses. """ import inspect @@ -12,28 +12,31 @@ from unittest.mock import Mock import pytest +from googleapiclient.errors import HttpError import syft.sync.connections.connection_router as connection_router_module -from syft.sync.sync import datasite_owner_syncer from syft.sync.connections.connection_router import ConnectionRouter -from syft.sync.connections.drive.drive_connections import ( +from syft.sync.connections.drive import mock_drive_service +from syft.sync.connections.drive.drive_collections import ( + CollectionFolder, DriveCollections, - DriveOwnerStore, - DriveP2PConnection, - DrivePeering, - DriveWatcherStore, ) -from syft.sync.connections.drive.gdrive_transport import ( - CollectionFolder, - GDriveConnection, +from syft.sync.connections.drive.drive_connections import DriveWatcherStore +from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore +from syft.sync.connections.drive.drive_p2p import DriveP2PConnection +from syft.sync.connections.drive.drive_peering import DrivePeering +from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.mock_drive_service import ( + MockDeleteRequest, + http_error, + pair_with_mock_service, ) -from syft.sync.connections.drive import mock_drive_service -from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service from syft.sync.encryption_policy import EncryptionPolicy from syft.sync.events.file_change_event import FileChangeEventsMessage from syft.sync.messages.proposed_filechange import ProposedFileChangesMessage from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.sync import datasite_owner_syncer from syft.sync.transport import ( Collections, MessageKind, @@ -69,18 +72,18 @@ def _events_name(ts: float) -> str: @pytest.fixture def pair(): - """DO and DS connections over one mock Drive, with the DS's channels open.""" - do_conn, ds_conn = pair_with_mock_service(DO, DS) - DrivePeering(ds_conn).open_channels(DO) - return do_conn, ds_conn + """DO and DS cached clients over one mock Drive, with the DS's channels open.""" + do_cached_client, ds_cached_client = pair_with_mock_service(DO, DS) + DrivePeering(ds_cached_client).open_channels(DO) + return do_cached_client, ds_cached_client # -- the transport ---------------------------------------------------------------- def test_drive_transport_does_every_job(): - do_conn, _ = pair_with_mock_service(DO, DS) - transport = SyftboxManager._drive_transport(do_conn) + do_cached_client, _ = pair_with_mock_service(DO, DS) + transport = SyftboxManager._drive_transport(do_cached_client) assert isinstance(transport, Transport) assert isinstance(transport.p2p, P2PConnection) assert isinstance(transport.owner_store, OwnerStore) @@ -94,12 +97,14 @@ def test_drive_transport_does_every_job(): def test_a_proposed_change_sent_through_p2p_reaches_the_owner_inbox(pair): - do_conn, ds_conn = pair - status = DriveP2PConnection(ds_conn).send( + do_cached_client, ds_cached_client = pair + status = DriveP2PConnection(ds_cached_client).send( DO, _proposal_name(1.0), b"job", MessageKind.PROPOSED_CHANGES ) assert status is SendStatus.DELIVERED - [message] = DriveP2PConnection(do_conn).receive(DS, MessageKind.PROPOSED_CHANGES) + [message] = DriveP2PConnection(do_cached_client).receive( + DS, MessageKind.PROPOSED_CHANGES + ) assert (message.sender, message.kind, message.data) == ( DS, MessageKind.PROPOSED_CHANGES, @@ -108,13 +113,13 @@ def test_a_proposed_change_sent_through_p2p_reaches_the_owner_inbox(pair): def test_inbox_messages_are_received_oldest_first_up_to_max(pair): - do_conn, ds_conn = pair + do_cached_client, ds_cached_client = pair names = {submitted: _proposal_name(submitted) for submitted in (2.0, 1.0, 3.0)} for submitted, name in names.items(): - DriveP2PConnection(ds_conn).send( + DriveP2PConnection(ds_cached_client).send( DO, name, f"job {submitted}".encode(), MessageKind.PROPOSED_CHANGES ) - p2p = DriveP2PConnection(do_conn) + p2p = DriveP2PConnection(do_cached_client) received = p2p.receive(DS, MessageKind.PROPOSED_CHANGES, max_messages=2) assert [(m.msg_id, m.data) for m in received] == [ @@ -158,20 +163,20 @@ def counted_get_media(self, *args, **kwargs): return calls -def _send_proposals(ds_conn, count: int) -> list[str]: +def _send_proposals(ds_cached_client, count: int) -> list[str]: names = [_proposal_name(float(i)) for i in range(count)] for name in names: - DriveP2PConnection(ds_conn).send( + DriveP2PConnection(ds_cached_client).send( DO, name, name.encode(), MessageKind.PROPOSED_CHANGES ) return names def test_receiving_inbox_lists_once_and_downloads_each_message(pair, monkeypatch): - do_conn, ds_conn = pair - p2p = DriveP2PConnection(do_conn) + do_cached_client, ds_cached_client = pair + p2p = DriveP2PConnection(do_cached_client) p2p.receive(DS, MessageKind.PROPOSED_CHANGES) # finds the inbox folder once - names = _send_proposals(ds_conn, 3) + names = _send_proposals(ds_cached_client, 3) calls = _count_drive_calls(monkeypatch) received = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) @@ -181,11 +186,11 @@ def test_receiving_inbox_lists_once_and_downloads_each_message(pair, monkeypatch def test_inbox_messages_download_at_the_same_time(pair, monkeypatch): - do_conn, ds_conn = pair - _send_proposals(ds_conn, 3) + do_cached_client, ds_cached_client = pair + _send_proposals(ds_cached_client, 3) running, most_running = [0], [0] lock = threading.Lock() - original_download = GDriveConnection.download_file + original_download = GDriveClient.download_file def slow_download(self, file_id: str) -> bytes: with lock: @@ -196,15 +201,15 @@ def slow_download(self, file_id: str) -> bytes: running[0] -= 1 return original_download(self, file_id) - monkeypatch.setattr(GDriveConnection, "download_file", slow_download) - DriveP2PConnection(do_conn).receive(DS, MessageKind.PROPOSED_CHANGES) + monkeypatch.setattr(GDriveClient, "download_file", slow_download) + DriveP2PConnection(do_cached_client).receive(DS, MessageKind.PROPOSED_CHANGES) assert most_running[0] > 1 def test_marking_inbox_message_processed_is_one_drive_call(pair, monkeypatch): - do_conn, ds_conn = pair - _send_proposals(ds_conn, 2) - p2p = DriveP2PConnection(do_conn) + do_cached_client, ds_cached_client = pair + _send_proposals(ds_cached_client, 2) + p2p = DriveP2PConnection(do_cached_client) first, second = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) p2p.mark_processed(first) # finds or creates the archive folder once @@ -216,9 +221,9 @@ def test_marking_inbox_message_processed_is_one_drive_call(pair, monkeypatch): def test_marking_inbox_messages_processed_is_one_batch_request(pair, monkeypatch): - do_conn, ds_conn = pair - _send_proposals(ds_conn, 4) - p2p = DriveP2PConnection(do_conn) + do_cached_client, ds_cached_client = pair + _send_proposals(ds_cached_client, 4) + p2p = DriveP2PConnection(do_cached_client) first, *rest = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) p2p.mark_processed(first) # finds or creates the archive folder once @@ -230,9 +235,9 @@ def test_marking_inbox_messages_processed_is_one_batch_request(pair, monkeypatch def test_move_that_fails_in_batch_warns_and_rest_are_archived(pair): - do_conn, ds_conn = pair - _send_proposals(ds_conn, 2) - p2p = DriveP2PConnection(do_conn) + do_cached_client, ds_cached_client = pair + _send_proposals(ds_cached_client, 2) + p2p = DriveP2PConnection(do_cached_client) first, second = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) gone = first.model_copy(update={"handle": "no-such-file"}) @@ -247,41 +252,43 @@ def test_move_that_fails_in_batch_warns_and_rest_are_archived(pair): def test_inbox_message_that_fails_to_download_is_skipped_with_warning( pair, monkeypatch ): - do_conn, ds_conn = pair - names = _send_proposals(ds_conn, 3) + do_cached_client, ds_cached_client = pair + names = _send_proposals(ds_cached_client, 3) [broken_id] = [ m["file_id"] - for m in do_conn.owner_get_inbox_file_metadatas(DS) + for m in DriveP2PConnection(do_cached_client)._list_inbox(DS) if m["file_name"] == names[1] ] - original_download = GDriveConnection.download_file + original_download = GDriveClient.download_file def download_all_but_broken(self, file_id: str) -> bytes: if file_id == broken_id: raise ValueError("file is gone") return original_download(self, file_id) - monkeypatch.setattr(GDriveConnection, "download_file", download_all_but_broken) + monkeypatch.setattr(GDriveClient, "download_file", download_all_but_broken) with pytest.warns(UserWarning, match=f"Skipped inbox message {names[1]}"): - received = DriveP2PConnection(do_conn).receive(DS, MessageKind.PROPOSED_CHANGES) + received = DriveP2PConnection(do_cached_client).receive( + DS, MessageKind.PROPOSED_CHANGES + ) assert [m.msg_id for m in received] == [names[0], names[2]] def test_receiving_inbox_while_drive_is_down_raises(pair, monkeypatch): - do_conn, ds_conn = pair - _send_proposals(ds_conn, 2) - monkeypatch.setattr(GDriveConnection, "download_file", _drive_down) + do_cached_client, ds_cached_client = pair + _send_proposals(ds_cached_client, 2) + monkeypatch.setattr(GDriveClient, "download_file", _drive_down) with pytest.raises(TransportUnavailableError): - DriveP2PConnection(do_conn).receive(DS, MessageKind.PROPOSED_CHANGES) + DriveP2PConnection(do_cached_client).receive(DS, MessageKind.PROPOSED_CHANGES) def test_processed_inbox_message_is_not_received_again(pair): - do_conn, ds_conn = pair - DriveP2PConnection(ds_conn).send( + do_cached_client, ds_cached_client = pair + DriveP2PConnection(ds_cached_client).send( DO, _proposal_name(1.0), b"job", MessageKind.PROPOSED_CHANGES ) - p2p = DriveP2PConnection(do_conn) + p2p = DriveP2PConnection(do_cached_client) [message] = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) p2p.mark_processed(message) @@ -289,13 +296,13 @@ def test_processed_inbox_message_is_not_received_again(pair): def test_file_events_sent_through_p2p_reach_the_scientist_outbox(pair): - do_conn, ds_conn = pair + do_cached_client, ds_cached_client = pair name = _events_name(1.0) - status = DriveP2PConnection(do_conn).send( + status = DriveP2PConnection(do_cached_client).send( DS, name, b"results", MessageKind.FILE_EVENTS ) assert status is SendStatus.DELIVERED - p2p = DriveP2PConnection(ds_conn) + p2p = DriveP2PConnection(ds_cached_client) listed = p2p.outbox_message_metadatas(DO, since_timestamp=None) assert [m["file_name"] for m in listed] == [name] assert p2p.fetch_message(listed[0]["file_id"]) == b"results" @@ -303,32 +310,34 @@ def test_file_events_sent_through_p2p_reach_the_scientist_outbox(pair): def test_a_p2p_copy_for_another_thread_reads_the_same_outbox(pair): - do_conn, ds_conn = pair - DriveP2PConnection(do_conn).send( + do_cached_client, ds_cached_client = pair + DriveP2PConnection(do_cached_client).send( DS, _events_name(1.0), b"results", MessageKind.FILE_EVENTS ) - p2p = DriveP2PConnection(ds_conn) + p2p = DriveP2PConnection(ds_cached_client) [listed] = p2p.outbox_message_metadatas(DO, since_timestamp=None) copy = p2p.clone_for_thread() - assert copy._conn is not ds_conn + assert copy._cached_client is not ds_cached_client assert copy.fetch_message(listed["file_id"]) == b"results" def test_an_unknown_message_kind_is_refused(pair): - do_conn, _ = pair + do_cached_client, _ = pair with pytest.raises(ValueError): - DriveP2PConnection(do_conn).send(DS, _events_name(1.0), b"x", "unknown-kind") + DriveP2PConnection(do_cached_client).send( + DS, _events_name(1.0), b"x", "unknown-kind" + ) def test_drive_never_has_sends_waiting_for_confirmation(pair): - do_conn, _ = pair - assert DriveP2PConnection(do_conn).delivered() == [] + do_cached_client, _ = pair + assert DriveP2PConnection(do_cached_client).delivered() == [] def test_receiving_file_events_through_p2p_is_not_built_yet(pair): - _, ds_conn = pair - p2p = DriveP2PConnection(ds_conn) + _, ds_cached_client = pair + p2p = DriveP2PConnection(ds_cached_client) message = RawMessage(sender=DO, msg_id="m", kind=MessageKind.FILE_EVENTS, data=b"") with pytest.raises(NotImplementedError): p2p.receive(DO, MessageKind.FILE_EVENTS) @@ -340,8 +349,8 @@ def test_receiving_file_events_through_p2p_is_not_built_yet(pair): def test_owner_event_log_is_read_oldest_first(pair): - do_conn, _ = pair - store = DriveOwnerStore(do_conn) + do_cached_client, _ = pair + store = DriveOwnerStore(do_cached_client) # Written out of order: Drive lists by name, newest first. second = store.append_event(_events_name(2.0), b"second") first = store.append_event(_events_name(1.0), b"first") @@ -356,25 +365,48 @@ def test_owner_event_log_is_read_oldest_first(pair): def test_owner_event_count_matches_the_listed_handles(pair): - do_conn, _ = pair - store = DriveOwnerStore(do_conn) + do_cached_client, _ = pair + store = DriveOwnerStore(do_cached_client) store.append_event(_events_name(1.0), b"event") store.append_event("syfteventsmessagev3_not-a-valid-name.tar.gz", b"junk") assert store.count_events_since(None) == len(store.list_event_handles()) == 1 +def test_owner_event_log_read_while_drive_is_down_raises(pair, monkeypatch): + do_cached_client, _ = pair + store = DriveOwnerStore(do_cached_client) + store.append_event(_events_name(2.0), b"event") + + monkeypatch.setattr(GDriveClient, "download_file", _drive_down) + with pytest.raises(TransportUnavailableError): + store.fetch_events_since(1.0) + + +def test_owner_event_log_download_error_raises(pair, monkeypatch): + do_cached_client, _ = pair + store = DriveOwnerStore(do_cached_client) + store.append_event(_events_name(2.0), b"event") + + def forbidden(self, file_id: str) -> bytes: + raise http_error(403, "cannotDownloadFile", "The file cannot be downloaded") + + monkeypatch.setattr(GDriveClient, "download_file", forbidden) + with pytest.raises(HttpError): + store.fetch_events_since(1.0) + + def test_owner_store_keeps_the_newest_full_checkpoint(pair): - do_conn, _ = pair - store = DriveOwnerStore(do_conn) + do_cached_client, _ = pair + store = DriveOwnerStore(do_cached_client) store.put_checkpoint("checkpoint_1.0.tar.gz", b"one") store.put_checkpoint("checkpoint_2.0.tar.gz", b"two") assert store.latest_checkpoint() == b"two" def test_owner_store_counts_and_deletes_incremental_checkpoints(pair): - do_conn, _ = pair - store = DriveOwnerStore(do_conn) + do_cached_client, _ = pair + store = DriveOwnerStore(do_cached_client) store.put_incremental("incremental_checkpoint_1_1.0.tar.gz", b"inc") assert store.incremental_count() == 1 assert store.next_incremental_sequence() == 2 @@ -384,8 +416,8 @@ def test_owner_store_counts_and_deletes_incremental_checkpoints(pair): def test_owner_store_overwrites_and_deletes_the_rolling_state(pair): - do_conn, _ = pair - store = DriveOwnerStore(do_conn) + do_cached_client, _ = pair + store = DriveOwnerStore(do_cached_client) store.put_rolling_state("rolling_state_1.0.tar.gz", b"state 1") store.put_rolling_state("rolling_state_2.0.tar.gz", b"state 2") assert store.get_rolling_state() == b"state 2" @@ -394,15 +426,15 @@ def test_owner_store_overwrites_and_deletes_the_rolling_state(pair): def test_owner_store_files_carry_the_syft_tag(pair): - do_conn, _ = pair - store = DriveOwnerStore(do_conn) + do_cached_client, _ = pair + store = DriveOwnerStore(do_cached_client) store.put_checkpoint("checkpoint_1.0.tar.gz", b"full") store.put_incremental("incremental_checkpoint_1_1.0.tar.gz", b"inc") store.put_rolling_state("rolling_state_1.0.tar.gz", b"state") do_files = [ f - for f in do_conn.drive_service._backing_store.files.values() + for f in do_cached_client.client.drive_service._backing_store.files.values() if f.owners[0]["emailAddress"] == DO ] untagged = [f.name for f in do_files if f.properties != {"syft": "true"}] @@ -414,17 +446,20 @@ def test_owner_store_files_carry_the_syft_tag(pair): def test_watcher_storage_on_drive_is_not_built_yet(pair): - _, ds_conn = pair + _, ds_cached_client = pair with pytest.raises(NotImplementedError): - DriveWatcherStore(ds_conn).latest_checkpoint(DO) + DriveWatcherStore().latest_checkpoint(DO) # -- Collections ------------------------------------------------------------------ def test_a_published_collection_is_listed_and_downloaded_by_its_reader(pair): - do_conn, ds_conn = pair - owner, reader = DriveCollections(do_conn), DriveCollections(ds_conn) + do_cached_client, ds_cached_client = pair + owner, reader = ( + DriveCollections(do_cached_client), + DriveCollections(ds_cached_client), + ) files = {"mock.csv": b"a,b\n1,2\n"} content_hash = CollectionFolder.compute_hash(files) @@ -446,30 +481,46 @@ def test_a_published_collection_is_listed_and_downloaded_by_its_reader(pair): assert owner.list_mine(DATASETS) == [] +def test_a_failed_collection_delete_raises(pair, monkeypatch): + do_cached_client, _ = pair + owner = DriveCollections(do_cached_client) + owner.create(DATASETS, "census", "hash", DO) + + def forbidden(self) -> None: + raise http_error(403, "forbidden", "The user may not delete this folder") + + monkeypatch.setattr(MockDeleteRequest, "execute", forbidden) + with pytest.raises(HttpError): + owner.delete(DATASETS, "census") + monkeypatch.undo() + + assert owner.list_mine(DATASETS) == ["census"] + + # -- Peering ---------------------------------------------------------------------- def test_peering_finds_the_peer_that_opened_channels(pair): - do_conn, _ = pair - peering = DrivePeering(do_conn) + do_cached_client, _ = pair + peering = DrivePeering(do_cached_client) assert peering.whoami() == DO assert DS in peering.incoming_requests() def test_peering_keeps_the_peer_registry(pair): - do_conn, _ = pair - peering = DrivePeering(do_conn) + do_cached_client, _ = pair + peering = DrivePeering(do_cached_client) peering.update_peer_registry(DS, "accepted") assert peering.read_peer_registry(force_download=True)[DS]["state"] == "accepted" def test_peer_registry_update_while_drive_is_down_keeps_other_peers(pair, monkeypatch): - do_conn, _ = pair - peering = DrivePeering(do_conn) + do_cached_client, _ = pair + peering = DrivePeering(do_cached_client) peering.update_peer_registry(DS, "accepted") - do_conn._peers_json_cache = None # a new process has no cache + peering._peers_json_cache = None # a new process has no cache - monkeypatch.setattr(GDriveConnection, "download_file", _drive_down) + monkeypatch.setattr(GDriveClient, "download_file", _drive_down) with pytest.raises(TransportUnavailableError): peering.update_peer_registry("other@example.com", "accepted") monkeypatch.undo() @@ -480,12 +531,12 @@ def test_peer_registry_update_while_drive_is_down_keeps_other_peers(pair, monkey def test_reading_own_key_bundle_creates_no_folder(pair): owner = DrivePeering(pair[0]) assert owner.read_own_key_bundle(DS) is None - assert pair[0]._find_encryption_bundles_folder_id() is None + assert owner._find_encryption_bundles_folder_id() is None def test_peering_publishes_a_key_bundle_both_sides_can_read(pair): - do_conn, ds_conn = pair - owner, scientist = DrivePeering(do_conn), DrivePeering(ds_conn) + do_cached_client, ds_cached_client = pair + owner, scientist = DrivePeering(do_cached_client), DrivePeering(ds_cached_client) bundle = '{"bundle": "do-public-key"}' owner.publish_key_bundle(DS, bundle) @@ -497,8 +548,8 @@ def test_peering_publishes_a_key_bundle_both_sides_can_read(pair): def test_peering_publishes_peer_info_the_peer_can_read(pair): - do_conn, ds_conn = pair - owner, scientist = DrivePeering(do_conn), DrivePeering(ds_conn) + do_cached_client, ds_cached_client = pair + owner, scientist = DrivePeering(do_cached_client), DrivePeering(ds_cached_client) info = PeerInfo( version_info=VersionInfo.current(), encryption_policy=EncryptionPolicy.REQUIRED, @@ -515,9 +566,11 @@ def test_peering_publishes_peer_info_the_peer_can_read(pair): def _router_with_mock_p2p(): - do_conn, _ = pair_with_mock_service(DO, DS) + do_cached_client, _ = pair_with_mock_service(DO, DS) p2p = Mock(spec=DriveP2PConnection) - transport = SyftboxManager._drive_transport(do_conn).model_copy(update={"p2p": p2p}) + transport = SyftboxManager._drive_transport(do_cached_client).model_copy( + update={"p2p": p2p} + ) router = ConnectionRouter(transport=transport, peer_store=_store_knowing(DO, DS)) return router, p2p diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index 33613382a58..5621bb15f22 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -12,7 +12,7 @@ from syft.sync.syftbox_manager import SyftboxManager from tests.unit.test_sync_manager import path_for_job from tests.unit.utils import grant_job_inbox_access -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of # ========================================================================= @@ -624,17 +624,17 @@ def test_tampered_inbox_message_is_archived_with_warning(): ds_manager._send_file_change(file_path, "will be tampered") # Tamper with raw bytes in the mock backing store - do_conn = gdrive_connection_of(do_manager._connection_router.transport) - inbox_id = do_conn._get_own_datasite_inbox_id(ds_manager.email) - files = do_conn.get_file_metadatas_from_folder(inbox_id) + do_cached_client = drive_cached_client_of(do_manager._connection_router.transport) + inbox_id = do_cached_client.get_own_datasite_inbox_id(ds_manager.email) + files = do_cached_client.get_file_metadatas_from_folder(inbox_id) assert len(files) > 0 file_id = files[0]["id"] - mock_file = do_conn.drive_service._backing_store.files[file_id] + mock_file = do_cached_client.client.drive_service._backing_store.files[file_id] tampered = bytearray(mock_file.content) tampered[-1] ^= 0xFF mock_file.content = bytes(tampered) with pytest.warns(UserWarning, match="unread"): do_manager.sync() - assert do_conn.get_file_metadatas_from_folder(inbox_id) == [] + assert do_cached_client.get_file_metadatas_from_folder(inbox_id) == [] diff --git a/tests/unit/test_encryption_policy.py b/tests/unit/test_encryption_policy.py index 4973f394661..63f963b0e56 100644 --- a/tests/unit/test_encryption_policy.py +++ b/tests/unit/test_encryption_policy.py @@ -9,6 +9,7 @@ import pytest from pydantic import ValidationError +from syft.sync.connections.drive.drive_peering import DrivePeering from syft.sync.connections.connection_router import ConnectionRouter from syft.sync.connections.drive import mock_drive_service from syft.sync.encryption_policy import ( @@ -169,11 +170,16 @@ def test_own_version_advertises_policy(): def test_missing_version_file_created_with_own_policy(): - a_conn, b_conn = mock_drive_service.pair_with_mock_service( + a_cached_client, b_cached_client = mock_drive_service.pair_with_mock_service( "a@example.com", "b@example.com" ) - a_conn.share_peer_info_file_with_peer("b@example.com", _peer_info(REQUIRED)) - assert b_conn.read_peer_info_file("a@example.com").encryption_policy == REQUIRED + DrivePeering(a_cached_client).grant_peer_info_read( + "b@example.com", _peer_info(REQUIRED) + ) + assert ( + DrivePeering(b_cached_client).read_peer_info("a@example.com").encryption_policy + == REQUIRED + ) def test_malformed_peer_bundle_raises_on_peer_load(): diff --git a/tests/unit/test_gdrive_client.py b/tests/unit/test_gdrive_client.py new file mode 100644 index 00000000000..8dc8226cc48 --- /dev/null +++ b/tests/unit/test_gdrive_client.py @@ -0,0 +1,39 @@ +"""GDriveClient on the mock Drive service.""" + +import pytest +from googleapiclient.errors import HttpError + +from syft.sync.connections.drive.mock_drive_service import ( + MockDeleteRequest, + MockPermissionCreateRequest, + http_error, + pair_with_mock_service, +) + + +def test_batch_delete_raises_error_whose_text_contains_404(monkeypatch): + do_cached_client, _ = pair_with_mock_service("do@example.com", "ds@example.com") + folder_id = do_cached_client.client.create_folder( + "folder", do_cached_client.get_syftbox_folder_id() + ) + + def forbidden(self) -> None: + raise http_error(403, "forbidden", "Cannot delete file 1Xy404abc") + + monkeypatch.setattr(MockDeleteRequest, "execute", forbidden) + with pytest.raises(HttpError): + do_cached_client.client.delete_multiple_files_by_ids([folder_id]) + + +def test_batch_share_raises_error_whose_text_says_already_shared(monkeypatch): + do_cached_client, _ = pair_with_mock_service("do@example.com", "ds@example.com") + folder_id = do_cached_client.client.create_folder( + "folder", do_cached_client.get_syftbox_folder_id() + ) + + def forbidden(self) -> None: + raise http_error(403, "forbidden", "alreadyShared: no access to this folder") + + monkeypatch.setattr(MockPermissionCreateRequest, "execute", forbidden) + with pytest.raises(HttpError): + do_cached_client.client.batch_add_permissions(folder_id, ["ds@example.com"]) diff --git a/tests/unit/test_gdrive_pagination.py b/tests/unit/test_gdrive_pagination.py index f2550dc07b6..3355b013004 100644 --- a/tests/unit/test_gdrive_pagination.py +++ b/tests/unit/test_gdrive_pagination.py @@ -1,27 +1,34 @@ from unittest.mock import Mock, patch from uuid import uuid4 -from syft.sync.connections.drive.gdrive_transport import GDriveConnection + +from syft.sync.connections.drive.drive_collections import DriveCollections +from syft.sync.connections.drive.drive_p2p import DriveP2PConnection +from syft.sync.connections.drive.drive_peering import DrivePeering +from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient +from syft.sync.connections.drive.mock_drive_service import ( + MockDriveBackingStore, + MockDriveService, +) from syft.sync.events.file_change_event import ( FileChangeEventsMessage, FileChangeEventsMessageFileName, ) -def create_mock_connection(): - conn = GDriveConnection(email="test@test.com", verbose=False) - conn.drive_service = Mock() - return conn +def create_mock_cached_client() -> DriveCachedClient: + return DriveCachedClient(GDriveClient(Mock(), Mock), "test@test.com") def test_single_page_works(): """Current implementation works for single page (< 100 files)""" - conn = create_mock_connection() + cached_client = create_mock_cached_client() - conn.drive_service.files().list().execute.return_value = { + cached_client.client.drive_service.files().list().execute.return_value = { "files": [{"name": "file1.txt", "id": "id1"}] } - files = conn.get_file_metadatas_from_folder("folder_id") + files = cached_client.get_file_metadatas_from_folder("folder_id") assert len(files) == 1 assert files[0]["name"] == "file1.txt" @@ -29,38 +36,40 @@ def test_single_page_works(): def test_multiple_pages_pagination(): """When there are >100 files, should fetch all pages""" - conn = create_mock_connection() + cached_client = create_mock_cached_client() - conn.drive_service.files().list().execute.side_effect = [ + cached_client.client.drive_service.files().list().execute.side_effect = [ {"files": [{"name": "file1.txt", "id": "id1"}], "nextPageToken": "token1"}, {"files": [{"name": "file2.txt", "id": "id2"}], "nextPageToken": "token2"}, {"files": [{"name": "file3.txt", "id": "id3"}]}, ] - files = conn.get_file_metadatas_from_folder("folder_id") + files = cached_client.get_file_metadatas_from_folder("folder_id") assert len(files) == 3 - assert conn.drive_service.files().list().execute.call_count == 3 + assert cached_client.client.drive_service.files().list().execute.call_count == 3 def test_empty_folder(): """Empty folder should return empty list""" - conn = create_mock_connection() + cached_client = create_mock_cached_client() - conn.drive_service.files().list().execute.return_value = {"files": []} + cached_client.client.drive_service.files().list().execute.return_value = { + "files": [] + } - files = conn.get_file_metadatas_from_folder("folder_id") + files = cached_client.get_file_metadatas_from_folder("folder_id") assert files == [] @patch( - "syft.sync.connections.drive.gdrive_transport.GDriveConnection._get_peer_datasite_outbox_id" + "syft.sync.connections.drive.drive_cached_client.DriveCachedClient.get_peer_datasite_outbox_id" ) -@patch("syft.sync.connections.drive.gdrive_transport.GDriveConnection.download_file") +@patch("syft.sync.connections.drive.gdrive_client.GDriveClient.download_file") def test_events_returned_in_chronological_order(mock_download_file, mock_get_folder_id): """Events should be returned in chronological order by timestamp""" - conn = create_mock_connection() + cached_client = create_mock_cached_client() # Create valid filenames with proper UUIDs uuid1 = uuid4() @@ -68,7 +77,7 @@ def test_events_returned_in_chronological_order(mock_download_file, mock_get_fol uuid3 = uuid4() # Mock Drive returns files in REVERSE chronological order (3.0, 2.0, 1.0) - conn.drive_service.files().list().execute.return_value = { + cached_client.client.drive_service.files().list().execute.return_value = { "files": [ {"id": "id3", "name": f"syfteventsmessagev3_3.0_{uuid3}.tar.gz"}, {"id": "id2", "name": f"syfteventsmessagev3_2.0_{uuid2}.tar.gz"}, @@ -76,7 +85,7 @@ def test_events_returned_in_chronological_order(mock_download_file, mock_get_fol ] } - # Mock _get_peer_datasite_outbox_id + # Mock get_peer_datasite_outbox_id mock_get_folder_id.return_value = "folder_id" # Mock download_file to return events with matching timestamps @@ -93,7 +102,12 @@ def mock_download(file_id): mock_download_file.side_effect = mock_download - events = conn.watcher_get_events_messages("peer@test.com", None) + events = [ + FileChangeEventsMessage.from_compressed_data(data) + for data in DriveP2PConnection(cached_client).outbox_messages( + "peer@test.com", None + ) + ] # Should return in chronological order (1.0, 2.0, 3.0) timestamps = [e.timestamp for e in events] @@ -106,42 +120,22 @@ def mock_download(file_id): def test_orderby_name_desc_applied(): """Drive API should be called with orderBy='name desc'""" - conn = create_mock_connection() + cached_client = create_mock_cached_client() - conn.drive_service.files().list().execute.return_value = {"files": []} + cached_client.client.drive_service.files().list().execute.return_value = { + "files": [] + } - conn.get_file_metadatas_from_folder("folder_id") + cached_client.get_file_metadatas_from_folder("folder_id") # Check that orderBy was passed to Drive API - call_kwargs = conn.drive_service.files().list.call_args[1] + call_kwargs = cached_client.client.drive_service.files().list.call_args[1] assert call_kwargs["orderBy"] == "name desc" -def test_extract_timestamp_from_event_filename(): - """Should extract timestamp from event filename""" - timestamp = GDriveConnection._extract_timestamp_from_filename( - "syfteventsmessagev3_1731506400.123_uuid.tar.gz" - ) - assert timestamp == 1731506400.123 - - -def test_extract_timestamp_from_job_filename(): - """Should extract timestamp from job/message filename""" - timestamp = GDriveConnection._extract_timestamp_from_filename( - "msgv2_1731506400.123_uuid.tar.gz" - ) - assert timestamp == 1731506400.123 - - -def test_extract_timestamp_returns_none_for_invalid(): - """Should return None for non-parseable filenames""" - timestamp = GDriveConnection._extract_timestamp_from_filename("invalid_file.txt") - assert timestamp is None - - def test_early_termination_stops_pagination(): """Should stop fetching pages when encountering old file""" - conn = create_mock_connection() + cached_client = create_mock_cached_client() # Create filenames using FileChangeEventsMessageFileName fname1 = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506450.0) @@ -153,7 +147,7 @@ def test_early_termination_stops_pagination(): fname5 = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506390.0) # Mock returns 3 potential pages, but should stop after page 2 - conn.drive_service.files().list().execute.side_effect = [ + cached_client.client.drive_service.files().list().execute.side_effect = [ { "files": [ {"name": fname1.as_string(), "id": "id1"}, @@ -177,12 +171,12 @@ def test_early_termination_stops_pagination(): }, ] - files = conn.get_file_metadatas_from_folder( + files = cached_client.get_file_metadatas_from_folder( "folder_id", since_timestamp=1731506400.0 ) # Should stop after 2nd page (when it finds a timestamp in an earlier second) - assert conn.drive_service.files().list().execute.call_count == 2 + assert cached_client.client.drive_service.files().list().execute.call_count == 2 # Only files with timestamp > since_timestamp are returned assert len(files) == 3 # fname1, fname2, fname3 (fname4 is excluded) # Page 3 should NOT be fetched @@ -190,12 +184,12 @@ def test_early_termination_stops_pagination(): def test_no_early_termination_when_all_files_new(): """Should fetch all pages when all files are newer than since_timestamp""" - conn = create_mock_connection() + cached_client = create_mock_cached_client() fname1 = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506450.0) fname2 = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506440.0) - conn.drive_service.files().list().execute.side_effect = [ + cached_client.client.drive_service.files().list().execute.side_effect = [ { "files": [{"name": fname1.as_string(), "id": "id1"}], "nextPageToken": "token1", @@ -206,20 +200,20 @@ def test_no_early_termination_when_all_files_new(): }, ] - files = conn.get_file_metadatas_from_folder( + files = cached_client.get_file_metadatas_from_folder( "folder_id", since_timestamp=1731506400.0 ) # Should fetch all pages (both files are > since_timestamp) - assert conn.drive_service.files().list().execute.call_count == 2 + assert cached_client.client.drive_service.files().list().execute.call_count == 2 assert len(files) == 2 def test_no_early_termination_when_since_timestamp_none(): """Should fetch all pages when since_timestamp is None""" - conn = create_mock_connection() + cached_client = create_mock_cached_client() - conn.drive_service.files().list().execute.side_effect = [ + cached_client.client.drive_service.files().list().execute.side_effect = [ { "files": [{"name": "file1.txt", "id": "id1"}], "nextPageToken": "token1", @@ -230,25 +224,68 @@ def test_no_early_termination_when_since_timestamp_none(): }, ] - files = conn.get_file_metadatas_from_folder("folder_id", since_timestamp=None) + files = cached_client.get_file_metadatas_from_folder( + "folder_id", since_timestamp=None + ) # Should fetch all pages (no early termination) - assert conn.drive_service.files().list().execute.call_count == 2 + assert cached_client.client.drive_service.files().list().execute.call_count == 2 assert len(files) == 2 def test_newer_message_listed_after_older_one_in_same_second_not_skipped(): - conn = create_mock_connection() + cached_client = create_mock_cached_client() older = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506400.9) newer = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506400.95) # Drive lists by "name desc": "…400.9_" sorts before "…400.95_". names = sorted([older.as_string(), newer.as_string()], reverse=True) - conn.drive_service.files().list().execute.return_value = { + cached_client.client.drive_service.files().list().execute.return_value = { "files": [{"name": name, "id": name} for name in names] } - files = conn.get_file_metadatas_from_folder( + files = cached_client.get_file_metadatas_from_folder( "folder_id", since_timestamp=1731506400.92 ) assert [f["name"] for f in files] == [newer.as_string()] + + +# Three pages of the mock Drive, which returns 100 files per page. +FILE_COUNT = 250 + + +def _mock_cached_clients(*emails: str) -> list[DriveCachedClient]: + store = MockDriveBackingStore() + return [ + DriveCachedClient.from_service(e, MockDriveService(store, e)) for e in emails + ] + + +def test_collection_listings_read_every_page(): + do_cached_client, ds_cached_client = _mock_cached_clients( + "do@example.com", "ds@example.com" + ) + owner, reader = ( + DriveCollections(do_cached_client), + DriveCollections(ds_cached_client), + ) + tags = [f"tag{i}" for i in range(FILE_COUNT)] + for tag in tags: + owner.create("syft_datasetcollection", tag, "hash", "do@example.com") + owner.grant_read("syft_datasetcollection", tag, "hash", ["ds@example.com"]) + + assert sorted(owner.list_mine("syft_datasetcollection")) == sorted(tags) + mine = owner.list_mine_with_permissions("syft_datasetcollection") + assert len(mine) == FILE_COUNT + shared = reader.list_shared_with_me("syft_datasetcollection") + assert len(shared) == FILE_COUNT + + +def test_peer_requests_read_every_page(): + peer_emails = [f"peer{i}@example.com" for i in range(FILE_COUNT)] + do_cached_client, *peers = _mock_cached_clients("do@example.com", *peer_emails) + for peer in peers: + DrivePeering(peer).open_channels("do@example.com") + + requests = DrivePeering(do_cached_client).incoming_requests() + assert sorted(requests) == sorted(peer_emails) diff --git a/tests/unit/test_login_email_check.py b/tests/unit/test_login_email_check.py index 6eb060be798..8cbb29f71ac 100644 --- a/tests/unit/test_login_email_check.py +++ b/tests/unit/test_login_email_check.py @@ -2,7 +2,7 @@ from syft.sync.login import _verify_token_matches_email from syft.sync.syftbox_manager import SyftboxManager -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of def test_get_authenticated_email_matches_constructor_email(): @@ -18,7 +18,9 @@ def test_verify_token_matches_email_passes_when_matching(): def _drive_service(manager): - return gdrive_connection_of(manager._connection_router.transport).drive_service + return drive_cached_client_of( + manager._connection_router.transport + ).client.drive_service def test_verify_token_matches_email_raises_on_mismatch(): diff --git a/tests/unit/test_mock_drive_service.py b/tests/unit/test_mock_drive_service.py index ea40cd20d4c..2c9ceb197e9 100644 --- a/tests/unit/test_mock_drive_service.py +++ b/tests/unit/test_mock_drive_service.py @@ -1,7 +1,7 @@ """Tests for MockDriveService and pair_with_mock_drive_service_connection. These tests verify that the mock Google Drive service provides a compatible -interface for testing GDriveConnection code paths without real API calls. +interface for testing the Drive code paths without real API calls. """ from syft.sync.syftbox_manager import SyftboxManager @@ -13,7 +13,7 @@ parse_gdrive_query, GOOGLE_FOLDER_MIME_TYPE, ) -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient class TestMockDriveBackingStore: @@ -431,25 +431,25 @@ def test_shared_file_access(self): assert len(list_result["files"]) == 1 -class TestGDriveConnectionWithMock: - """Tests for GDriveConnection with mock service.""" +class TestDriveCachedClientWithMock: + """Tests for DriveCachedClient with mock service.""" def test_from_mock_service(self): store = MockDriveBackingStore() mock_service = MockDriveService(store, "user@example.com") - connection = GDriveConnection.from_service("user@example.com", mock_service) + cached_client = DriveCachedClient.from_service("user@example.com", mock_service) - assert connection.email == "user@example.com" - assert connection._is_setup is True + assert cached_client.email == "user@example.com" + assert cached_client.client.drive_service._backing_store is store def test_create_folder(self): store = MockDriveBackingStore() mock_service = MockDriveService(store, "user@example.com") - connection = GDriveConnection.from_service("user@example.com", mock_service) + cached_client = DriveCachedClient.from_service("user@example.com", mock_service) - syftbox_id = connection.create_syftbox_folder() - folder_id = connection.create_folder("TestFolder", syftbox_id) + syftbox_id = cached_client.create_syftbox_folder() + folder_id = cached_client.client.create_folder("TestFolder", syftbox_id) assert store.files[syftbox_id].properties == {"syft": "true"} assert store.files[folder_id].properties == {"syft": "true"} @@ -466,9 +466,9 @@ def test_download_file(self): store.add_file(file) mock_service = MockDriveService(store, "user@example.com") - connection = GDriveConnection.from_service("user@example.com", mock_service) + cached_client = DriveCachedClient.from_service("user@example.com", mock_service) - content = connection.download_file(file.id) + content = cached_client.client.download_file(file.id) assert content == b"Test content" diff --git a/tests/unit/test_outbox_compaction.py b/tests/unit/test_outbox_compaction.py index ff43ddfbfe1..acebbac9f4b 100644 --- a/tests/unit/test_outbox_compaction.py +++ b/tests/unit/test_outbox_compaction.py @@ -3,7 +3,7 @@ from syft.sync.events.file_change_event import FileChangeEventsMessage from syft.sync.syftbox_manager import SyftboxManager from tests.unit.utils import get_mock_event -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of def test_compact_outboxes_if_needed_collapses_to_single_message(): @@ -21,13 +21,17 @@ def test_compact_outboxes_if_needed_collapses_to_single_message(): ds_manager.email, msg ) - conn_do = gdrive_connection_of(do_manager._connection_router.transport) - outbox_id = conn_do._get_own_datasite_outbox_id(ds_manager.email) - pre = conn_do.get_file_metadatas_from_folder(outbox_id, since_timestamp=None) + do_cached_client = drive_cached_client_of(do_manager._connection_router.transport) + outbox_id = do_cached_client.get_own_datasite_outbox_id(ds_manager.email) + pre = do_cached_client.get_file_metadatas_from_folder( + outbox_id, since_timestamp=None + ) assert len(pre) == n_messages result = do_manager.compact_outboxes_if_needed() assert result[ds_manager.email] == n_messages - post = conn_do.get_file_metadatas_from_folder(outbox_id, since_timestamp=None) + post = do_cached_client.get_file_metadatas_from_folder( + outbox_id, since_timestamp=None + ) assert len(post) == 1 diff --git a/tests/unit/test_p2p_folder_lookup.py b/tests/unit/test_p2p_folder_lookup.py index 3714e55ec8f..afa8c7d2c2b 100644 --- a/tests/unit/test_p2p_folder_lookup.py +++ b/tests/unit/test_p2p_folder_lookup.py @@ -11,7 +11,8 @@ from unittest.mock import Mock -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient +from syft.sync.connections.drive.gdrive_client import GDriveClient ME = "alice@example.com" PEER = "bob@example.com" @@ -21,15 +22,14 @@ def _name(version: str, datasite: str, folder_type: str, peer: str) -> str: return f"syft_datasite#{version}#{datasite}#{folder_type}#{peer}" -def _conn(found): - conn = GDriveConnection(email=ME, verbose=False) - conn.drive_service = Mock() - conn._find_folders = Mock(return_value=found) - return conn +def _cached_client(found): + cached_client = DriveCachedClient(GDriveClient(Mock(), Mock), ME) + cached_client.client.find_folders = Mock(return_value=found) + return cached_client -def _lookup(conn): - return conn._find_p2p_folder_id( +def _lookup(cached_client): + return cached_client.find_p2p_folder_id( datasite_email=PEER, folder_type="inbox", peer_email=ME, owner_email=ME ) @@ -39,12 +39,12 @@ def test_a_folder_of_another_minor_version_is_found(): # the peer kept writing into the first. 0.2.0 differs in the minor from the # current client version, which is what the filter used to reject. old = _name("0.2.0", PEER, "inbox", ME) - assert _lookup(_conn([("old", old)])) == "old" + assert _lookup(_cached_client([("old", old)])) == "old" def test_a_folder_of_an_older_major_version_is_found(): old = _name("0.0.9", PEER, "inbox", ME) - assert _lookup(_conn([("old", old)])) == "old" + assert _lookup(_cached_client([("old", old)])) == "old" def test_the_highest_version_wins_when_several_exist(): @@ -53,7 +53,7 @@ def test_the_highest_version_wins_when_several_exist(): ("v2", _name("0.2.0", PEER, "inbox", ME)), ("v0", _name("0.0.9", PEER, "inbox", ME)), ] - assert _lookup(_conn(folders)) == "v2" + assert _lookup(_cached_client(folders)) == "v2" def test_versions_order_by_number_not_by_string(): @@ -61,7 +61,7 @@ def test_versions_order_by_number_not_by_string(): ("nine", _name("0.1.9", PEER, "inbox", ME)), ("ten", _name("0.1.10", PEER, "inbox", ME)), ] - assert _lookup(_conn(folders)) == "ten" + assert _lookup(_cached_client(folders)) == "ten" def test_several_folders_no_longer_raise(): @@ -69,22 +69,22 @@ def test_several_folders_no_longer_raise(): ("a", _name("0.1.117", PEER, "inbox", ME)), ("b", _name("0.1.118", PEER, "inbox", ME)), ] - assert _lookup(_conn(folders)) is not None + assert _lookup(_cached_client(folders)) is not None def test_a_folder_of_another_peer_is_ignored(): other = _name("0.1.117", PEER, "inbox", "carol@example.com") - assert _lookup(_conn([("other", other)])) is None + assert _lookup(_cached_client([("other", other)])) is None def test_a_folder_of_another_type_is_ignored(): outbox = _name("0.1.117", PEER, "outbox", ME) - assert _lookup(_conn([("outbox", outbox)])) is None + assert _lookup(_cached_client([("outbox", outbox)])) is None def test_no_folder_returns_none(): - assert _lookup(_conn([])) is None + assert _lookup(_cached_client([])) is None def test_a_name_that_does_not_parse_is_ignored(): - assert _lookup(_conn([("junk", "not_a_p2p_folder")])) is None + assert _lookup(_cached_client([("junk", "not_a_p2p_folder")])) is None diff --git a/tests/unit/test_peer_request_owner.py b/tests/unit/test_peer_request_owner.py index 36df24d72e2..c4b5c60ceb0 100644 --- a/tests/unit/test_peer_request_owner.py +++ b/tests/unit/test_peer_request_owner.py @@ -6,9 +6,8 @@ import pytest -from syft.sync.connections.drive.gdrive_transport import ( +from syft.sync.connections.drive.drive_cached_client import ( GDRIVE_P2P_FOLDER_DATASITE_PREFIX, - PEERS_META_KEY, GdriveP2PFolder, ) from syft.sync.connections.drive.mock_drive_service import ( @@ -17,9 +16,10 @@ MockPermission, ) from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.transport.naming import PEERS_META_KEY from syft.sync.version.peer_manager import PeerState from syft.version import SYFT_VERSION -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of DO_EMAIL = "do@test.com" DS_EMAIL = "ds@test.com" @@ -47,9 +47,9 @@ def share_request_folder( victim: SyftboxManager, owner_email: str | None, claimed_email: str ) -> None: """Put a request folder for `victim` in Drive and share it with `victim`.""" - backing_store = gdrive_connection_of( + backing_store = drive_cached_client_of( victim._connection_router.transport - ).drive_service._backing_store + ).client.drive_service._backing_store owners = [{"emailAddress": owner_email}] if owner_email else [] folder = MockDriveFile( name=p2p_name(victim.email, claimed_email), diff --git a/tests/unit/test_peers_json_version.py b/tests/unit/test_peers_json_version.py index 179c41fa38c..b887422f206 100644 --- a/tests/unit/test_peers_json_version.py +++ b/tests/unit/test_peers_json_version.py @@ -6,7 +6,7 @@ fails, and skips it, so the reserved key is invisible to a client that predates it. -The record itself is safe either way. The only writer is `_update_peer_state`, +The record itself is safe either way. The only writer is `update_peer_registry`, which changes one entry of the raw map and writes the rest back, so a peer this client cannot read is not erased for the other side. """ @@ -14,70 +14,68 @@ import logging from unittest.mock import Mock, patch -from syft.sync.connections.drive.gdrive_transport import ( - PEERS_META_KEY, - SYFT_PEERS_VERSION, - GDriveConnection, -) +from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient +from syft.sync.connections.drive.drive_peering import DrivePeering from syft.sync.peers.peer import PeerState -from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform +from syft.sync.transport.naming import PEERS_META_KEY, SYFT_PEERS_VERSION PEER = "bob@example.com" -def _conn(peers_data): - conn = GDriveConnection(email="alice@example.com", verbose=False) - conn.drive_service = Mock() - conn._peers_json_cache = dict(peers_data) - return conn +def _peering(peers_data): + cached_client = DriveCachedClient(GDriveClient(Mock(), Mock), "alice@example.com") + peering = DrivePeering(cached_client) + peering._peers_json_cache = dict(peers_data) + return peering -def _router(conn): +def _router(peering): router = Mock() - router.transport = SyftboxManager._drive_transport(conn) + router.transport.peering = peering + router.transport.platform = GdriveFilesPlatform() from syft.sync.connections.connection_router import ConnectionRouter return ConnectionRouter.get_all_peers_from_json.__get__(router, ConnectionRouter) def test_a_write_stamps_the_reserved_entry(): - conn = _conn({PEER: {"state": "accepted"}}) + peering = _peering({PEER: {"state": "accepted"}}) with ( - patch.object(GDriveConnection, "_get_peers_file_id", return_value="file-id"), + patch.object(DrivePeering, "_get_peers_file_id", return_value="file-id"), patch.object( - GDriveConnection, "get_syftbox_folder_id", return_value="folder-id" - ), - patch.object( - GDriveConnection, "create_file_payload", return_value=(Mock(), None) + DriveCachedClient, "get_syftbox_folder_id", return_value="folder-id" ), + patch.object(GDriveClient, "create_file_payload", return_value=(Mock(), None)), ): - conn._write_peers_json({PEER: {"state": "accepted"}}) + peering._write_peers_json({PEER: {"state": "accepted"}}) - assert conn._peers_json_cache[PEERS_META_KEY] == {"version": SYFT_PEERS_VERSION} - assert conn._peers_json_cache[PEER] == {"state": "accepted"} + assert peering._peers_json_cache[PEERS_META_KEY] == {"version": SYFT_PEERS_VERSION} + assert peering._peers_json_cache[PEER] == {"state": "accepted"} def test_the_reserved_entry_is_not_a_peer(): - conn = _conn( + peering = _peering( { PEERS_META_KEY: {"version": SYFT_PEERS_VERSION}, PEER: {"state": "accepted"}, } ) - peers = _router(conn)() + peers = _router(peering)() assert [p.email for p in peers] == [PEER] def test_a_known_state_loads(): - conn = _conn({PEER: {"state": "rejected"}}) - peers = _router(conn)() + peering = _peering({PEER: {"state": "rejected"}}) + peers = _router(peering)() assert peers[0].state == PeerState.REJECTED def test_an_unknown_state_is_skipped_and_logged(caplog): - conn = _conn({PEER: {"state": "quarantined"}}) + peering = _peering({PEER: {"state": "quarantined"}}) with caplog.at_level(logging.WARNING, logger="syft"): - peers = _router(conn)() + peers = _router(peering)() assert peers == [] assert any(PEER in r.getMessage() for r in caplog.records) assert any("quarantined" in r.getMessage() for r in caplog.records) @@ -85,18 +83,18 @@ def test_an_unknown_state_is_skipped_and_logged(caplog): def test_a_file_without_the_reserved_entry_still_loads(): # Written before the reserved key existed. - conn = _conn({PEER: {"state": "accepted"}}) - peers = _router(conn)() + peering = _peering({PEER: {"state": "accepted"}}) + peers = _router(peering)() assert [p.email for p in peers] == [PEER] def test_a_reserved_entry_from_a_newer_client_does_not_stop_the_read(caplog): - conn = _conn( + peering = _peering( { PEERS_META_KEY: {"version": SYFT_PEERS_VERSION + 1}, PEER: {"state": "accepted"}, } ) with caplog.at_level(logging.WARNING, logger="syft"): - peers = _router(conn)() + peers = _router(peering)() assert [p.email for p in peers] == [PEER] diff --git a/tests/unit/test_rolling_state.py b/tests/unit/test_rolling_state.py index 1ef1dbce756..b3723dedf20 100644 --- a/tests/unit/test_rolling_state.py +++ b/tests/unit/test_rolling_state.py @@ -7,10 +7,25 @@ """ import time + +import pytest +from googleapiclient.errors import HttpError + +from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore +from syft.sync.checkpoints.rolling_state import ( + ROLLING_STATE_FILENAME_PREFIX, + RollingState, +) +from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.mock_drive_service import ( + MockListRequest, + http_error, + pair_with_mock_service, +) from syft.sync.syftbox_manager import SyftboxManager -from syft.sync.checkpoints.rolling_state import RollingState from syft.sync.sync.constants import CACHE_DIR, ROLLING_STATE_FILENAME -from tests.unit.utils import get_mock_event +from syft.sync.transport import TransportUnavailableError +from tests.unit.utils import fail_next_listing, get_mock_event def test_upload_and_get_rolling_state(): @@ -270,3 +285,51 @@ def test_no_rolling_state_when_owner_state_disabled(): assert not (syncer.syftbox_folder / CACHE_DIR / ROLLING_STATE_FILENAME).exists() # Nothing is tracked in memory either - there is nothing to track it for. assert syncer._rolling_state is None + + +def test_rolling_state_read_retries_drive_error(monkeypatch): + do_cached_client, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = DriveOwnerStore(do_cached_client) + store.put_rolling_state("rolling_state_1.0.tar.gz", b"state") + + fail_next_listing(monkeypatch, ROLLING_STATE_FILENAME_PREFIX) + assert store.get_rolling_state() == b"state" + + +def test_rolling_state_read_while_drive_is_down_raises(monkeypatch): + do_cached_client, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = DriveOwnerStore(do_cached_client) + store.put_rolling_state("rolling_state_1.0.tar.gz", b"state") + + def drive_down(self, file_id: str) -> bytes: + raise ConnectionResetError("Drive is down") + + monkeypatch.setattr(GDriveClient, "download_file", drive_down) + with pytest.raises(TransportUnavailableError): + store.get_rolling_state() + + +def test_rolling_state_listing_while_drive_is_down_raises(monkeypatch): + do_cached_client, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = DriveOwnerStore(do_cached_client) + store.put_rolling_state("rolling_state_1.0.tar.gz", b"state") + + def drive_down(self) -> dict: + raise ConnectionRefusedError("Drive is down") + + monkeypatch.setattr(MockListRequest, "execute", drive_down) + with pytest.raises(TransportUnavailableError): + store.get_rolling_state() + + +def test_rolling_state_download_error_raises(monkeypatch): + do_cached_client, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = DriveOwnerStore(do_cached_client) + store.put_rolling_state("rolling_state_1.0.tar.gz", b"state") + + def forbidden(self, file_id: str) -> bytes: + raise http_error(403, "cannotDownloadFile", "The file cannot be downloaded") + + monkeypatch.setattr(GDriveClient, "download_file", forbidden) + with pytest.raises(HttpError): + store.get_rolling_state() diff --git a/tests/unit/test_sender_binding.py b/tests/unit/test_sender_binding.py index cd084e4dc75..87929d02ed1 100644 --- a/tests/unit/test_sender_binding.py +++ b/tests/unit/test_sender_binding.py @@ -12,7 +12,7 @@ ProposedFileChangesMessage, ) from syft.sync.syftbox_manager import SyftboxManager -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of def send_spoofed_message( @@ -46,8 +46,8 @@ def sync_do(do_manager: SyftboxManager) -> None: def written_paths(do_manager: SyftboxManager) -> list[str]: - connection = do_manager.datasite_owner_syncer.event_cache.file_connection - return [str(path) for path, _ in connection.get_items()] + do_cached_client = do_manager.datasite_owner_syncer.event_cache.file_connection + return [str(path) for path, _ in do_cached_client.get_items()] def test_sender_email_in_body_cannot_claim_the_datasite_owner(): @@ -81,11 +81,12 @@ def test_spoofed_sender_does_not_name_an_archive_folder(): send_spoofed_message(ds_manager, do_email, claimed_email, "any.txt") sync_do(do_manager) - connection = gdrive_connection_of( + do_cached_client = drive_cached_client_of( do_manager.datasite_owner_syncer.connection_router.transport ) folder_names = [ - f.name for f in connection.drive_service._backing_store.files.values() + f.name + for f in do_cached_client.client.drive_service._backing_store.files.values() ] assert not [name for name in folder_names if claimed_email in name] diff --git a/tests/unit/test_shared_router.py b/tests/unit/test_shared_router.py index 288e32c1325..add45eb5bf8 100644 --- a/tests/unit/test_shared_router.py +++ b/tests/unit/test_shared_router.py @@ -6,8 +6,11 @@ import pytest from syft.sync.connections.base_connection import ConnectionConfig -from syft.sync.connections.drive.drive_connections import gdrive_connection_of -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.drive_connections import ( + DriveWatcherStore, + drive_cached_client_of, +) +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.grdrive_config import GdriveConnectionConfig from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service from syft.sync.syftbox_manager import SyftboxManager, SyftboxManagerConfig @@ -25,14 +28,14 @@ def _drive_config(email: str) -> GdriveConnectionConfig: return GdriveConnectionConfig(email=email, token_path=None) -def _mock_connection(email: str) -> GDriveConnection: - connection, _ = pair_with_mock_service(email, "peer@example.com") - return connection +def _mock_cached_client(email: str) -> DriveCachedClient: + cached_client, _ = pair_with_mock_service(email, "peer@example.com") + return cached_client def test_every_part_of_manager_shares_one_router(): config = _config_with_both_roles() - transport = SyftboxManager._drive_transport(_mock_connection(config.email)) + transport = SyftboxManager._drive_transport(_mock_cached_client(config.email)) manager = SyftboxManager.from_config(config, transport) router = manager._connection_router @@ -46,40 +49,50 @@ def test_every_part_of_manager_shares_one_router(): assert manager._peer_store is router.peer_store +def _drive_cached_clients(transport) -> list[DriveCachedClient]: + """The cached client of each component that uses Drive. The watcher store does not + use it yet.""" + return [ + component._cached_client + for component in transport.components() + if not isinstance(component, DriveWatcherStore) + ] + + def test_manager_builds_one_drive_connection_from_its_config(monkeypatch): built = [] - def build_mock_connection(connection_config): - built.append(_mock_connection(connection_config.email)) + def build_mock_cached_client(connection_config): + built.append(_mock_cached_client(connection_config.email)) return built[-1] - monkeypatch.setattr(GDriveConnection, "from_config", build_mock_connection) + monkeypatch.setattr(DriveCachedClient, "from_config", build_mock_cached_client) config = _config_with_both_roles() config.transport_config = TransportConfig.one_connection( _drive_config(config.email) ) manager = SyftboxManager.from_config(config) - [connection] = built - for component in manager._connection_router.transport.components(): - assert component._conn is connection + [built_client] = built + for cached_client in _drive_cached_clients(manager._connection_router.transport): + assert cached_client is built_client -def test_every_transport_component_shares_one_mock_connection(): +def test_every_drive_component_shares_one_cached_client(): _, do_manager = SyftboxManager.pair_with_mock_drive_service_connection() transport = do_manager._connection_router.transport - connection = gdrive_connection_of(transport) + cached_client = drive_cached_client_of(transport) - for component in transport.components(): - assert component._conn is connection + for component_cached_client in _drive_cached_clients(transport): + assert component_cached_client is cached_client -def test_gdrive_connection_of_transport_not_on_drive_raises(): - transport = SyftboxManager._drive_transport(_mock_connection("do@example.com")) +def test_drive_cached_client_of_transport_not_on_drive_raises(): + transport = SyftboxManager._drive_transport(_mock_cached_client("do@example.com")) transport = transport.model_copy(update={"owner_store": Mock(spec=OwnerStore)}) with pytest.raises(TypeError, match="not built on Google Drive"): - gdrive_connection_of(transport) + drive_cached_client_of(transport) def test_manager_without_transport_config_fails_at_init(): diff --git a/tests/unit/test_sync_manager.py b/tests/unit/test_sync_manager.py index a9f763cdaed..07afb3083c0 100644 --- a/tests/unit/test_sync_manager.py +++ b/tests/unit/test_sync_manager.py @@ -3,8 +3,9 @@ import pytest import yaml +from syft.sync.connections.drive.drive_peering import DrivePeering from syft.sync.connections.connection_router import ConnectionRouter -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.encryption_policy import EncryptionPolicy from syft.sync.peers.peer import Peer from syft.sync.peers.peer_store import PeerStore @@ -24,7 +25,7 @@ get_mock_proposed_events_messages, grant_job_inbox_access, ) -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of def path_for_job(do_email: str, ds_email: str, filename: str = "test.job") -> str: @@ -149,7 +150,7 @@ def test_sync_existing_datasite_state_do(): ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( use_in_memory_cache=False ) - connection_ds = gdrive_connection_of(ds_manager._connection_router.transport) + ds_p2p = ds_manager._connection_router.transport.p2p events_messages = get_mock_events_messages(2) for message in events_messages: @@ -168,7 +169,7 @@ def test_sync_existing_datasite_state_do(): n_files_in_cache = len(do_manager.datasite_owner_syncer.event_cache.file_connection) hashes_in_cache = len(do_manager.datasite_owner_syncer.event_cache.file_hashes) - n_outbox = connection_ds.watcher_get_outbox_file_metadatas(do_manager.email, None) + n_outbox = ds_p2p.outbox_message_metadatas(do_manager.email, None) assert n_messages_in_cache >= 1 # At least 1 message with the 2 file changes # 2 data files + private/crypto_keys.json + private/syft_settings.json assert n_files_in_cache == 4 @@ -187,13 +188,11 @@ def test_sync_existing_inbox_state_do(): do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( ds_manager.email ) - connection_ds = gdrive_connection_of(ds_manager._connection_router.transport) - proposed_events_messages = get_mock_proposed_events_messages( 2, email=ds_manager.email ) for message in proposed_events_messages: - connection_ds.watcher_send_proposed_file_changes_message( + ds_manager._connection_router.watcher_send_proposed_file_changes_message( do_manager.email, message ) @@ -487,14 +486,16 @@ def _add_known_peer(do_manager, email: str) -> None: def _peer_router(ds_manager, do_manager, email: str) -> ConnectionRouter: """Make ``email`` a second peer of the DO, and return its router.""" - conn = GDriveConnection.from_service( + cached_client = DriveCachedClient.from_service( email, - gdrive_connection_of(ds_manager._connection_router.transport).drive_service, + drive_cached_client_of( + ds_manager._connection_router.transport + ).client.drive_service, ) - conn.add_peer(do_manager.email) + DrivePeering(cached_client).open_channels(do_manager.email) _add_known_peer(do_manager, email) return ConnectionRouter( - transport=SyftboxManager._drive_transport(conn), + transport=SyftboxManager._drive_transport(cached_client), peer_store=_store_knowing(email, do_manager.email), ) diff --git a/tests/unit/test_version_negotiation.py b/tests/unit/test_version_negotiation.py index 7211011c447..d1d0d19f0b9 100644 --- a/tests/unit/test_version_negotiation.py +++ b/tests/unit/test_version_negotiation.py @@ -5,7 +5,7 @@ import httplib2 import pytest from googleapiclient.errors import HttpError -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.syftbox_manager import SyftboxManager from syft.sync.transport import TransportUnavailableError from syft.sync.version.exceptions import ( @@ -15,7 +15,7 @@ from syft.sync.version.peer_info import PeerInfo from syft.sync.version.peer_manager import CompatAction from syft.sync.version.version_info import CompatibilityStatus, VersionInfo -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of def build_client_version(client_version: str) -> VersionInfo: @@ -228,7 +228,7 @@ def test_load_peer_version_without_permission(self): assert version is None def test_load_peer_versions_parallel(self): - from syft.sync.connections.drive.gdrive_transport import ( + from syft.sync.connections.drive.drive_peering import ( SYFT_VERSION_FILE, ) from syft.sync.connections.drive.mock_drive_service import ( @@ -240,9 +240,9 @@ def test_load_peer_versions_parallel(self): SyftboxManager.pair_with_mock_drive_service_connection() ) - backing_store = gdrive_connection_of( + backing_store = drive_cached_client_of( ds_manager._connection_router.transport - ).drive_service._backing_store + ).client.drive_service._backing_store third_peer_email = "third_peer@test.com" third_peer_info = PeerInfo( @@ -284,7 +284,7 @@ def test_drive_error_during_peer_version_read_raises(self, monkeypatch): ds_manager, do_manager = ( SyftboxManager.pair_with_mock_drive_service_connection() ) - monkeypatch.setattr(GDriveConnection, "download_file", _raises(DRIVE_DOWN)) + monkeypatch.setattr(GDriveClient, "download_file", _raises(DRIVE_DOWN)) with pytest.raises(TransportUnavailableError): ds_manager.peer_manager.load_peer_infos_parallel([do_manager.email]) @@ -293,14 +293,14 @@ def test_drive_error_during_peer_version_read_aborts_sync(self, monkeypatch): ds_manager, do_manager = ( SyftboxManager.pair_with_mock_drive_service_connection() ) - monkeypatch.setattr(GDriveConnection, "download_file", _raises(DRIVE_DOWN)) + monkeypatch.setattr(GDriveClient, "download_file", _raises(DRIVE_DOWN)) with pytest.raises(TransportUnavailableError): do_manager.sync() def test_drive_error_during_own_version_read_raises(self, monkeypatch): _, do_manager = SyftboxManager.pair_with_mock_drive_service_connection() - monkeypatch.setattr(GDriveConnection, "download_file", _raises(DRIVE_DOWN)) + monkeypatch.setattr(GDriveClient, "download_file", _raises(DRIVE_DOWN)) with pytest.raises(TransportUnavailableError): do_manager.peer_manager.connection_router.read_own_peer_info_file() @@ -310,7 +310,7 @@ def test_offline_during_peer_version_read_raises_unavailable(self, monkeypatch): SyftboxManager.pair_with_mock_drive_service_connection() ) no_dns = httplib2.ServerNotFoundError("no DNS") - monkeypatch.setattr(GDriveConnection, "download_file", _raises(no_dns)) + monkeypatch.setattr(GDriveClient, "download_file", _raises(no_dns)) with pytest.raises(TransportUnavailableError): ds_manager.peer_manager.load_peer_infos_parallel([do_manager.email]) @@ -323,7 +323,7 @@ def test_drive_refusal_during_peer_version_read_skips_peer(self, monkeypatch): httplib2.Response({"status": 403}), b'{"error": {"code": 403, "errors": [{"reason": "forbidden"}]}}', ) - monkeypatch.setattr(GDriveConnection, "download_file", _raises(forbidden)) + monkeypatch.setattr(GDriveClient, "download_file", _raises(forbidden)) infos = ds_manager.peer_manager.load_peer_infos_parallel([do_manager.email]) @@ -333,7 +333,7 @@ def test_unreadable_peer_version_file_skips_peer(self, monkeypatch): ds_manager, do_manager = ( SyftboxManager.pair_with_mock_drive_service_connection() ) - monkeypatch.setattr(GDriveConnection, "download_file", _returns(b"{trunc")) + monkeypatch.setattr(GDriveClient, "download_file", _returns(b"{trunc")) infos = ds_manager.peer_manager.load_peer_infos_parallel([do_manager.email]) diff --git a/tests/unit/test_versioned_folder_adopt.py b/tests/unit/test_versioned_folder_adopt.py index 8b899d4e0f3..fbb2cad084e 100644 --- a/tests/unit/test_versioned_folder_adopt.py +++ b/tests/unit/test_versioned_folder_adopt.py @@ -19,28 +19,27 @@ import pytest from syft.sync.checkpoints.checkpoint import CHECKPOINT_VERSION, Checkpoint -from syft.sync.connections.drive.gdrive_transport import ( - GDriveConnection, +from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.drive_cached_client import ( + DriveCachedClient, _partition_by_version, - execute_with_retries, ) +from syft.sync.connections.drive.gdrive_retry import execute_with_retries from syft.sync.syftbox_manager import SyftboxManager -from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of EMAIL = "alice@example.com" -def _conn(): - conn = GDriveConnection(email=EMAIL, verbose=False) - conn.drive_service = Mock() - return conn +def _cached_client(): + return DriveCachedClient(GDriveClient(Mock(), Mock), EMAIL) -def _renames(conn): - """Return the (fileId, new name) pairs the connection sent to Drive.""" +def _renames(cached_client): + """Return the (fileId, new name) pairs the cached client sent to Drive.""" return [ (kwargs["fileId"], kwargs["body"]["name"]) - for _, kwargs in conn.drive_service.files().update.call_args_list + for _, kwargs in cached_client.client.drive_service.files().update.call_args_list if "body" in kwargs and "name" in kwargs.get("body", {}) ] @@ -81,65 +80,65 @@ def test_partition_returns_empty_for_a_bad_current_version(): def test_a_compatible_folder_wins_and_nothing_is_renamed(): - conn = _conn() + cached_client = _cached_client() folders = [("same", f"0.2.5#{EMAIL}"), ("old", f"0.1.9#{EMAIL}")] - got = conn._find_or_adopt_versioned_folder( + got = cached_client.find_or_adopt_versioned_folder( folders, current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) assert got == "same" - assert _renames(conn) == [] + assert _renames(cached_client) == [] def test_an_older_folder_is_adopted_by_rename(): - conn = _conn() + cached_client = _cached_client() folders = [("old", f"0.1.9#{EMAIL}")] - got = conn._find_or_adopt_versioned_folder( + got = cached_client.find_or_adopt_versioned_folder( folders, current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) assert got == "old", "the client must keep the folder that holds the data" - assert _renames(conn) == [("old", f"0.2.7#{EMAIL}")] + assert _renames(cached_client) == [("old", f"0.2.7#{EMAIL}")] def test_the_highest_older_folder_is_adopted(): - conn = _conn() + cached_client = _cached_client() folders = [ ("v1", f"0.1.9#{EMAIL}"), ("v2", f"0.1.20#{EMAIL}"), ("v0", f"0.0.4#{EMAIL}"), ] - got = conn._find_or_adopt_versioned_folder( + got = cached_client.find_or_adopt_versioned_folder( folders, current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) assert got == "v2" - assert _renames(conn) == [("v2", f"0.2.7#{EMAIL}")] + assert _renames(cached_client) == [("v2", f"0.2.7#{EMAIL}")] def test_a_newer_folder_stops_the_client(): # A new folder here would hide data that this client cannot read. Report the # version to install instead. - conn = _conn() + cached_client = _cached_client() folders = [("new", f"0.3.0#{EMAIL}")] with pytest.raises(RuntimeError, match="0.3.0"): - conn._find_or_adopt_versioned_folder( + cached_client.find_or_adopt_versioned_folder( folders, current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) - assert _renames(conn) == [] + assert _renames(cached_client) == [] def test_no_folder_returns_none_so_the_caller_creates_one(): - conn = _conn() - got = conn._find_or_adopt_versioned_folder( + cached_client = _cached_client() + got = cached_client.find_or_adopt_versioned_folder( [], current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) assert got is None - assert _renames(conn) == [] + assert _renames(cached_client) == [] def test_two_compatible_folders_still_raise(): - conn = _conn() + cached_client = _cached_client() folders = [("a", f"0.2.1#{EMAIL}"), ("b", f"0.2.2#{EMAIL}")] with pytest.raises(RuntimeError): - conn._find_or_adopt_versioned_folder( + cached_client.find_or_adopt_versioned_folder( folders, current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) @@ -151,7 +150,7 @@ def _checkpoints_owner_with_version(checkpoint_version: int): """A DO whose checkpoints folder is named for an earlier client version. Uploads one checkpoint at `checkpoint_version`, then renames the folder as a - pre-upgrade client would have left it. Returns (router, connection, + pre-upgrade client would have left it. Returns (router, cached_client, older_name). """ _, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( @@ -159,27 +158,27 @@ def _checkpoints_owner_with_version(checkpoint_version: int): sync_automatically=False, ) router = do_manager._connection_router - connection = gdrive_connection_of(router.transport) + cached_client = drive_cached_client_of(router.transport) router.upload_checkpoint( Checkpoint(email=do_manager.email, version=checkpoint_version) ) - folder_id = connection._get_checkpoints_folder_id() + folder_id = router.transport.owner_store._get_checkpoints_folder_id() older_name = f"{do_manager.email}-0.0.1-checkpoints" execute_with_retries( - connection.drive_service.files().update( + cached_client.client.drive_service.files().update( fileId=folder_id, body={"name": older_name} ) ) - return router, connection, older_name + return router, cached_client, older_name -def _checkpoints_folder_names(connection) -> list[str]: +def _checkpoints_folder_names(cached_client) -> list[str]: return [ name - for _, name in connection._find_folders( - name_contains=[f"{connection.email}-", "-checkpoints"], - parent_id=connection.get_syftbox_folder_id(), + for _, name in cached_client.client.find_folders( + name_contains=[f"{cached_client.email}-", "-checkpoints"], + parent_id=cached_client.get_syftbox_folder_id(), ) ] @@ -187,16 +186,16 @@ def _checkpoints_folder_names(connection) -> list[str]: def test_an_adopted_folder_serves_a_checkpoint_from_an_earlier_client(): # Version 0 predates the field, which is the shape an earlier client wrote. # The rename finds the folder and the file still loads. - router, connection, older_name = _checkpoints_owner_with_version(0) + router, cached_client, older_name = _checkpoints_owner_with_version(0) checkpoint = router.get_latest_checkpoint() assert checkpoint is not None, "the adopted folder must still serve its files" assert checkpoint.version == 0 - assert _checkpoints_folder_names(connection) == [ - connection._get_checkpoints_folder_name() + assert _checkpoints_folder_names(cached_client) == [ + router.transport.owner_store._get_checkpoints_folder_name() ], "the read must go through the adopt, which renames the folder" - assert older_name not in _checkpoints_folder_names(connection) + assert older_name not in _checkpoints_folder_names(cached_client) def test_an_adopted_folder_still_refuses_a_later_checkpoint(capsys): @@ -204,7 +203,7 @@ def test_an_adopted_folder_still_refuses_a_later_checkpoint(capsys): # the rename does not vouch for the content. A checkpoint from a later client # is refused inside an adopted folder exactly as it is anywhere else, and the # caller falls back to a download of all events. - router, connection, _ = _checkpoints_owner_with_version(CHECKPOINT_VERSION + 1) + router, cached_client, _ = _checkpoints_owner_with_version(CHECKPOINT_VERSION + 1) assert router.get_latest_checkpoint() is None assert "Failed to load checkpoint" in capsys.readouterr().out diff --git a/tests/unit/test_versioned_folder_lookup.py b/tests/unit/test_versioned_folder_lookup.py index b2a991f57fb..98f6241a1f0 100644 --- a/tests/unit/test_versioned_folder_lookup.py +++ b/tests/unit/test_versioned_folder_lookup.py @@ -7,7 +7,7 @@ folders) and _sorted_by_version (P2P lookup), each tested separately. """ -from syft.sync.connections.drive.gdrive_transport import ( +from syft.sync.connections.drive.drive_cached_client import ( _extract_version_from_name, _looks_like_version, _partition_by_version, diff --git a/tests/unit/utils.py b/tests/unit/utils.py index 1045405dfac..798fcc31f0f 100644 --- a/tests/unit/utils.py +++ b/tests/unit/utils.py @@ -1,13 +1,20 @@ -from syft.sync.events.file_change_event import FileChangeEventsMessage -from syft.sync.events.file_change_event import FileChangeEvent -from syft.sync.events.file_change_event import FileChangeEventsMessageFileName -from pathlib import Path -import uuid import random -from typing import List import time -from syft.sync.messages.proposed_filechange import ProposedFileChangesMessage -from syft.sync.messages.proposed_filechange import ProposedFileChange +import uuid +from pathlib import Path +from typing import List + +from syft.sync.connections.drive import gdrive_retry +from syft.sync.connections.drive.mock_drive_service import MockListRequest +from syft.sync.events.file_change_event import ( + FileChangeEvent, + FileChangeEventsMessage, + FileChangeEventsMessageFileName, +) +from syft.sync.messages.proposed_filechange import ( + ProposedFileChange, + ProposedFileChangesMessage, +) from syft.sync.utils.syftbox_utils import get_event_hash_from_content from syft.sync.connections.connection_router import ConnectionRouter from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service @@ -15,6 +22,22 @@ from syft.sync.syftbox_manager import SyftboxManager +def fail_next_listing(monkeypatch, name_prefix: str) -> None: + """Make the next mock Drive listing of files named ``name_prefix`` fail once, + with an error the Drive client retries.""" + monkeypatch.setattr(gdrive_retry.time, "sleep", lambda _: None) + execute = MockListRequest.execute + failed = [] + + def fail_once(request: MockListRequest) -> dict: + if not failed and f"name contains '{name_prefix}'" in request._q: + failed.append(request) + raise ConnectionResetError("Drive dropped the connection") + return execute(request) + + monkeypatch.setattr(MockListRequest, "execute", fail_once) + + def get_mock_event(path: str = "email@email.com/test.job") -> FileChangeEvent: email = path.split("/")[0] file_path = path.split("/")[-1] @@ -226,8 +249,8 @@ def router_with_mock_connection(email: str) -> ConnectionRouter: """A connection router over a mock Drive, for a part a test builds on its own. In a manager every part shares the router ``SyftboxManager.from_config`` builds.""" - connection, _ = pair_with_mock_service(email, "peer@example.com") + cached_client, _ = pair_with_mock_service(email, "peer@example.com") return ConnectionRouter( - transport=SyftboxManager._drive_transport(connection), + transport=SyftboxManager._drive_transport(cached_client), peer_store=PeerStore(email=email), )