From d23191073dc785402076daa730dba04827201220 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Tue, 27 Oct 2026 20:22:17 -0300 Subject: [PATCH 01/28] fix: retry the checkpoint and rolling-state Drive calls --- .../connections/drive/gdrive_transport.py | 54 +++++++++++-------- tests/unit/test_checkpoints.py | 25 ++++++++- tests/unit/test_rolling_state.py | 18 ++++++- tests/unit/utils.py | 39 +++++++++++--- 4 files changed, 102 insertions(+), 34 deletions(-) diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 25c90276081..abb9de393c6 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -1997,8 +1997,8 @@ def download_raw_latest_checkpoint(self) -> bytes | None: 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() + results = execute_with_retries( + self.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) @@ -2038,8 +2038,8 @@ def delete_all_checkpoints(self, exclude_file_id: str | None = None): 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() + results = execute_with_retries( + self.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) @@ -2050,7 +2050,9 @@ def delete_all_checkpoints(self, exclude_file_id: str | None = None): if item["id"] == exclude_file_id: continue # Skip the newly uploaded checkpoint try: - self.drive_service.files().delete(fileId=item["id"]).execute() + execute_with_retries( + self.drive_service.files().delete(fileId=item["id"]) + ) except Exception as e: print(f"Warning: Failed to delete checkpoint {item['name']}: {e}") @@ -2081,8 +2083,8 @@ def download_all_raw_incremental_checkpoints(self) -> list[bytes]: 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() + results = execute_with_retries( + self.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) @@ -2111,7 +2113,9 @@ def get_incremental_checkpoint_count(self) -> int: 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() + results = execute_with_retries( + self.drive_service.files().list(q=query, fields="files(id)") + ) return len(results.get("files", [])) def get_next_incremental_sequence_number(self) -> int: @@ -2124,8 +2128,8 @@ def get_next_incremental_sequence_number(self) -> int: 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() + results = execute_with_retries( + self.drive_service.files().list(q=query, fields="files(name)") ) items = results.get("files", []) @@ -2151,14 +2155,16 @@ def delete_all_incremental_checkpoints(self) -> None: 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() + results = execute_with_retries( + self.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) for item in items: try: - self.drive_service.files().delete(fileId=item["id"]).execute() + execute_with_retries( + self.drive_service.files().delete(fileId=item["id"]) + ) except Exception as e: print( f"Warning: Failed to delete incremental checkpoint {item['name']}: {e}" @@ -2279,10 +2285,12 @@ def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: # 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() + execute_with_retries( + self.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 @@ -2318,8 +2326,8 @@ def download_raw_rolling_state(self) -> bytes | None: 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() + results = execute_with_retries( + self.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) @@ -2360,15 +2368,17 @@ def delete_rolling_state(self) -> None: 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() + results = execute_with_retries( + self.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) # Delete each rolling state file for item in items: try: - self.drive_service.files().delete(fileId=item["id"]).execute() + execute_with_retries( + self.drive_service.files().delete(fileId=item["id"]) + ) except Exception as e: print(f"Warning: Failed to delete rolling state {item['name']}: {e}") diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index 6f54a66cf50..a6b917e09ff 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -12,15 +12,18 @@ import pytest -from syft.sync.sync import datasite_owner_syncer 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.mock_drive_service import 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 +743,21 @@ 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_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = do_conn.as_transport().owner_store + 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_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = do_conn.as_transport().owner_store + store.put_incremental("incremental_checkpoint_1_1.0.tar.gz", b"inc") + + fail_next_listing(monkeypatch, INCREMENTAL_CHECKPOINT_PREFIX) + assert store.incremental_count() == 1 diff --git a/tests/unit/test_rolling_state.py b/tests/unit/test_rolling_state.py index 1ef1dbce756..50a97d24460 100644 --- a/tests/unit/test_rolling_state.py +++ b/tests/unit/test_rolling_state.py @@ -7,10 +7,15 @@ """ import time + +from syft.sync.checkpoints.rolling_state import ( + ROLLING_STATE_FILENAME_PREFIX, + RollingState, +) +from syft.sync.connections.drive.mock_drive_service import 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 tests.unit.utils import fail_next_listing, get_mock_event def test_upload_and_get_rolling_state(): @@ -270,3 +275,12 @@ 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_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = do_conn.as_transport().owner_store + 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" diff --git a/tests/unit/utils.py b/tests/unit/utils.py index 2d7c1f3da8d..06aea921a2d 100644 --- a/tests/unit/utils.py +++ b/tests/unit/utils.py @@ -1,16 +1,39 @@ -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 +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] From 99003b5b540b86449eaea44d787bd4fba16c3fd5 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Tue, 27 Oct 2026 20:45:50 -0300 Subject: [PATCH 02/28] fix: raise when Drive is unreachable while reading the owner state --- syft/sync/connections/drive/gdrive_retry.py | 8 +++++++- syft/sync/connections/drive/gdrive_transport.py | 14 ++++++++++++++ tests/unit/test_rolling_state.py | 17 +++++++++++++++++ 3 files changed, 38 insertions(+), 1 deletion(-) diff --git a/syft/sync/connections/drive/gdrive_retry.py b/syft/sync/connections/drive/gdrive_retry.py index e861994eca3..95658d77d9b 100644 --- a/syft/sync/connections/drive/gdrive_retry.py +++ b/syft/sync/connections/drive/gdrive_retry.py @@ -72,6 +72,12 @@ def is_retryable_error(error: Exception) -> bool: return False +def is_unreachable(error: Exception) -> bool: + """True when ``error`` means Drive cannot be reached: an error in + UNREACHABLE_ERRORS, or a transient error that outlasted the retries.""" + return isinstance(error, UNREACHABLE_ERRORS) or is_retryable_error(error) + + @contextmanager def raise_unavailable(message: str) -> Iterator[None]: """Raise TransportUnavailableError(message) when Drive cannot be reached. @@ -83,7 +89,7 @@ def raise_unavailable(message: str) -> Iterator[None]: try: yield except (HttpError, *RETRYABLE_TRANSPORT_ERRORS, *UNREACHABLE_ERRORS) as e: - if isinstance(e, UNREACHABLE_ERRORS) or is_retryable_error(e): + if is_unreachable(e): raise TransportUnavailableError(message) from e raise diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index abb9de393c6..f51ca53fab5 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -34,6 +34,7 @@ from syft.sync.connections.drive.gdrive_retry import ( batch_execute_with_retries, execute_with_retries, + is_unreachable, next_chunk_with_retries, raise_unavailable, ) @@ -50,6 +51,7 @@ MessageFileName, ProposedFileChangesMessage, ) +from syft.sync.transport.exceptions import TransportUnavailableError from syft.sync.transport.naming import ( # noqa: F401 (re-exported) PEERS_META_KEY, SYFT_PEERS_VERSION, @@ -2019,6 +2021,10 @@ def download_raw_latest_checkpoint(self) -> bytes | None: try: return self.download_file(latest_file["id"]) except Exception as e: + if is_unreachable(e): + raise TransportUnavailableError( + "Could not read the latest checkpoint from Drive." + ) from e print(f"Warning: Failed to load checkpoint: {e}") return None @@ -2097,6 +2103,10 @@ def download_all_raw_incremental_checkpoints(self) -> list[bytes]: file_data = self.download_file(item["id"]) result.append(file_data) except Exception as e: + if is_unreachable(e): + raise TransportUnavailableError( + "Could not read the incremental checkpoints from Drive." + ) from e print( f"Warning: Failed to load incremental checkpoint {item['name']}: {e}" ) @@ -2350,6 +2360,10 @@ def download_raw_rolling_state(self) -> bytes | None: try: return self.download_file(latest_file["id"]) except Exception as e: + if is_unreachable(e): + raise TransportUnavailableError( + "Could not read the rolling state from Drive." + ) from e print(f"Warning: Failed to load rolling state: {e}") self._rolling_state_file_id = None return None diff --git a/tests/unit/test_rolling_state.py b/tests/unit/test_rolling_state.py index 50a97d24460..3b9c8cbe85f 100644 --- a/tests/unit/test_rolling_state.py +++ b/tests/unit/test_rolling_state.py @@ -8,13 +8,17 @@ import time +import pytest + from syft.sync.checkpoints.rolling_state import ( ROLLING_STATE_FILENAME_PREFIX, RollingState, ) +from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service from syft.sync.syftbox_manager import SyftboxManager from syft.sync.sync.constants import CACHE_DIR, ROLLING_STATE_FILENAME +from syft.sync.transport import TransportUnavailableError from tests.unit.utils import fail_next_listing, get_mock_event @@ -284,3 +288,16 @@ def test_rolling_state_read_retries_drive_error(monkeypatch): 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_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = do_conn.as_transport().owner_store + 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(GDriveConnection, "download_file", drive_down) + with pytest.raises(TransportUnavailableError): + store.get_rolling_state() From 1b236335e50ecb5a7ad01be402af8200e6b2dc35 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Tue, 27 Oct 2026 21:04:45 -0300 Subject: [PATCH 03/28] fix: read every page of the peer request and collection listings --- .../connections/drive/gdrive_transport.py | 66 ++++++++++--------- tests/unit/test_gdrive_pagination.py | 45 +++++++++++++ 2 files changed, 81 insertions(+), 30 deletions(-) diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index f51ca53fab5..8f3537d56f2 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -639,19 +639,17 @@ def get_peer_requests(self) -> list[str]: 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)", - ) + folders = self._list_all_files( + 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="id,name,owners", ) all_folder_peers = set() - for f in results.get("files", []): + for f in folders: try: folder = GdriveP2PFolder.from_name(f["name"]) except ValueError: @@ -673,9 +671,11 @@ def get_peer_requests(self) -> list[str]: 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"]: + undecided = peers_data.get(peer_email, {}).get("state") not in [ + "accepted", + "rejected", + ] + if undecided: pending_peers.append(peer_email) return pending_peers @@ -1404,6 +1404,26 @@ def _find_folder_by_name( 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, *, @@ -1675,12 +1695,8 @@ def owner_list_collections(self, prefix: str) -> list[str]: 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", []): + for folder in self._list_all_files(query, fields="name"): try: cf = CollectionFolder.from_name(prefix, folder["name"]) except ValueError: @@ -1699,14 +1715,8 @@ def owner_list_all_collections_with_permissions( 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", []): + for folder in self._list_all_files(query, fields="id,name,appProperties"): try: cf = CollectionFolder.from_name(prefix, folder["name"]) except ValueError: @@ -1749,12 +1759,8 @@ def watcher_list_collections(self, prefix: str) -> list[dict]: 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", []): + for folder in self._list_all_files(query, fields="name, owners"): try: cf = CollectionFolder.from_name(prefix, folder["name"]) except ValueError: diff --git a/tests/unit/test_gdrive_pagination.py b/tests/unit/test_gdrive_pagination.py index f2550dc07b6..c5c7ebb933f 100644 --- a/tests/unit/test_gdrive_pagination.py +++ b/tests/unit/test_gdrive_pagination.py @@ -1,6 +1,11 @@ from unittest.mock import Mock, patch from uuid import uuid4 + from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.mock_drive_service import ( + MockDriveBackingStore, + MockDriveService, +) from syft.sync.events.file_change_event import ( FileChangeEventsMessage, FileChangeEventsMessageFileName, @@ -252,3 +257,43 @@ def test_newer_message_listed_after_older_one_in_same_second_not_skipped(): ) 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_connections(*emails: str) -> list[GDriveConnection]: + store = MockDriveBackingStore() + return [ + GDriveConnection.from_service(e, MockDriveService(store, e)) for e in emails + ] + + +def test_collection_listings_read_every_page(): + do_conn, ds_conn = _mock_connections("do@example.com", "ds@example.com") + owner, reader = do_conn.as_transport(), ds_conn.as_transport() + tags = [f"tag{i}" for i in range(FILE_COUNT)] + for tag in tags: + owner.collections.create( + "syft_datasetcollection", tag, "hash", "do@example.com" + ) + owner.collections.grant_read( + "syft_datasetcollection", tag, "hash", ["ds@example.com"] + ) + + assert sorted(owner.collections.list_mine("syft_datasetcollection")) == sorted(tags) + mine = owner.collections.list_mine_with_permissions("syft_datasetcollection") + assert len(mine) == FILE_COUNT + shared = reader.collections.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_conn, *peers = _mock_connections("do@example.com", *peer_emails) + for peer in peers: + peer.as_transport().peering.open_channels("do@example.com") + + requests = do_conn.as_transport().peering.incoming_requests() + assert sorted(requests) == sorted(peer_emails) From 1327132689d2ba472728ad9808cef12699dea17b Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Tue, 27 Oct 2026 21:51:05 -0300 Subject: [PATCH 04/28] refactor: drop the naming re-export and timestamp wrapper --- .../connections/drive/gdrive_transport.py | 11 +++------- tests/unit/test_gdrive_pagination.py | 22 ------------------- tests/unit/test_peer_request_owner.py | 2 +- tests/unit/test_peers_json_version.py | 7 ++---- 4 files changed, 6 insertions(+), 36 deletions(-) diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 8f3537d56f2..e087cf2858e 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -52,7 +52,7 @@ ProposedFileChangesMessage, ) from syft.sync.transport.exceptions import TransportUnavailableError -from syft.sync.transport.naming import ( # noqa: F401 (re-exported) +from syft.sync.transport.naming import ( PEERS_META_KEY, SYFT_PEERS_VERSION, timestamp_from_message_name, @@ -980,11 +980,6 @@ def owner_get_archive_folder_id(self, sender_email: str) -> str: else: return self.create_archive_folder(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, @@ -1026,7 +1021,7 @@ def get_file_metadatas_from_folder( should_stop = False for file in page_files: - timestamp = self._extract_timestamp_from_filename(file["name"]) + timestamp = timestamp_from_message_name(file["name"]) if timestamp is not None: if timestamp > since_timestamp: @@ -2216,7 +2211,7 @@ def get_events_count_since_checkpoint( for f in file_metadatas: if not f["name"].startswith("syfteventsmessagev3_"): continue - event_timestamp = self._extract_timestamp_from_filename(f["name"]) + event_timestamp = timestamp_from_message_name(f["name"]) if event_timestamp is not None and event_timestamp > checkpoint_timestamp: count += 1 return count diff --git a/tests/unit/test_gdrive_pagination.py b/tests/unit/test_gdrive_pagination.py index c5c7ebb933f..90dde76832e 100644 --- a/tests/unit/test_gdrive_pagination.py +++ b/tests/unit/test_gdrive_pagination.py @@ -122,28 +122,6 @@ def test_orderby_name_desc_applied(): 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() diff --git a/tests/unit/test_peer_request_owner.py b/tests/unit/test_peer_request_owner.py index d93d74e5092..80a25aadf8d 100644 --- a/tests/unit/test_peer_request_owner.py +++ b/tests/unit/test_peer_request_owner.py @@ -8,7 +8,6 @@ from syft.sync.connections.drive.gdrive_transport import ( GDRIVE_P2P_FOLDER_DATASITE_PREFIX, - PEERS_META_KEY, GdriveP2PFolder, ) from syft.sync.connections.drive.mock_drive_service import ( @@ -17,6 +16,7 @@ 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 diff --git a/tests/unit/test_peers_json_version.py b/tests/unit/test_peers_json_version.py index 17178da8c18..15a580eadfe 100644 --- a/tests/unit/test_peers_json_version.py +++ b/tests/unit/test_peers_json_version.py @@ -14,12 +14,9 @@ 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_transport import GDriveConnection from syft.sync.peers.peer import PeerState +from syft.sync.transport.naming import PEERS_META_KEY, SYFT_PEERS_VERSION PEER = "bob@example.com" From f5f734deea12d25e46e4eb1203d89d8db89d70c3 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Tue, 27 Oct 2026 22:09:57 -0300 Subject: [PATCH 05/28] refactor: route the syncers' Drive calls through the router --- syft/sync/connections/connection_router.py | 19 ++++++++++++++++ syft/sync/sync/datasite_owner_syncer.py | 26 +++++++++------------- syft/sync/sync/datasite_watcher_syncer.py | 7 +++--- 3 files changed, 33 insertions(+), 19 deletions(-) diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index b1e0e3e119d..3a3b8f8c03b 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -239,6 +239,25 @@ def connection_for_parallel_download(self) -> SyftboxPlatformConnection: """Create a new connection for thread-safe parallel downloads.""" return self.copy_connection(self.connection_for_datasite_watcher()) + def transport_for_thread(self) -> Transport: + """The transport jobs over a new copy of the connection, for one thread.""" + return self.connection_for_parallel_download().as_transport() + + def download_outbox_message(self, file_id: str) -> bytes: + """The bytes of one outbox message file (thread-safe).""" + return self.connection_for_parallel_download().download_file(file_id) + + def owner_list_outbox_files(self, recipient_email: str) -> List[dict] | None: + """The files in my outbox for ``recipient_email``, or None when that + outbox does not exist.""" + connection = self.connection_for_outbox() + folder_id = connection._get_own_datasite_outbox_id(recipient_email) + if folder_id is None: + return None + return connection.get_file_metadatas_from_folder( + folder_id, since_timestamp=None + ) + # ========================================================================= # EVENT LOG (self-encrypted at rest when encryption enabled) # ========================================================================= diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 8e01312751b..19fd3ba1a36 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -633,18 +633,18 @@ def _get_file_metadatas_with_new_connection( self, collection: FileCollection, spec: CollectionSyncSpec ) -> list: """Get file metadatas for a collection using a new connection for thread safety.""" - connection = self.connection_router.connection_for_parallel_download() - return connection.watcher_get_collection_file_metadatas( + transport = self.connection_router.transport_for_thread() + return transport.collections.list_files( spec.wire_prefix(collection.variant), - tag=collection.tag, - content_hash=collection.content_hash, - owner_email=self.email, + collection.tag, + collection.content_hash, + self.email, ) def _download_file_with_new_connection(self, file_id: str) -> bytes: """Download a file using a new connection for thread safety.""" - connection = self.connection_router.connection_for_parallel_download() - return connection.watcher_download_collection_file(file_id) + transport = self.connection_router.transport_for_thread() + return transport.collections.fetch_file(file_id) def _create_current_filestate_event(self, path: str) -> FileChangeEvent: """An event that sets ``path`` on a peer to its current local state. @@ -941,8 +941,7 @@ def _download_outbox_message_with_new_connection( self, file_id: str ) -> FileChangeEventsMessage: """Download and decompress one outbox events-message file.""" - connection = self.connection_router.connection_for_parallel_download() - raw = connection.download_file(file_id) + raw = self.connection_router.download_outbox_message(file_id) return FileChangeEventsMessage.from_compressed_data(raw) def compact_outbox_if_needed( @@ -968,17 +967,14 @@ def compact_outbox_if_needed( ) return 0 - conn = self.connection_router.connection_for_outbox() - folder_id = conn._get_own_datasite_outbox_id(recipient_email) - if folder_id is None: + metas = self.connection_router.owner_list_outbox_files(recipient_email) + if metas is None: logger.info( f"WARNING: outbox compaction for {recipient_email} skipped — " f"could not resolve outbox folder id" ) return 0 - metas = conn.get_file_metadatas_from_folder(folder_id, since_timestamp=None) - # Keep only well-formed events-message filenames; track timestamps so # we apply events in order in the merged message. parsed: list[tuple[float, str]] = [] @@ -1018,7 +1014,7 @@ def compact_outbox_if_needed( self.connection_router.owner_write_event_messages_to_outbox( recipient_email, FileChangeEventsMessage(events=batch) ) - conn.delete_multiple_files_by_ids(file_ids) + self.connection_router.delete_multiple_files_by_ids(file_ids) return len(parsed) def write_file_filesystem(self, path: str, content: str): diff --git a/syft/sync/sync/datasite_watcher_syncer.py b/syft/sync/sync/datasite_watcher_syncer.py index fecf0aeddd3..1cca9147d56 100644 --- a/syft/sync/sync/datasite_watcher_syncer.py +++ b/syft/sync/sync/datasite_watcher_syncer.py @@ -178,16 +178,15 @@ def download_events_message_with_new_connection( Returns None when the message is skipped (see ``ConnectionRouter.watcher_read_events_message``). """ - connection = self.connection_router.connection_for_parallel_download() - raw = connection.download_file(file_id) + raw = self.connection_router.download_outbox_message(file_id) return self.connection_router.watcher_read_events_message(peer_email, raw) def download_collection_file_with_new_connection( self, file_id: str, owner_email: str ) -> bytes: """Download collection file using a new connection (thread-safe).""" - connection = self.connection_router.connection_for_parallel_download() - data = connection.watcher_download_collection_file(file_id) + transport = self.connection_router.transport_for_thread() + data = transport.collections.fetch_file(file_id) data = self.connection_router.peer_store.decrypt_dataset_if_needed( owner_email, data ) From 2377c713ff7f406af85fccbf10b0e2abaccfe705 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 09:44:45 -0300 Subject: [PATCH 06/28] refactor: hold the Drive service in GDriveClient --- .../tests/test_dataset_upload_private.py | 4 +- packages/syft-rds/tests/test_sync_manager.py | 16 +-- .../tests/test_version_mismatch_flow.py | 6 +- syft/sync/connections/drive/gdrive_client.py | 14 +++ .../connections/drive/gdrive_transport.py | 114 +++++++++--------- syft/sync/syftbox_manager.py | 6 +- tests/unit/test_dataset_collection_listing.py | 5 +- tests/unit/test_delete_syftbox.py | 10 +- tests/unit/test_drive_connections.py | 2 +- tests/unit/test_encryption.py | 2 +- tests/unit/test_gdrive_pagination.py | 31 ++--- tests/unit/test_login_email_check.py | 2 +- tests/unit/test_p2p_folder_lookup.py | 3 +- tests/unit/test_peer_request_owner.py | 2 +- tests/unit/test_peers_json_version.py | 3 +- tests/unit/test_sender_binding.py | 2 +- tests/unit/test_sync_manager.py | 2 +- tests/unit/test_version_negotiation.py | 2 +- tests/unit/test_versioned_folder_adopt.py | 7 +- 19 files changed, 131 insertions(+), 102 deletions(-) create mode 100644 syft/sync/connections/drive/gdrive_client.py diff --git a/packages/syft-rds/tests/test_dataset_upload_private.py b/packages/syft-rds/tests/test_dataset_upload_private.py index ef90e08f0d1..61b3fb717dc 100644 --- a/packages/syft-rds/tests/test_dataset_upload_private.py +++ b/packages/syft-rds/tests/test_dataset_upload_private.py @@ -149,7 +149,7 @@ def test_ds_cannot_find_private_folders_via_gdrive_query(self): ds_manager.sync_engine._connection_router.connections[0] ) results = ( - ds_connection.drive_service.files() + ds_connection.client.drive_service.files() .list( q=f"{collection_name_query(PRIVATE_DATASET_COLLECTION_PREFIX)} and trashed=false", fields="files(id,name)", @@ -163,7 +163,7 @@ def test_ds_cannot_find_private_folders_via_gdrive_query(self): do_manager.sync_engine._connection_router.connections[0] ) do_results = ( - do_connection.drive_service.files() + do_connection.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 d375856d18e..c5bc03fef0f 100644 --- a/packages/syft-rds/tests/test_sync_manager.py +++ b/packages/syft-rds/tests/test_sync_manager.py @@ -233,7 +233,7 @@ def test_dataset_only_mock_data_uploaded(): files = ds_manager.sync_engine._connection_router.connections[ 0 - ].drive_service._backing_store.files + ].client.drive_service._backing_store.files list(files.keys()) file_objs = list(files.values()) file_objs_ex_dataset_yaml = [ @@ -1038,7 +1038,7 @@ def test_ds_dataset_cache_aware_sync(): # Get the mock backing store and directories for creating second pair mock_backing_store = ds_manager.sync_engine._connection_router.connections[ 0 - ].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 @@ -1057,10 +1057,10 @@ def test_ds_dataset_cache_aware_sync(): # Replace mock backing store to share dataset collections ds_manager2.sync_engine._connection_router.connections[ 0 - ].drive_service._backing_store = mock_backing_store + ].client.drive_service._backing_store = mock_backing_store do_manager2.sync_engine._connection_router.connections[ 0 - ].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() @@ -1138,7 +1138,7 @@ def test_do_dataset_cache_aware_sync(): # Get the mock backing store and directories for creating second pair mock_backing_store = ds_manager.sync_engine._connection_router.connections[ 0 - ].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 @@ -1157,10 +1157,10 @@ def test_do_dataset_cache_aware_sync(): # Replace mock backing store to share dataset collections ds_manager2.sync_engine._connection_router.connections[ 0 - ].drive_service._backing_store = mock_backing_store + ].client.drive_service._backing_store = mock_backing_store do_manager2.sync_engine._connection_router.connections[ 0 - ].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() @@ -1205,7 +1205,7 @@ def test_in_memory_connection_load_state(): # Get the backing store that will persist across "restarts" backing_store = do_manager1.sync_engine._connection_router.connections[ 0 - ].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 diff --git a/packages/syft-rds/tests/test_version_mismatch_flow.py b/packages/syft-rds/tests/test_version_mismatch_flow.py index 5ebb44ae328..8b90b94bca5 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -33,7 +33,9 @@ def _find_p2p_folders(connection, peer_email): " and trashed=false" ) results = ( - connection.drive_service.files().list(q=q, fields="files(id, name)").execute() + connection.client.drive_service.files() + .list(q=q, fields="files(id, name)") + .execute() ) return results.get("files", []) @@ -47,7 +49,7 @@ def _find_versioned_p2p_folders(connection, peer_email, version): def _get_backing_store(manager): """Extract the mock backing store from a manager's connection.""" conn = manager.peer_manager.connection_router.connections[0] - return conn.drive_service._backing_store + return conn.client.drive_service._backing_store def _reinitialize_manager( diff --git a/syft/sync/connections/drive/gdrive_client.py b/syft/sync/connections/drive/gdrive_client.py new file mode 100644 index 00000000000..763c7008728 --- /dev/null +++ b/syft/sync/connections/drive/gdrive_client.py @@ -0,0 +1,14 @@ +"""The Google Drive API client that the Drive transport jobs share. + +One client per connection. It holds the Drive service, which is not safe to +use from more than one thread. +""" + +from dataclasses import dataclass +from typing import Any + + +@dataclass +class GDriveClient: + # A googleapiclient Drive v3 service, or a MockDriveService in tests. + drive_service: Any diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index e087cf2858e..d1f0af4d096 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -31,6 +31,7 @@ SyftboxPlatformConnection, ) from syft.sync.connections.drive.drive_connections import drive_transport +from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_retry import ( batch_execute_with_retries, execute_with_retries, @@ -330,7 +331,7 @@ class GDriveConnection(SyftboxPlatformConnection): class Config: arbitrary_types_allowed = True - drive_service: Any = None + client: GDriveClient | None = None credentials: GoogleCredentials | None = None verbose: bool = True email: str @@ -433,12 +434,11 @@ def setup( 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( + if drive_service is None: + drive_service = build_drive_service( self.credentials, environment=self.environment ) + self.client = GDriveClient(drive_service) if warm_syftbox_folder_id_cache: self.get_personal_syftbox_folder_id() @@ -458,8 +458,10 @@ def copy(self) -> "GDriveConnection": MockDriveService, ) - if isinstance(self.drive_service, MockDriveService): - new_conn = GDriveConnection.from_service(self.email, self.drive_service) + if isinstance(self.client.drive_service, MockDriveService): + new_conn = GDriveConnection.from_service( + self.email, self.client.drive_service + ) else: new_conn = GDriveConnection.from_token_path( self.email, self.token_path, warm_syftbox_folder_id_cache=False @@ -493,9 +495,9 @@ def environment(self) -> Environment: return check_env() def get_authenticated_email(self) -> str: - """Return the email of the Google account behind drive_service.""" + """Return the email of the Google account behind the Drive service.""" about = execute_with_retries( - self.drive_service.about().get(fields="user(emailAddress)") + self.client.drive_service.about().get(fields="user(emailAddress)") ) return about["user"]["emailAddress"] @@ -548,7 +550,7 @@ def _get_peers_file_id(self) -> str | None: 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)") + self.client.drive_service.files().list(q=query, fields="files(id)") ) items = results.get("files", []) return items[0]["id"] if items else None @@ -607,7 +609,7 @@ def _write_peers_json(self, peers_data: dict[str, dict[str, Any]]): else: # Update existing file execute_with_retries( - self.drive_service.files().update( + self.client.drive_service.files().update( fileId=file_id, media_body=file_payload ) ) @@ -821,7 +823,7 @@ def owner_write_raw_bytes_to_outbox( f"Moving file to correct folder..." ) execute_with_retries( - self.drive_service.files().update( + self.client.drive_service.files().update( fileId=file_id, addParents=outbox_folder_id, fields="id, parents", @@ -861,12 +863,12 @@ def owner_remove_proposed_filechange_message_from_inbox( 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.""" file_info = execute_with_retries( - self.drive_service.files().get(fileId=file_id, fields="parents") + self.client.drive_service.files().get(fileId=file_id, fields="parents") ) previous_parents = ",".join(file_info.get("parents", [])) archive_folder_id = self.owner_get_archive_folder_id(sender_email) execute_with_retries( - self.drive_service.files().update( + self.client.drive_service.files().update( fileId=file_id, addParents=archive_folder_id, removeParents=previous_parents, @@ -878,7 +880,7 @@ def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: 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( + self.client.drive_service.permissions().list( fileId=file_id, fields="permissions(emailAddress)" ) ) @@ -899,7 +901,7 @@ def add_permission(self, file_id: str, recipient: str, write=False): "emailAddress": recipient, } execute_with_retries( - self.drive_service.permissions().create( + self.client.drive_service.permissions().create( fileId=file_id, body=permission, sendNotificationEmail=True ) ) @@ -964,7 +966,7 @@ def get_archive_folder_id_from_drive(self, sender_email: str) -> str | None: 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)") + self.client.drive_service.files().list(q=query, fields="files(id)") ) items = results.get("files", []) return items[0]["id"] if items else None @@ -1005,7 +1007,7 @@ def get_file_metadatas_from_folder( while True: results = execute_with_retries( - self.drive_service.files().list( + self.client.drive_service.files().list( q=query, fields="files(id, name, size, mimeType, modifiedTime), nextPageToken", pageSize=page_size, @@ -1246,7 +1248,7 @@ def watcher_send_raw_bytes_to_inbox( f"Moving file to correct folder..." ) execute_with_retries( - self.drive_service.files().update( + self.client.drive_service.files().update( fileId=file_id, addParents=inbox_id, fields="id, parents", @@ -1280,7 +1282,7 @@ def reset_caches(self): 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 + self.client.drive_service, syftbox_folder_id ) def delete_multiple_files_by_ids( @@ -1318,14 +1320,16 @@ def callback(request_id, response, exception): BATCH_SIZE = 100 for i in range(0, len(file_ids), BATCH_SIZE): chunk = file_ids[i : i + BATCH_SIZE] - batch = self.drive_service.new_batch_http_request(callback=callback) + batch = self.client.drive_service.new_batch_http_request(callback=callback) for file_id in chunk: - batch.add(self.drive_service.files().delete(fileId=file_id)) + batch.add(self.client.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)) + execute_with_retries( + self.client.drive_service.files().delete(fileId=file_id) + ) except Exception as e: if raise_on_error: raise e @@ -1349,7 +1353,7 @@ def find_orphaned_syft_files(self) -> list[str]: page_token = None while True: results = execute_with_retries( - self.drive_service.files().list( + self.client.drive_service.files().list( q=query, fields="files(id), nextPageToken", pageToken=page_token ) ) @@ -1394,7 +1398,9 @@ def _find_folder_by_name( 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) + self.client.drive_service.files().list( + q=query, fields="files(id)", pageSize=1 + ) ) items = results.get("files", []) return items[0]["id"] if items else None @@ -1408,7 +1414,7 @@ def _list_all_files(self, query: str, fields: str) -> list[dict]: page_token = None while True: page = execute_with_retries( - self.drive_service.files().list( + self.client.drive_service.files().list( q=query, fields=f"files({fields}), nextPageToken", pageToken=page_token, @@ -1446,7 +1452,7 @@ def _find_folders( page_token = None while True: page = execute_with_retries( - self.drive_service.files().list( + self.client.drive_service.files().list( q=query, fields="files(id, name), nextPageToken", pageToken=page_token, @@ -1511,7 +1517,7 @@ def _find_or_adopt_versioned_folder( folder_id, name = older[-1] execute_with_retries( - self.drive_service.files().update( + self.client.drive_service.files().update( fileId=folder_id, body={"name": current_name} ) ) @@ -1525,7 +1531,7 @@ def _find_or_adopt_versioned_folder( return folder_id def download_file(self, file_id: str) -> bytes: - request = self.drive_service.files().get_media(fileId=file_id) + request = self.client.drive_service.files().get_media(fileId=file_id) file_buffer = io.BytesIO() downloader = MediaIoBaseDownload( @@ -1557,7 +1563,7 @@ def _create_syft_file( ) -> dict: """Create a file or folder with the syft tag.""" return execute_with_retries( - self.drive_service.files().create( + self.client.drive_service.files().create( body={ **body, "properties": {SYFT_FILE_PROPERTY_KEY: SYFT_FILE_PROPERTY_VALUE}, @@ -1570,7 +1576,7 @@ def _create_syft_file( 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)") + self.client.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) return items[0]["id"] if items else None @@ -1581,7 +1587,7 @@ def get_inbox_proposed_event_id_from_name( inbox_folder_id = self._get_own_datasite_inbox_id(sender_email) query = f"name='{name}' and '{inbox_folder_id}' in parents and trashed=false" results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id, name)") + self.client.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) return items[0]["id"] if items else None @@ -1625,7 +1631,7 @@ def owner_tag_collection_as_any( """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( + self.client.drive_service.files().update( fileId=folder_id, body={"appProperties": {"syft_shared_with_any": "true"}}, ) @@ -1651,7 +1657,7 @@ def callback(request_id, response, exception): BATCH_SIZE = 100 for i in range(0, len(users), BATCH_SIZE): chunk = users[i : i + BATCH_SIZE] - batch = self.drive_service.new_batch_http_request(callback=callback) + batch = self.client.drive_service.new_batch_http_request(callback=callback) for user_email in chunk: permission = { "type": "user", @@ -1659,7 +1665,7 @@ def callback(request_id, response, exception): "emailAddress": user_email, } batch.add( - self.drive_service.permissions().create( + self.client.drive_service.permissions().create( fileId=file_id, body=permission, sendNotificationEmail=True, @@ -1839,7 +1845,7 @@ def _get_own_peer_info_file_id(self) -> Optional[str]: 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)") + self.client.drive_service.files().list(q=query, fields="files(id)") ) items = results.get("files", []) return items[0]["id"] if items else None @@ -1864,7 +1870,7 @@ def write_peer_info_file(self, peer_info: "PeerInfo") -> None: else: # Update existing file execute_with_retries( - self.drive_service.files().update( + self.client.drive_service.files().update( fileId=file_id, media_body=file_payload ) ) @@ -1875,7 +1881,7 @@ def _get_peer_info_file_id_of(self, peer_email: str) -> Optional[str]: 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)") + self.client.drive_service.files().list(q=query, fields="files(id)") ) items = results.get("files", []) return items[0]["id"] if items else None @@ -2001,7 +2007,7 @@ def download_raw_latest_checkpoint(self) -> bytes | None: f"and name contains '{CHECKPOINT_FILENAME_PREFIX}'" ) results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id, name)") + self.client.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) @@ -2046,7 +2052,7 @@ def delete_all_checkpoints(self, exclude_file_id: str | None = None): f"and name contains '{CHECKPOINT_FILENAME_PREFIX}'" ) results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id, name)") + self.client.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) @@ -2058,7 +2064,7 @@ def delete_all_checkpoints(self, exclude_file_id: str | None = None): continue # Skip the newly uploaded checkpoint try: execute_with_retries( - self.drive_service.files().delete(fileId=item["id"]) + self.client.drive_service.files().delete(fileId=item["id"]) ) except Exception as e: print(f"Warning: Failed to delete checkpoint {item['name']}: {e}") @@ -2091,7 +2097,7 @@ def download_all_raw_incremental_checkpoints(self) -> list[bytes]: f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" ) results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id, name)") + self.client.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) @@ -2125,7 +2131,7 @@ def get_incremental_checkpoint_count(self) -> int: f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" ) results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id)") + self.client.drive_service.files().list(q=query, fields="files(id)") ) return len(results.get("files", [])) @@ -2140,7 +2146,7 @@ def get_next_incremental_sequence_number(self) -> int: f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" ) results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(name)") + self.client.drive_service.files().list(q=query, fields="files(name)") ) items = results.get("files", []) @@ -2167,14 +2173,14 @@ def delete_all_incremental_checkpoints(self) -> None: f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" ) results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id, name)") + self.client.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) for item in items: try: execute_with_retries( - self.drive_service.files().delete(fileId=item["id"]) + self.client.drive_service.files().delete(fileId=item["id"]) ) except Exception as e: print( @@ -2297,7 +2303,7 @@ def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: if self._rolling_state_file_id is not None: try: execute_with_retries( - self.drive_service.files().update( + self.client.drive_service.files().update( fileId=self._rolling_state_file_id, media_body=payload, ) @@ -2338,7 +2344,7 @@ def download_raw_rolling_state(self) -> bytes | None: f"and name contains '{ROLLING_STATE_FILENAME_PREFIX}'" ) results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id, name)") + self.client.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) @@ -2384,7 +2390,7 @@ def delete_rolling_state(self) -> None: f"and name contains '{ROLLING_STATE_FILENAME_PREFIX}'" ) results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id, name)") + self.client.drive_service.files().list(q=query, fields="files(id, name)") ) items = results.get("files", []) @@ -2392,7 +2398,7 @@ def delete_rolling_state(self) -> None: for item in items: try: execute_with_retries( - self.drive_service.files().delete(fileId=item["id"]) + self.client.drive_service.files().delete(fileId=item["id"]) ) except Exception as e: print(f"Warning: Failed to delete rolling state {item['name']}: {e}") @@ -2434,13 +2440,13 @@ def write_encryption_bundle(self, peer_email: str, bundle_json: str) -> None: # 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)") + self.client.drive_service.files().list(q=query, fields="files(id)") ) items = results.get("files", []) if items: execute_with_retries( - self.drive_service.files().update( + self.client.drive_service.files().update( fileId=items[0]["id"], media_body=file_payload ) ) @@ -2464,7 +2470,7 @@ def read_own_encryption_bundle(self, peer_email: str) -> str | 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)") + self.client.drive_service.files().list(q=query, fields="files(id)") ) items = results.get("files", []) if not items: @@ -2480,7 +2486,7 @@ def read_peer_encryption_bundle(self, peer_email: str) -> str | None: 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)") + self.client.drive_service.files().list(q=query, fields="files(id)") ) items = results.get("files", []) if not items: diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index 5664e35f9df..c3c61507d80 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -1081,7 +1081,7 @@ def _add_connection(self, connection: SyftboxPlatformConnection): if not ( isinstance(connection, GDriveConnection) and isinstance( - connection.drive_service, mock_drive_service.MockDriveService + connection.client.drive_service, mock_drive_service.MockDriveService ) ): raise ValueError( @@ -1343,11 +1343,11 @@ def _copy(self): if not isinstance(self._connection_router.connections[0], GDriveConnection): raise ValueError("Only GDriveConnections can be copied") if isinstance( - self._connection_router.connections[0].drive_service, + self._connection_router.connections[0].client.drive_service, mock_drive_service.MockDriveService, ): # Create new connection pointing to the same backing store - drive_service = self._connection_router.connections[0].drive_service + drive_service = self._connection_router.connections[0].client.drive_service new_do_connection = GDriveConnection.from_service(self.email, drive_service) new_manager._add_connection(new_do_connection) return new_manager diff --git a/tests/unit/test_dataset_collection_listing.py b/tests/unit/test_dataset_collection_listing.py index 5c74a2448f5..ce4f532bd77 100644 --- a/tests/unit/test_dataset_collection_listing.py +++ b/tests/unit/test_dataset_collection_listing.py @@ -9,6 +9,7 @@ import pytest +from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX @@ -18,9 +19,9 @@ def _conn(files): conn = GDriveConnection(email="alice@example.com", verbose=False) - conn.drive_service = Mock() + conn.client = GDriveClient(Mock()) conn._syftbox_folder_id = "syftbox-id" - conn.drive_service.files().list().execute.return_value = {"files": files} + conn.client.drive_service.files().list().execute.return_value = {"files": files} return conn diff --git a/tests/unit/test_delete_syftbox.py b/tests/unit/test_delete_syftbox.py index 35aefb338ed..324d4025074 100644 --- a/tests/unit/test_delete_syftbox.py +++ b/tests/unit/test_delete_syftbox.py @@ -131,7 +131,9 @@ def _query_files(connection, 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() + connection.client.drive_service.files() + .list(q=q, fields="files(id, name)") + .execute() ) return results.get("files", []) @@ -171,19 +173,19 @@ 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( + do_conn.client.drive_service.files().delete( fileId=do_conn.get_syftbox_folder_id() ).execute() found = _connection(ds_manager).find_orphaned_syft_files() - assert archive_id not in do_conn.drive_service._backing_store.files + assert archive_id not in do_conn.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 = _connection(ds_manager).client.drive_service user_file = drive.files().create(body={"name": "checkpoint_final.pt"}).execute() found = _connection(ds_manager).find_orphaned_syft_files() diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index b4c8214d280..2473de6933a 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -193,7 +193,7 @@ def test_owner_store_files_carry_the_syft_tag(pair): do_files = [ f - for f in do_conn.drive_service._backing_store.files.values() + for f in do_conn.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"}] diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index 51567c8ef2a..cb2488a6155 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -628,7 +628,7 @@ def test_tampered_inbox_message_logs_error(caplog): assert len(files) > 0 file_id = files[0]["id"] - mock_file = do_conn.drive_service._backing_store.files[file_id] + mock_file = do_conn.client.drive_service._backing_store.files[file_id] tampered = bytearray(mock_file.content) tampered[-1] ^= 0xFF mock_file.content = bytes(tampered) diff --git a/tests/unit/test_gdrive_pagination.py b/tests/unit/test_gdrive_pagination.py index 90dde76832e..9ff94ebbefd 100644 --- a/tests/unit/test_gdrive_pagination.py +++ b/tests/unit/test_gdrive_pagination.py @@ -1,6 +1,7 @@ from unittest.mock import Mock, patch from uuid import uuid4 +from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.connections.drive.mock_drive_service import ( MockDriveBackingStore, @@ -14,7 +15,7 @@ def create_mock_connection(): conn = GDriveConnection(email="test@test.com", verbose=False) - conn.drive_service = Mock() + conn.client = GDriveClient(Mock()) return conn @@ -22,7 +23,7 @@ def test_single_page_works(): """Current implementation works for single page (< 100 files)""" conn = create_mock_connection() - conn.drive_service.files().list().execute.return_value = { + conn.client.drive_service.files().list().execute.return_value = { "files": [{"name": "file1.txt", "id": "id1"}] } @@ -36,7 +37,7 @@ def test_multiple_pages_pagination(): """When there are >100 files, should fetch all pages""" conn = create_mock_connection() - conn.drive_service.files().list().execute.side_effect = [ + conn.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"}]}, @@ -45,14 +46,14 @@ def test_multiple_pages_pagination(): files = conn.get_file_metadatas_from_folder("folder_id") assert len(files) == 3 - assert conn.drive_service.files().list().execute.call_count == 3 + assert conn.client.drive_service.files().list().execute.call_count == 3 def test_empty_folder(): """Empty folder should return empty list""" conn = create_mock_connection() - conn.drive_service.files().list().execute.return_value = {"files": []} + conn.client.drive_service.files().list().execute.return_value = {"files": []} files = conn.get_file_metadatas_from_folder("folder_id") @@ -73,7 +74,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 = { + conn.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"}, @@ -113,12 +114,12 @@ def test_orderby_name_desc_applied(): """Drive API should be called with orderBy='name desc'""" conn = create_mock_connection() - conn.drive_service.files().list().execute.return_value = {"files": []} + conn.client.drive_service.files().list().execute.return_value = {"files": []} conn.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 = conn.client.drive_service.files().list.call_args[1] assert call_kwargs["orderBy"] == "name desc" @@ -136,7 +137,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 = [ + conn.client.drive_service.files().list().execute.side_effect = [ { "files": [ {"name": fname1.as_string(), "id": "id1"}, @@ -165,7 +166,7 @@ def test_early_termination_stops_pagination(): ) # 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 conn.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 @@ -178,7 +179,7 @@ def test_no_early_termination_when_all_files_new(): fname1 = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506450.0) fname2 = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506440.0) - conn.drive_service.files().list().execute.side_effect = [ + conn.client.drive_service.files().list().execute.side_effect = [ { "files": [{"name": fname1.as_string(), "id": "id1"}], "nextPageToken": "token1", @@ -194,7 +195,7 @@ def test_no_early_termination_when_all_files_new(): ) # Should fetch all pages (both files are > since_timestamp) - assert conn.drive_service.files().list().execute.call_count == 2 + assert conn.client.drive_service.files().list().execute.call_count == 2 assert len(files) == 2 @@ -202,7 +203,7 @@ def test_no_early_termination_when_since_timestamp_none(): """Should fetch all pages when since_timestamp is None""" conn = create_mock_connection() - conn.drive_service.files().list().execute.side_effect = [ + conn.client.drive_service.files().list().execute.side_effect = [ { "files": [{"name": "file1.txt", "id": "id1"}], "nextPageToken": "token1", @@ -216,7 +217,7 @@ def test_no_early_termination_when_since_timestamp_none(): files = conn.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 conn.client.drive_service.files().list().execute.call_count == 2 assert len(files) == 2 @@ -226,7 +227,7 @@ def test_newer_message_listed_after_older_one_in_same_second_not_skipped(): 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 = { + conn.client.drive_service.files().list().execute.return_value = { "files": [{"name": name, "id": name} for name in names] } diff --git a/tests/unit/test_login_email_check.py b/tests/unit/test_login_email_check.py index 2d828f1dc50..ed8926e965e 100644 --- a/tests/unit/test_login_email_check.py +++ b/tests/unit/test_login_email_check.py @@ -17,7 +17,7 @@ def test_verify_token_matches_email_passes_when_matching(): def _drive_service(manager): - return manager._connection_router.connections[0].drive_service + return manager._connection_router.connections[0].client.drive_service def test_verify_token_matches_email_raises_on_mismatch(): diff --git a/tests/unit/test_p2p_folder_lookup.py b/tests/unit/test_p2p_folder_lookup.py index 3714e55ec8f..3b3fb9b2a2b 100644 --- a/tests/unit/test_p2p_folder_lookup.py +++ b/tests/unit/test_p2p_folder_lookup.py @@ -11,6 +11,7 @@ from unittest.mock import Mock +from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_transport import GDriveConnection ME = "alice@example.com" @@ -23,7 +24,7 @@ def _name(version: str, datasite: str, folder_type: str, peer: str) -> str: def _conn(found): conn = GDriveConnection(email=ME, verbose=False) - conn.drive_service = Mock() + conn.client = GDriveClient(Mock()) conn._find_folders = Mock(return_value=found) return conn diff --git a/tests/unit/test_peer_request_owner.py b/tests/unit/test_peer_request_owner.py index 80a25aadf8d..d486d9ffaf4 100644 --- a/tests/unit/test_peer_request_owner.py +++ b/tests/unit/test_peer_request_owner.py @@ -48,7 +48,7 @@ def share_request_folder( """Put a request folder for `victim` in Drive and share it with `victim`.""" backing_store = victim._connection_router.connections[ 0 - ].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 15a580eadfe..fb64776b1e0 100644 --- a/tests/unit/test_peers_json_version.py +++ b/tests/unit/test_peers_json_version.py @@ -14,6 +14,7 @@ import logging from unittest.mock import Mock, patch +from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.peers.peer import PeerState from syft.sync.transport.naming import PEERS_META_KEY, SYFT_PEERS_VERSION @@ -23,7 +24,7 @@ def _conn(peers_data): conn = GDriveConnection(email="alice@example.com", verbose=False) - conn.drive_service = Mock() + conn.client = GDriveClient(Mock()) conn._peers_json_cache = dict(peers_data) return conn diff --git a/tests/unit/test_sender_binding.py b/tests/unit/test_sender_binding.py index 883d7f4d5e3..3ff59c951c5 100644 --- a/tests/unit/test_sender_binding.py +++ b/tests/unit/test_sender_binding.py @@ -82,7 +82,7 @@ def test_spoofed_sender_does_not_name_an_archive_folder(): connection = do_manager.datasite_owner_syncer.connection_router.connections[0] folder_names = [ - f.name for f in connection.drive_service._backing_store.files.values() + f.name for f in connection.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_sync_manager.py b/tests/unit/test_sync_manager.py index 4f84b621823..aaaee0fcd3f 100644 --- a/tests/unit/test_sync_manager.py +++ b/tests/unit/test_sync_manager.py @@ -487,7 +487,7 @@ 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( - email, ds_manager._connection_router.connections[0].drive_service + email, ds_manager._connection_router.connections[0].client.drive_service ) conn.add_peer(do_manager.email) _add_known_peer(do_manager, email) diff --git a/tests/unit/test_version_negotiation.py b/tests/unit/test_version_negotiation.py index dfe0cc6c063..417a6f044de 100644 --- a/tests/unit/test_version_negotiation.py +++ b/tests/unit/test_version_negotiation.py @@ -241,7 +241,7 @@ def test_load_peer_versions_parallel(self): backing_store = ds_manager._connection_router.connections[ 0 - ].drive_service._backing_store + ].client.drive_service._backing_store third_peer_email = "third_peer@test.com" third_peer_info = PeerInfo( diff --git a/tests/unit/test_versioned_folder_adopt.py b/tests/unit/test_versioned_folder_adopt.py index 44b786be8f5..1837ad848c4 100644 --- a/tests/unit/test_versioned_folder_adopt.py +++ b/tests/unit/test_versioned_folder_adopt.py @@ -18,6 +18,7 @@ import pytest +from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.checkpoints.checkpoint import CHECKPOINT_VERSION, Checkpoint from syft.sync.connections.drive.gdrive_transport import ( GDriveConnection, @@ -31,7 +32,7 @@ def _conn(): conn = GDriveConnection(email=EMAIL, verbose=False) - conn.drive_service = Mock() + conn.client = GDriveClient(Mock()) return conn @@ -39,7 +40,7 @@ def _renames(conn): """Return the (fileId, new name) pairs the connection sent to Drive.""" return [ (kwargs["fileId"], kwargs["body"]["name"]) - for _, kwargs in conn.drive_service.files().update.call_args_list + for _, kwargs in conn.client.drive_service.files().update.call_args_list if "body" in kwargs and "name" in kwargs.get("body", {}) ] @@ -166,7 +167,7 @@ def _checkpoints_owner_with_version(checkpoint_version: int): folder_id = connection._get_checkpoints_folder_id() older_name = f"{do_manager.email}-0.0.1-checkpoints" execute_with_retries( - connection.drive_service.files().update( + connection.client.drive_service.files().update( fileId=folder_id, body={"name": older_name} ) ) From 2d1803c613dedc8c77b90affdf467a8d87ca80b1 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 10:08:59 -0300 Subject: [PATCH 07/28] refactor: move the Drive file operations into GDriveClient --- .../tests/test_version_mismatch_flow.py | 2 +- syft/gdrive_utils.py | 16 +- syft/sync/connections/drive/gdrive_client.py | 124 ++++++++++ .../connections/drive/gdrive_transport.py | 229 +++++------------- tests/unit/test_drive_connections.py | 4 +- tests/unit/test_gdrive_pagination.py | 2 +- tests/unit/test_mock_drive_service.py | 2 +- tests/unit/test_peers_json_version.py | 4 +- tests/unit/test_rolling_state.py | 4 +- tests/unit/test_version_negotiation.py | 14 +- 10 files changed, 211 insertions(+), 190 deletions(-) diff --git a/packages/syft-rds/tests/test_version_mismatch_flow.py b/packages/syft-rds/tests/test_version_mismatch_flow.py index 8b90b94bca5..88433d48304 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -7,9 +7,9 @@ from unittest.mock import patch +from syft.sync.connections.drive.gdrive_client import GOOGLE_FOLDER_MIME_TYPE from syft.sync.connections.drive.gdrive_transport import ( GDRIVE_P2P_FOLDER_DATASITE_PREFIX, - GOOGLE_FOLDER_MIME_TYPE, GDriveConnection, ) from syft.sync.connections.drive.mock_drive_service import ( diff --git a/syft/gdrive_utils.py b/syft/gdrive_utils.py index b84201edb55..bba196110b5 100644 --- a/syft/gdrive_utils.py +++ b/syft/gdrive_utils.py @@ -4,29 +4,27 @@ 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, - GOOGLE_FOLDER_MIME_TYPE, -) -from syft.sync.connections.drive.gdrive_transport import SCOPES as GDRIVE_SCOPES +from syft.sync.config.config import settings +from syft.sync.connections.drive.gdrive_client import GOOGLE_FOLDER_MIME_TYPE from syft.sync.connections.drive.gdrive_retry import ( execute_with_retries, next_chunk_with_retries, ) +from syft.sync.connections.drive.gdrive_transport import SCOPES as GDRIVE_SCOPES +from syft.sync.connections.drive.gdrive_transport import build_drive_service +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) diff --git a/syft/sync/connections/drive/gdrive_client.py b/syft/sync/connections/drive/gdrive_client.py index 763c7008728..36b7aa88ede 100644 --- a/syft/sync/connections/drive/gdrive_client.py +++ b/syft/sync/connections/drive/gdrive_client.py @@ -4,11 +4,135 @@ use from more than one thread. """ +import io +import json +import pickle from dataclasses import dataclass from typing import Any +from googleapiclient.http import MediaIoBaseDownload, MediaIoBaseUpload + +from syft.sync.connections.drive.gdrive_retry import ( + batch_execute_with_retries, + execute_with_retries, + next_chunk_with_retries, +) + +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" + @dataclass class GDriveClient: # A googleapiclient Drive v3 service, or a MockDriveService in tests. drive_service: Any + + 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, 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 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 + + # Google Drive batch API has a limit of 100 requests per batch + BATCH_SIZE = 100 + for i in range(0, len(file_ids), BATCH_SIZE): + chunk = file_ids[i : i + 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) diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index d1f0af4d096..fdf62f707ca 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -1,9 +1,7 @@ """Google Drive Files transport layer implementation""" -import io import json import logging -import pickle import re from functools import lru_cache from pathlib import Path @@ -12,7 +10,7 @@ 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 googleapiclient.http import build_http from pydantic import BaseModel, field_validator from syft_migration import MigrationError @@ -31,12 +29,16 @@ SyftboxPlatformConnection, ) from syft.sync.connections.drive.drive_connections import drive_transport -from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.gdrive_client import ( + GOOGLE_FOLDER_MIME_TYPE, + SYFT_FILE_PROPERTY_KEY, + SYFT_FILE_PROPERTY_VALUE, + GDriveClient, +) from syft.sync.connections.drive.gdrive_retry import ( batch_execute_with_retries, execute_with_retries, is_unreachable, - next_chunk_with_retries, raise_unavailable, ) from syft.sync.connections.drive.gdrive_utils import ( @@ -72,7 +74,6 @@ GOOGLE_API_TIMEOUT = 120 # 2 minutes SYFTBOX_FOLDER = "SyftBox" -GOOGLE_FOLDER_MIME_TYPE = "application/vnd.google-apps.folder" SCOPES = ["https://www.googleapis.com/auth/drive"] @@ -107,9 +108,6 @@ def build_drive_service( 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): @@ -505,11 +503,11 @@ 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) + return self.client.create_folder(folder_name, syftbox_folder_id) def create_syftbox_folder(self) -> str: """Creates /SyftBox""" - folder = self._create_syft_file( + folder = self.client.create_syft_file( {"name": SYFTBOX_FOLDER, "mimeType": GOOGLE_FOLDER_MIME_TYPE} ) return folder["id"] @@ -520,7 +518,7 @@ def create_archive_folder(self, sender_email: str) -> str: ) 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) + return self.client.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. @@ -569,7 +567,7 @@ def _download_peers_json(self) -> dict[str, dict[str, Any]]: file_id = self._get_peers_file_id() if file_id is None: return {} - file_data = self.download_file(file_id) + file_data = self.client.download_file(file_id) try: return json.loads(file_data.decode("utf-8")) except ValueError as e: @@ -595,7 +593,7 @@ def _write_peers_json(self, peers_data: dict[str, dict[str, Any]]): # Convert to JSON bytes json_data = json.dumps(peers_data, indent=2) - file_payload, _ = self.create_file_payload(json_data) + file_payload, _ = self.client.create_file_payload(json_data) if file_id is None: # Create new file @@ -603,7 +601,7 @@ def _write_peers_json(self, peers_data: dict[str, dict[str, Any]]): "name": SYFT_PEERS_FILE, "parents": [syftbox_folder_id], } - result = self._create_syft_file(file_metadata, file_payload) + result = self.client.create_syft_file(file_metadata, file_payload) self._peers_json_cache = peers_data return result.get("id") else: @@ -708,7 +706,7 @@ def watcher_download_raw_events_from_outbox( 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])) + res.append(self.client.download_file(name_to_id[file_name])) return res def watcher_get_events_messages( @@ -754,16 +752,16 @@ def owner_write_raw_bytes_to_syftbox(self, filename: str, data: bytes) -> str: "name": filename, "parents": [personal_syftbox_folder_id], } - file_payload, _ = self.create_file_payload(data) + file_payload, _ = self.client.create_file_payload(data) - res = self._create_syft_file(file_metadata, file_payload) + res = self.client.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) + return self.client.download_file(file_id) def owner_get_all_accepted_event_file_ids( self, since_timestamp: float | None = None @@ -793,7 +791,7 @@ def owner_download_all_raw_events_from_syftbox(self) -> list[bytes]: f for f in file_metadatas if f["name"] == fname_obj.as_string() ][0]["id"] try: - file_data = self.download_file(gdrive_id) + file_data = self.client.download_file(gdrive_id) except Exception as e: print(f"Warning: could not download event {fname_obj.as_string()}: {e}") continue @@ -807,10 +805,10 @@ def owner_write_raw_bytes_to_outbox( if outbox_folder_id is None: raise ValueError(f"Outbox folder for {recipient} not found") - file_payload, _ = self.create_file_payload(data) + file_payload, _ = self.client.create_file_payload(data) file_metadata = {"name": filename, "parents": [outbox_folder_id]} - result = self._create_syft_file( + result = self.client.create_syft_file( file_metadata, file_payload, fields="id, parents" ) @@ -917,7 +915,7 @@ def _create_peer_datasite_folder(self, peer_email: str, folder_type: str) -> str 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) + folder_id = self.client.create_folder(folder_name, parent_id) if folder_type == "inbox": self.peer_datasite_inbox_cache[peer_email] = folder_id else: @@ -1112,7 +1110,7 @@ def owner_download_next_raw_proposed_message_from_inbox( first_file_id = [ x for x in file_metadatas if x["name"] == first_file_name.as_string() ][0]["id"] - raw_data = self.download_file(first_file_id) + raw_data = self.client.download_file(first_file_id) return raw_data, first_file_id def owner_get_next_proposed_filechange_message( @@ -1234,10 +1232,12 @@ def watcher_send_raw_bytes_to_inbox( if inbox_id is None: raise Exception(f"Inbox folder for {recipient}'s datasite not found") - payload, _ = self.create_file_payload(data) + payload, _ = self.client.create_file_payload(data) file_metadata = {"name": filename, "parents": [inbox_id]} - result = self._create_syft_file(file_metadata, payload, fields="id, parents") + result = self.client.create_syft_file( + file_metadata, payload, fields="id, parents" + ) file_id = result.get("id") actual_parents = result.get("parents", []) @@ -1285,55 +1285,23 @@ def gather_all_file_and_folder_ids(self) -> list[str]: self.client.drive_service, syftbox_folder_id ) + def download_file(self, file_id: str) -> bytes: + return self.client.download_file(file_id) + + def delete_file_by_id(self, file_id: str, raise_on_error: bool = False): + self.client.delete_file_by_id(file_id, raise_on_error=raise_on_error) + 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')}" - ) - - # Google Drive batch API has a limit of 100 requests per batch - BATCH_SIZE = 100 - for i in range(0, len(file_ids), BATCH_SIZE): - chunk = file_ids[i : i + BATCH_SIZE] - batch = self.client.drive_service.new_batch_http_request(callback=callback) - for file_id in chunk: - batch.add(self.client.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.client.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}") + self.client.delete_multiple_files_by_ids( + file_ids, + ignore_permissions_errors=ignore_permissions_errors, + ignore_file_not_found=ignore_file_not_found, + ) def find_orphaned_syft_files(self) -> list[str]: """Find syft files of this account in its My Drive root, except /SyftBox. @@ -1362,32 +1330,6 @@ def find_orphaned_syft_files(self) -> list[str]: 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]: @@ -1530,49 +1472,6 @@ def _find_or_adopt_versioned_folder( ) return folder_id - def download_file(self, file_id: str) -> bytes: - request = self.client.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.client.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( @@ -1621,7 +1520,7 @@ def owner_create_collection_folder( return folder_id # Create new folder - folder_id = self.create_folder(folder_name, syftbox_folder_id) + folder_id = self.client.create_folder(folder_name, syftbox_folder_id) self.collection_folder_id_cache[folder_name] = folder_id return folder_id @@ -1680,11 +1579,11 @@ def owner_upload_collection_files( 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_payload, _ = self.client.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) + self.client.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. @@ -1742,7 +1641,7 @@ def owner_delete_collection(self, prefix: str, tag: str) -> None: 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) + self.client.delete_file_by_id(c.folder_id) folder_name = CollectionFolder( prefix=prefix, tag=c.tag, @@ -1796,7 +1695,7 @@ def watcher_download_collection( for file_meta in file_metadatas: file_id = file_meta["id"] file_name = file_meta["name"] - files[file_name] = self.download_file(file_id) + files[file_name] = self.client.download_file(file_id) return files @@ -1817,7 +1716,7 @@ def watcher_get_collection_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) + return self.client.download_file(file_id) def _get_collection_folder_id( self, prefix: str, tag: str, content_hash: str @@ -1858,7 +1757,7 @@ def write_peer_info_file(self, peer_info: "PeerInfo") -> None: # Convert to JSON string json_data = peer_info.to_json() - file_payload, _ = self.create_file_payload(json_data) + file_payload, _ = self.client.create_file_payload(json_data) if file_id is None: # Create new file @@ -1866,7 +1765,7 @@ def write_peer_info_file(self, peer_info: "PeerInfo") -> None: "name": SYFT_VERSION_FILE, "parents": [syftbox_folder_id], } - self._create_syft_file(file_metadata, file_payload) + self.client.create_syft_file(file_metadata, file_payload) else: # Update existing file execute_with_retries( @@ -1927,7 +1826,7 @@ def _read_peer_info_file( file_id = find_file_id() if file_id is None: return None - file_data = self.download_file(file_id) + file_data = self.client.download_file(file_id) try: return PeerInfo.from_json(file_data.decode("utf-8")) except (ValueError, MigrationError) as e: @@ -1974,7 +1873,7 @@ def _get_or_create_checkpoints_folder_id(self) -> str: # 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) + return self.client.create_folder(folder_name, syftbox_folder_id) def upload_raw_checkpoint(self, filename: str, data: bytes) -> str: """Upload raw checkpoint bytes to Google Drive. @@ -1982,14 +1881,14 @@ def upload_raw_checkpoint(self, filename: str, data: bytes) -> str: Uploads the new checkpoint first, then deletes old ones. """ folder_id = self._get_or_create_checkpoints_folder_id() - payload, _ = self.create_file_payload(data) + payload, _ = self.client.create_file_payload(data) file_metadata = { "name": filename, "parents": [folder_id], } - result = self._create_syft_file(file_metadata, payload) + result = self.client.create_syft_file(file_metadata, payload) # Only delete old checkpoints after successful upload self.delete_all_checkpoints(exclude_file_id=result.get("id")) @@ -2026,7 +1925,7 @@ def download_raw_latest_checkpoint(self) -> bytes | None: return None try: - return self.download_file(latest_file["id"]) + return self.client.download_file(latest_file["id"]) except Exception as e: if is_unreachable(e): raise TransportUnavailableError( @@ -2076,14 +1975,14 @@ def delete_all_checkpoints(self, exclude_file_id: str | None = None): 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) + payload, _ = self.client.create_file_payload(data) file_metadata = { "name": filename, "parents": [folder_id], } - result = self._create_syft_file(file_metadata, payload) + result = self.client.create_syft_file(file_metadata, payload) return result.get("id") def download_all_raw_incremental_checkpoints(self) -> list[bytes]: @@ -2107,7 +2006,7 @@ def download_all_raw_incremental_checkpoints(self) -> list[bytes]: result = [] for item in items: try: - file_data = self.download_file(item["id"]) + file_data = self.client.download_file(item["id"]) result.append(file_data) except Exception as e: if is_unreachable(e): @@ -2243,7 +2142,7 @@ def download_raw_events_since_timestamp( ][0]["id"] try: - file_data = self.download_file(gdrive_id) + file_data = self.client.download_file(gdrive_id) result.append(file_data) except Exception as e: print(f"Warning: Failed to download event: {e}") @@ -2288,7 +2187,7 @@ def _get_or_create_rolling_state_folder_id(self) -> str: # 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) + folder_id = self.client.create_folder(folder_name, syftbox_folder_id) self._rolling_state_folder_id = folder_id return folder_id @@ -2297,7 +2196,7 @@ def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: Optimized to use update() if file ID is cached. """ - payload, _ = self.create_file_payload(data) + payload, _ = self.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: @@ -2325,7 +2224,7 @@ def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: "parents": [folder_id], } - result = self._create_syft_file(file_metadata, payload) + result = self.client.create_syft_file(file_metadata, payload) file_id = result.get("id") self._rolling_state_file_id = file_id return file_id @@ -2365,7 +2264,7 @@ def download_raw_rolling_state(self) -> bytes | None: self._rolling_state_file_id = latest_file["id"] try: - return self.download_file(latest_file["id"]) + return self.client.download_file(latest_file["id"]) except Exception as e: if is_unreachable(e): raise TransportUnavailableError( @@ -2424,7 +2323,9 @@ def _get_or_create_encryption_bundles_folder_id(self) -> str: 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()) + folder_id = self.client.create_folder( + folder_name, self.get_syftbox_folder_id() + ) self._encryption_bundles_folder_id = folder_id return folder_id @@ -2435,7 +2336,7 @@ 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) + file_payload, _ = self.client.create_file_payload(bundle_json) # Check if file already exists query = f"name='{filename}' and '{folder_id}' in parents and trashed=false" @@ -2452,7 +2353,7 @@ def write_encryption_bundle(self, peer_email: str, bundle_json: str) -> None: ) else: file_metadata = {"name": filename, "parents": [folder_id]} - self._create_syft_file(file_metadata, file_payload) + self.client.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.""" @@ -2475,7 +2376,7 @@ def read_own_encryption_bundle(self, peer_email: str) -> str | None: items = results.get("files", []) if not items: return None - return self.download_file(items[0]["id"]).decode("utf-8") + return self.client.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. @@ -2492,7 +2393,7 @@ def read_peer_encryption_bundle(self, peer_email: str) -> str | None: if not items: return None try: - data = self.download_file(items[0]["id"]) + data = self.client.download_file(items[0]["id"]) except Exception as e: print(f"Warning: could not download the bundle of {peer_email}: {e}") return None diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 2473de6933a..345ab5d9bc3 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -13,9 +13,9 @@ 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.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_transport import ( CollectionFolder, - GDriveConnection, ) from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service from syft.sync.encryption_policy import EncryptionPolicy @@ -260,7 +260,7 @@ def test_peer_registry_update_while_drive_is_down_keeps_other_peers(pair, monkey peering.update_peer_registry(DS, "accepted") do_conn._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() diff --git a/tests/unit/test_gdrive_pagination.py b/tests/unit/test_gdrive_pagination.py index 9ff94ebbefd..fdd23cc271d 100644 --- a/tests/unit/test_gdrive_pagination.py +++ b/tests/unit/test_gdrive_pagination.py @@ -63,7 +63,7 @@ def test_empty_folder(): @patch( "syft.sync.connections.drive.gdrive_transport.GDriveConnection._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() diff --git a/tests/unit/test_mock_drive_service.py b/tests/unit/test_mock_drive_service.py index ea40cd20d4c..f7262a15a92 100644 --- a/tests/unit/test_mock_drive_service.py +++ b/tests/unit/test_mock_drive_service.py @@ -449,7 +449,7 @@ def test_create_folder(self): connection = GDriveConnection.from_service("user@example.com", mock_service) syftbox_id = connection.create_syftbox_folder() - folder_id = connection.create_folder("TestFolder", syftbox_id) + folder_id = connection.client.create_folder("TestFolder", syftbox_id) assert store.files[syftbox_id].properties == {"syft": "true"} assert store.files[folder_id].properties == {"syft": "true"} diff --git a/tests/unit/test_peers_json_version.py b/tests/unit/test_peers_json_version.py index fb64776b1e0..1be702af753 100644 --- a/tests/unit/test_peers_json_version.py +++ b/tests/unit/test_peers_json_version.py @@ -44,9 +44,7 @@ def test_a_write_stamps_the_reserved_entry(): patch.object( GDriveConnection, "get_syftbox_folder_id", return_value="folder-id" ), - patch.object( - GDriveConnection, "create_file_payload", return_value=(Mock(), None) - ), + patch.object(GDriveClient, "create_file_payload", return_value=(Mock(), None)), ): conn._write_peers_json({PEER: {"state": "accepted"}}) diff --git a/tests/unit/test_rolling_state.py b/tests/unit/test_rolling_state.py index 3b9c8cbe85f..64ba7096ff3 100644 --- a/tests/unit/test_rolling_state.py +++ b/tests/unit/test_rolling_state.py @@ -14,7 +14,7 @@ ROLLING_STATE_FILENAME_PREFIX, RollingState, ) -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service from syft.sync.syftbox_manager import SyftboxManager from syft.sync.sync.constants import CACHE_DIR, ROLLING_STATE_FILENAME @@ -298,6 +298,6 @@ def test_rolling_state_read_while_drive_is_down_raises(monkeypatch): def drive_down(self, file_id: str) -> bytes: raise ConnectionResetError("Drive is down") - monkeypatch.setattr(GDriveConnection, "download_file", drive_down) + monkeypatch.setattr(GDriveClient, "download_file", drive_down) with pytest.raises(TransportUnavailableError): store.get_rolling_state() diff --git a/tests/unit/test_version_negotiation.py b/tests/unit/test_version_negotiation.py index 417a6f044de..a6af55d91d1 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 ( @@ -283,7 +283,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]) @@ -292,14 +292,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() @@ -309,7 +309,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]) @@ -322,7 +322,7 @@ def test_drive_refusal_during_peer_version_read_not_unavailable(self, monkeypatc 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)) with pytest.raises(HttpError): ds_manager.peer_manager.load_peer_infos_parallel([do_manager.email]) @@ -331,7 +331,7 @@ def test_unreadable_peer_version_file_raises(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")) with pytest.raises(ValueError): ds_manager.peer_manager.load_peer_infos_parallel([do_manager.email]) From fcfc5b4d64c6503519f3bf57203092840c74f142 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 11:15:50 -0300 Subject: [PATCH 08/28] fix: raise on owner-state read errors and match delete errors by status --- syft/sync/connections/drive/gdrive_client.py | 32 +++++----- syft/sync/connections/drive/gdrive_retry.py | 8 +-- .../connections/drive/gdrive_transport.py | 60 ++++--------------- .../connections/drive/mock_drive_service.py | 20 +++++-- tests/unit/test_checkpoints.py | 20 ++++++- tests/unit/test_drive_connections.py | 29 ++++++++- tests/unit/test_gdrive_client.py | 22 +++++++ tests/unit/test_rolling_state.py | 33 +++++++++- 8 files changed, 146 insertions(+), 78 deletions(-) create mode 100644 tests/unit/test_gdrive_client.py diff --git a/syft/sync/connections/drive/gdrive_client.py b/syft/sync/connections/drive/gdrive_client.py index 36b7aa88ede..d016f879285 100644 --- a/syft/sync/connections/drive/gdrive_client.py +++ b/syft/sync/connections/drive/gdrive_client.py @@ -10,6 +10,7 @@ from dataclasses import dataclass from typing import Any +from googleapiclient.errors import HttpError from googleapiclient.http import MediaIoBaseDownload, MediaIoBaseUpload from syft.sync.connections.drive.gdrive_retry import ( @@ -24,6 +25,12 @@ SYFT_FILE_PROPERTY_VALUE = "true" +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. @@ -113,20 +120,17 @@ def delete_multiple_files_by_ids( 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 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 # Google Drive batch API has a limit of 100 requests per batch BATCH_SIZE = 100 diff --git a/syft/sync/connections/drive/gdrive_retry.py b/syft/sync/connections/drive/gdrive_retry.py index 95658d77d9b..e861994eca3 100644 --- a/syft/sync/connections/drive/gdrive_retry.py +++ b/syft/sync/connections/drive/gdrive_retry.py @@ -72,12 +72,6 @@ def is_retryable_error(error: Exception) -> bool: return False -def is_unreachable(error: Exception) -> bool: - """True when ``error`` means Drive cannot be reached: an error in - UNREACHABLE_ERRORS, or a transient error that outlasted the retries.""" - return isinstance(error, UNREACHABLE_ERRORS) or is_retryable_error(error) - - @contextmanager def raise_unavailable(message: str) -> Iterator[None]: """Raise TransportUnavailableError(message) when Drive cannot be reached. @@ -89,7 +83,7 @@ def raise_unavailable(message: str) -> Iterator[None]: try: yield except (HttpError, *RETRYABLE_TRANSPORT_ERRORS, *UNREACHABLE_ERRORS) as e: - if is_unreachable(e): + if isinstance(e, UNREACHABLE_ERRORS) or is_retryable_error(e): raise TransportUnavailableError(message) from e raise diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index fdf62f707ca..1ea22a68a5c 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -38,7 +38,6 @@ from syft.sync.connections.drive.gdrive_retry import ( batch_execute_with_retries, execute_with_retries, - is_unreachable, raise_unavailable, ) from syft.sync.connections.drive.gdrive_utils import ( @@ -54,7 +53,6 @@ MessageFileName, ProposedFileChangesMessage, ) -from syft.sync.transport.exceptions import TransportUnavailableError from syft.sync.transport.naming import ( PEERS_META_KEY, SYFT_PEERS_VERSION, @@ -779,6 +777,7 @@ def owner_get_all_accepted_event_file_ids( ) return [f["id"] for f in oldest_first] + @raise_unavailable("Could not read the event log from Drive.") 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() @@ -790,12 +789,7 @@ def owner_download_all_raw_events_from_syftbox(self) -> list[bytes]: gdrive_id = [ f for f in file_metadatas if f["name"] == fname_obj.as_string() ][0]["id"] - try: - file_data = self.client.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) + result.append(self.client.download_file(gdrive_id)) return result def owner_write_raw_bytes_to_outbox( @@ -1895,6 +1889,7 @@ def upload_raw_checkpoint(self, filename: str, data: bytes) -> str: return result.get("id") + @raise_unavailable("Could not read the latest checkpoint from Drive.") def download_raw_latest_checkpoint(self) -> bytes | None: """Download the latest checkpoint as raw bytes, or None.""" folder_id = self._get_checkpoints_folder_id() @@ -1924,15 +1919,7 @@ def download_raw_latest_checkpoint(self) -> bytes | None: if latest_file is None: return None - try: - return self.client.download_file(latest_file["id"]) - except Exception as e: - if is_unreachable(e): - raise TransportUnavailableError( - "Could not read the latest checkpoint from Drive." - ) from e - print(f"Warning: Failed to load checkpoint: {e}") - return None + return self.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). @@ -1985,6 +1972,7 @@ def upload_raw_incremental_checkpoint(self, filename: str, data: bytes) -> str: result = self.client.create_syft_file(file_metadata, payload) return result.get("id") + @raise_unavailable("Could not read the incremental checkpoints from Drive.") def download_all_raw_incremental_checkpoints(self) -> list[bytes]: """Download all incremental checkpoints as raw bytes.""" folder_id = self._get_checkpoints_folder_id() @@ -2003,21 +1991,7 @@ def download_all_raw_incremental_checkpoints(self) -> list[bytes]: if not items: return [] - result = [] - for item in items: - try: - file_data = self.client.download_file(item["id"]) - result.append(file_data) - except Exception as e: - if is_unreachable(e): - raise TransportUnavailableError( - "Could not read the incremental checkpoints from Drive." - ) from e - print( - f"Warning: Failed to load incremental checkpoint {item['name']}: {e}" - ) - continue - return result + return [self.client.download_file(item["id"]) for item in items] def get_incremental_checkpoint_count(self) -> int: """Get the number of incremental checkpoints on Google Drive.""" @@ -2121,6 +2095,7 @@ def get_events_count_since_checkpoint( count += 1 return count + @raise_unavailable("Could not read the event log from Drive.") def download_raw_events_since_timestamp( self, since_timestamp: float ) -> list[bytes]: @@ -2141,12 +2116,7 @@ def download_raw_events_since_timestamp( f for f in file_metadatas if f["name"] == fname_obj.as_string() ][0]["id"] - try: - file_data = self.client.download_file(gdrive_id) - result.append(file_data) - except Exception as e: - print(f"Warning: Failed to download event: {e}") - continue + result.append(self.client.download_file(gdrive_id)) return result @@ -2229,6 +2199,7 @@ def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: self._rolling_state_file_id = file_id return file_id + @raise_unavailable("Could not read the rolling state from Drive.") def download_raw_rolling_state(self) -> bytes | None: """Download the latest rolling state as raw bytes, or None. @@ -2261,18 +2232,9 @@ def download_raw_rolling_state(self) -> bytes | None: if latest_file is None: return None + data = self.client.download_file(latest_file["id"]) self._rolling_state_file_id = latest_file["id"] - - try: - return self.client.download_file(latest_file["id"]) - except Exception as e: - if is_unreachable(e): - raise TransportUnavailableError( - "Could not read the rolling state from Drive." - ) from e - print(f"Warning: Failed to load rolling state: {e}") - self._rolling_state_file_id = None - return None + return data def delete_rolling_state(self) -> None: """Delete all existing rolling state files and clear cache.""" diff --git a/syft/sync/connections/drive/mock_drive_service.py b/syft/sync/connections/drive/mock_drive_service.py index 58c60fa0876..162fd763176 100644 --- a/syft/sync/connections/drive/mock_drive_service.py +++ b/syft/sync/connections/drive/mock_drive_service.py @@ -21,10 +21,13 @@ connection = GDriveConnection.from_mock_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 @@ -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) diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index a6b917e09ff..b33bc12adfc 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -11,6 +11,7 @@ from uuid import uuid4 import pytest +from googleapiclient.errors import HttpError from syft.sync.checkpoints.checkpoint import ( CHECKPOINT_FILENAME_PREFIX, @@ -19,7 +20,11 @@ IncrementalCheckpoint, ) from syft.sync.checkpoints.rolling_state import RollingState -from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service +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 syft.sync.sync import datasite_owner_syncer @@ -761,3 +766,16 @@ def test_incremental_count_retries_drive_error(monkeypatch): fail_next_listing(monkeypatch, INCREMENTAL_CHECKPOINT_PREFIX) assert store.incremental_count() == 1 + + +def test_latest_checkpoint_download_error_raises(monkeypatch): + do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = do_conn.as_transport().owner_store + 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() diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 345ab5d9bc3..d406d9979e9 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -9,6 +9,7 @@ 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 @@ -17,7 +18,10 @@ from syft.sync.connections.drive.gdrive_transport import ( CollectionFolder, ) -from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service +from syft.sync.connections.drive.mock_drive_service import ( + http_error, + 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 @@ -155,6 +159,29 @@ def test_owner_event_count_matches_the_listed_handles(pair): 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_conn, _ = pair + store = do_conn.as_transport().owner_store + 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_conn, _ = pair + store = do_conn.as_transport().owner_store + 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 = do_conn.as_transport().owner_store diff --git a/tests/unit/test_gdrive_client.py b/tests/unit/test_gdrive_client.py new file mode 100644 index 00000000000..aa241634a8f --- /dev/null +++ b/tests/unit/test_gdrive_client.py @@ -0,0 +1,22 @@ +"""GDriveClient on the mock Drive service.""" + +import pytest +from googleapiclient.errors import HttpError + +from syft.sync.connections.drive.mock_drive_service import ( + MockDeleteRequest, + http_error, + pair_with_mock_service, +) + + +def test_batch_delete_raises_error_whose_text_contains_404(monkeypatch): + do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") + folder_id = do_conn.client.create_folder("folder", do_conn.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_conn.client.delete_multiple_files_by_ids([folder_id]) diff --git a/tests/unit/test_rolling_state.py b/tests/unit/test_rolling_state.py index 64ba7096ff3..1f6b74d962f 100644 --- a/tests/unit/test_rolling_state.py +++ b/tests/unit/test_rolling_state.py @@ -9,13 +9,18 @@ import time import pytest +from googleapiclient.errors import HttpError 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 pair_with_mock_service +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.sync.constants import CACHE_DIR, ROLLING_STATE_FILENAME from syft.sync.transport import TransportUnavailableError @@ -301,3 +306,29 @@ def drive_down(self, file_id: str) -> bytes: 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_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = do_conn.as_transport().owner_store + 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_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") + store = do_conn.as_transport().owner_store + 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() From 5f53243fafbd28f3ce6bb6f6ee700ae4c95c1abe Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 11:36:21 -0300 Subject: [PATCH 09/28] refactor: move the Drive permission calls and folder search into GDriveClient --- .../tests/test_version_mismatch_flow.py | 8 +- syft/sync/connections/drive/gdrive_client.py | 119 ++++++++++- .../connections/drive/gdrive_transport.py | 190 +++--------------- tests/unit/test_p2p_folder_lookup.py | 2 +- tests/unit/test_versioned_folder_adopt.py | 2 +- 5 files changed, 151 insertions(+), 170 deletions(-) diff --git a/packages/syft-rds/tests/test_version_mismatch_flow.py b/packages/syft-rds/tests/test_version_mismatch_flow.py index 88433d48304..ab0d2697ca6 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -180,7 +180,7 @@ def test_version_mismatch_continues_and_repairs(): ) 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_conn.client.find_folder_by_name( old_personal_name, parent_id=do_conn.get_syftbox_folder_id(), owner_email=do_manager.email, @@ -215,7 +215,7 @@ def test_version_mismatch_continues_and_repairs(): # Personal folder is adopted (same Drive id, new name), not recreated. do_conn_new = do_manager.peer_manager.connection_router.connections[0] new_personal_name = f"{NEW_VERSION}#{do_email}" - new_personal_id = do_conn_new._find_folder_by_name( + new_personal_id = do_conn_new.client.find_folder_by_name( new_personal_name, parent_id=do_conn_new.get_syftbox_folder_id(), owner_email=do_email, @@ -223,7 +223,7 @@ def test_version_mismatch_continues_and_repairs(): assert new_personal_id is not None assert new_personal_id == old_personal_id assert ( - do_conn_new._find_folder_by_name( + do_conn_new.client.find_folder_by_name( old_personal_name, parent_id=do_conn_new.get_syftbox_folder_id(), owner_email=do_email, @@ -251,7 +251,7 @@ def test_version_mismatch_continues_and_repairs(): ds_conn_new = ds_manager.peer_manager.connection_router.connections[0] ds_personal_name = f"{NEW_VERSION}#{ds_email}" - ds_personal_id = ds_conn_new._find_folder_by_name( + ds_personal_id = ds_conn_new.client.find_folder_by_name( ds_personal_name, parent_id=ds_conn_new.get_syftbox_folder_id(), owner_email=ds_email, diff --git a/syft/sync/connections/drive/gdrive_client.py b/syft/sync/connections/drive/gdrive_client.py index d016f879285..7963fcd5d82 100644 --- a/syft/sync/connections/drive/gdrive_client.py +++ b/syft/sync/connections/drive/gdrive_client.py @@ -8,7 +8,7 @@ import json import pickle from dataclasses import dataclass -from typing import Any +from typing import Any, Optional from googleapiclient.errors import HttpError from googleapiclient.http import MediaIoBaseDownload, MediaIoBaseUpload @@ -140,3 +140,120 @@ def callback(request_id, response, exception): 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): + # Ignore "already shared" errors + if exception and "alreadyShared" not in str(exception): + raise exception + + BATCH_SIZE = 100 + for i in range(0, len(users), BATCH_SIZE): + chunk = users[i : i + 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 index 1ea22a68a5c..889a86a5e73 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -36,7 +36,6 @@ GDriveClient, ) from syft.sync.connections.drive.gdrive_retry import ( - batch_execute_with_retries, execute_with_retries, raise_unavailable, ) @@ -529,13 +528,13 @@ def create_peer_datasite_folders(self, peer_email: str): 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) + self.client.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) + self.client.add_permission(outbox_id, peer_email, write=True) def add_peer(self, peer_email: str): """Alias for create_peer_datasite_folders.""" @@ -637,7 +636,7 @@ def get_peer_requests(self) -> list[str]: peers who created folders for our datasite. Filters out already-accepted or rejected peers from SYFT_peers.json. """ - folders = self._list_all_files( + folders = self.client.list_all_files( f"name contains '{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#' " f"and name contains '#{self.email}#' " f"and trashed=false " @@ -869,35 +868,6 @@ def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: ) ) - def _has_permission(self, file_id: str, email: str) -> bool: - """Check if user already has permission on the file.""" - perms = execute_with_retries( - self.client.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.client.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"): @@ -920,7 +890,7 @@ 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( + folders = self.client.find_folders( name_contains=[f"#{self.email}"], parent_id=self.get_syftbox_folder_id(), owner_email=self.email, @@ -1125,7 +1095,7 @@ def _find_p2p_folder_id( peer_email: str, owner_email: str, ) -> str | None: - folders = self._find_folders( + folders = self.client.find_folders( name_contains=[ f"{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#", f"#{datasite_email}#{folder_type}#{peer_email}", @@ -1311,94 +1281,7 @@ def find_orphaned_syft_files(self) -> list[str]: 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.client.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 _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.client.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.client.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) - - out: list[tuple[str, str]] = [] - page_token = None - while True: - page = execute_with_retries( - self.client.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 + return [f["id"] for f in self.client.list_all_files(query, fields="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. @@ -1508,7 +1391,9 @@ def owner_create_collection_folder( 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) + folder_id = self.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 @@ -1537,34 +1422,7 @@ def owner_share_collection( 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 - - BATCH_SIZE = 100 - for i in range(0, len(users), BATCH_SIZE): - chunk = users[i : i + BATCH_SIZE] - batch = self.client.drive_service.new_batch_http_request(callback=callback) - for user_email in chunk: - permission = { - "type": "user", - "role": "reader", - "emailAddress": user_email, - } - batch.add( - self.client.drive_service.permissions().create( - fileId=file_id, - body=permission, - sendNotificationEmail=True, - ) - ) - batch_execute_with_retries(batch) + self.client.batch_add_permissions(folder_id, users) def owner_upload_collection_files( self, prefix: str, tag: str, content_hash: str, files: dict[str, bytes] @@ -1590,7 +1448,7 @@ def owner_list_collections(self, prefix: str) -> list[str]: f"and 'me' in owners and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" ) result: list[str] = [] - for folder in self._list_all_files(query, fields="name"): + for folder in self.client.list_all_files(query, fields="name"): try: cf = CollectionFolder.from_name(prefix, folder["name"]) except ValueError: @@ -1610,7 +1468,7 @@ def owner_list_all_collections_with_permissions( f"and 'me' in owners and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" ) collections = [] - for folder in self._list_all_files(query, fields="id,name,appProperties"): + for folder in self.client.list_all_files(query, fields="id,name,appProperties"): try: cf = CollectionFolder.from_name(prefix, folder["name"]) except ValueError: @@ -1654,7 +1512,7 @@ def watcher_list_collections(self, prefix: str) -> list[dict]: f"and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" ) result = [] - for folder in self._list_all_files(query, fields="name, owners"): + for folder in self.client.list_all_files(query, fields="name, owners"): try: cf = CollectionFolder.from_name(prefix, folder["name"]) except ValueError: @@ -1679,7 +1537,9 @@ def watcher_download_collection( 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) + folder_id = self.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") @@ -1700,7 +1560,9 @@ def watcher_get_collection_file_metadatas( 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) + folder_id = self.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") @@ -1723,7 +1585,9 @@ def _get_collection_folder_id( 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) + folder_id = self.client.find_folder_by_name( + folder_name, parent_id=syftbox_folder_id + ) if not folder_id: raise ValueError( @@ -1839,7 +1703,7 @@ def share_peer_info_file_with_peer( file_id = self._get_own_peer_info_file_id() if file_id: - self.add_permission(file_id, peer_email, write=False) + self.client.add_permission(file_id, peer_email, write=False) # ========================================================================= # CHECKPOINT METHODS @@ -1851,7 +1715,7 @@ def _get_checkpoints_folder_name(self) -> str: def _get_checkpoints_folder_id(self) -> str | None: """Find the checkpoints folder ID from Google Drive.""" - folders = self._find_folders( + folders = self.client.find_folders( name_contains=[f"{self.email}-", "-checkpoints"], parent_id=self.get_syftbox_folder_id(), ) @@ -2138,7 +2002,7 @@ def _get_rolling_state_folder_id(self, use_cache: bool = True) -> str | None: if use_cache and self._rolling_state_folder_id is not None: return self._rolling_state_folder_id - folders = self._find_folders( + folders = self.client.find_folders( name_contains=[f"{self.email}-", "-rolling-state"], parent_id=self.get_syftbox_folder_id(), ) @@ -2273,7 +2137,7 @@ def _find_encryption_bundles_folder_id(self) -> str | None: 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_id = self.client.find_folder_by_name( folder_name, parent_id=self.get_syftbox_folder_id() ) if folder_id: @@ -2320,7 +2184,7 @@ def write_encryption_bundle(self, peer_email: str, bundle_json: str) -> None: 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) + self.client.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. diff --git a/tests/unit/test_p2p_folder_lookup.py b/tests/unit/test_p2p_folder_lookup.py index 3b3fb9b2a2b..14468abcd7c 100644 --- a/tests/unit/test_p2p_folder_lookup.py +++ b/tests/unit/test_p2p_folder_lookup.py @@ -25,7 +25,7 @@ def _name(version: str, datasite: str, folder_type: str, peer: str) -> str: def _conn(found): conn = GDriveConnection(email=ME, verbose=False) conn.client = GDriveClient(Mock()) - conn._find_folders = Mock(return_value=found) + conn.client.find_folders = Mock(return_value=found) return conn diff --git a/tests/unit/test_versioned_folder_adopt.py b/tests/unit/test_versioned_folder_adopt.py index 1837ad848c4..3b172710332 100644 --- a/tests/unit/test_versioned_folder_adopt.py +++ b/tests/unit/test_versioned_folder_adopt.py @@ -177,7 +177,7 @@ def _checkpoints_owner_with_version(checkpoint_version: int): def _checkpoints_folder_names(connection) -> list[str]: return [ name - for _, name in connection._find_folders( + for _, name in connection.client.find_folders( name_contains=[f"{connection.email}-", "-checkpoints"], parent_id=connection.get_syftbox_folder_id(), ) From a3c0fd50ca5acb085d46dac33104a2b78f627ad3 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 14:18:54 -0300 Subject: [PATCH 10/28] refactor: move the syft folder layout into DriveLayout --- packages/syft-rds/tests/test_sync_manager.py | 2 +- .../tests/test_version_mismatch_flow.py | 18 +- syft/sync/connections/connection_router.py | 8 +- syft/sync/connections/drive/gdrive_layout.py | 466 +++++++++++++++ .../connections/drive/gdrive_transport.py | 534 ++---------------- tests/unit/test_dataset_collection_listing.py | 5 +- tests/unit/test_delete_syftbox.py | 8 +- tests/unit/test_encryption.py | 2 +- tests/unit/test_gdrive_client.py | 4 +- tests/unit/test_gdrive_pagination.py | 7 +- tests/unit/test_mock_drive_service.py | 2 +- tests/unit/test_outbox_compaction.py | 2 +- tests/unit/test_p2p_folder_lookup.py | 5 +- tests/unit/test_peer_request_owner.py | 2 +- tests/unit/test_peers_json_version.py | 7 +- tests/unit/test_versioned_folder_adopt.py | 20 +- tests/unit/test_versioned_folder_lookup.py | 2 +- 17 files changed, 568 insertions(+), 526 deletions(-) create mode 100644 syft/sync/connections/drive/gdrive_layout.py diff --git a/packages/syft-rds/tests/test_sync_manager.py b/packages/syft-rds/tests/test_sync_manager.py index c5bc03fef0f..ccd6adba4fc 100644 --- a/packages/syft-rds/tests/test_sync_manager.py +++ b/packages/syft-rds/tests/test_sync_manager.py @@ -283,7 +283,7 @@ def test_jobs(): # 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 = do_manager.sync_engine._connection_router.connections[0] - inbox_folder_id = connection_do._get_own_datasite_inbox_id(ds_manager.email) + inbox_folder_id = connection_do.layout.get_own_datasite_inbox_id(ds_manager.email) inbox_file_metadatas = connection_do.get_file_metadatas_from_folder(inbox_folder_id) assert len(inbox_file_metadatas) == 1 diff --git a/packages/syft-rds/tests/test_version_mismatch_flow.py b/packages/syft-rds/tests/test_version_mismatch_flow.py index ab0d2697ca6..3a1e23e331e 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -8,10 +8,8 @@ from unittest.mock import patch from syft.sync.connections.drive.gdrive_client import GOOGLE_FOLDER_MIME_TYPE -from syft.sync.connections.drive.gdrive_transport import ( - GDRIVE_P2P_FOLDER_DATASITE_PREFIX, - GDriveConnection, -) +from syft.sync.connections.drive.gdrive_layout import GDRIVE_P2P_FOLDER_DATASITE_PREFIX +from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.connections.drive.mock_drive_service import ( MockDriveService, ) @@ -182,7 +180,7 @@ def test_version_mismatch_continues_and_repairs(): old_personal_name = f"{SYFT_VERSION}#{do_manager.email}" old_personal_id = do_conn.client.find_folder_by_name( old_personal_name, - parent_id=do_conn.get_syftbox_folder_id(), + parent_id=do_conn.layout.get_syftbox_folder_id(), owner_email=do_manager.email, ) assert old_personal_id is not None @@ -199,6 +197,10 @@ def test_version_mismatch_continues_and_repairs(): "syft.sync.connections.drive.gdrive_transport.SYFT_VERSION", NEW_VERSION, ), + patch( + "syft.sync.connections.drive.gdrive_layout.SYFT_VERSION", + NEW_VERSION, + ), patch("syft.sync.login_utils.SYFT_VERSION", NEW_VERSION), patch("syft.sync.version.version_info.SYFT_VERSION", NEW_VERSION), ): @@ -217,7 +219,7 @@ def test_version_mismatch_continues_and_repairs(): new_personal_name = f"{NEW_VERSION}#{do_email}" new_personal_id = do_conn_new.client.find_folder_by_name( new_personal_name, - parent_id=do_conn_new.get_syftbox_folder_id(), + parent_id=do_conn_new.layout.get_syftbox_folder_id(), owner_email=do_email, ) assert new_personal_id is not None @@ -225,7 +227,7 @@ def test_version_mismatch_continues_and_repairs(): assert ( do_conn_new.client.find_folder_by_name( old_personal_name, - parent_id=do_conn_new.get_syftbox_folder_id(), + parent_id=do_conn_new.layout.get_syftbox_folder_id(), owner_email=do_email, ) is None @@ -253,7 +255,7 @@ def test_version_mismatch_continues_and_repairs(): ds_personal_name = f"{NEW_VERSION}#{ds_email}" ds_personal_id = ds_conn_new.client.find_folder_by_name( ds_personal_name, - parent_id=ds_conn_new.get_syftbox_folder_id(), + parent_id=ds_conn_new.layout.get_syftbox_folder_id(), owner_email=ds_email, ) assert ds_personal_id is not None diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index 3a3b8f8c03b..f6d05760c89 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -250,13 +250,7 @@ def download_outbox_message(self, file_id: str) -> bytes: def owner_list_outbox_files(self, recipient_email: str) -> List[dict] | None: """The files in my outbox for ``recipient_email``, or None when that outbox does not exist.""" - connection = self.connection_for_outbox() - folder_id = connection._get_own_datasite_outbox_id(recipient_email) - if folder_id is None: - return None - return connection.get_file_metadatas_from_folder( - folder_id, since_timestamp=None - ) + return self.connection_for_outbox().owner_list_outbox_files(recipient_email) # ========================================================================= # EVENT LOG (self-encrypted at rest when encryption enabled) diff --git a/syft/sync/connections/drive/gdrive_layout.py b/syft/sync/connections/drive/gdrive_layout.py new file mode 100644 index 00000000000..6a3a6c97983 --- /dev/null +++ b/syft/sync/connections/drive/gdrive_layout.py @@ -0,0 +1,466 @@ +"""Where syft keeps its folders in one account's Google Drive. + +The layout 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 pydantic import BaseModel, field_validator + +from syft.sync.connections.drive.gdrive_client import ( + GOOGLE_FOLDER_MIME_TYPE, + GDriveClient, +) +from syft.sync.connections.drive.gdrive_retry import execute_with_retries +from syft.sync.version.version_info import _parse_semver +from syft.version import SYFT_VERSION + +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) + + +@dataclass +class DriveLayout: + 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) + + def copy_caches_to(self, other: "DriveLayout") -> None: + """Seed the folder ids of ``other`` from this layout.""" + 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 in self.archive_folder_id_cache: + return self.archive_folder_id_cache[sender_email] + else: + archive_folder_id = self.get_archive_folder_id_from_drive(sender_email) + if archive_folder_id: + self.archive_folder_id_cache[sender_email] = archive_folder_id + return archive_folder_id + else: + return self.create_archive_folder(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) diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 889a86a5e73..a0f0b11d22f 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -11,7 +11,7 @@ from google_auth_httplib2 import AuthorizedHttp from googleapiclient.discovery import build from googleapiclient.http import build_http -from pydantic import BaseModel, field_validator +from pydantic import BaseModel from syft_migration import MigrationError from syft.sync.checkpoints.checkpoint import ( @@ -35,6 +35,12 @@ SYFT_FILE_PROPERTY_VALUE, GDriveClient, ) +from syft.sync.connections.drive.gdrive_layout import ( + GDRIVE_P2P_FOLDER_DATASITE_PREFIX, + SYFTBOX_FOLDER, + DriveLayout, + GdriveP2PFolder, +) from syft.sync.connections.drive.gdrive_retry import ( execute_with_retries, raise_unavailable, @@ -58,7 +64,6 @@ 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: @@ -70,7 +75,6 @@ # Timeout for Google API requests (in seconds) GOOGLE_API_TIMEOUT = 120 # 2 minutes -SYFTBOX_FOLDER = "SyftBox" SCOPES = ["https://www.googleapis.com/auth/drive"] @@ -101,52 +105,11 @@ def build_drive_service( 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" -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""" @@ -156,15 +119,6 @@ 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. @@ -233,125 +187,19 @@ def compute_hash(files: dict[str, bytes]) -> str: 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 - client: GDriveClient | None = None + layout: DriveLayout | None = 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] = {} @@ -433,12 +281,16 @@ def setup( drive_service = build_drive_service( self.credentials, environment=self.environment ) - self.client = GDriveClient(drive_service) + self.layout = DriveLayout(GDriveClient(drive_service), self.email) if warm_syftbox_folder_id_cache: - self.get_personal_syftbox_folder_id() + self.layout.get_personal_syftbox_folder_id() self._is_setup = True + @property + def client(self) -> GDriveClient: + return self.layout.client + def as_transport(self): """This connection as the transport jobs. See ``drive_connections.py``.""" return drive_transport(self) @@ -467,19 +319,13 @@ def copy(self) -> "GDriveConnection": 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 + self.layout.copy_caches_to(other.layout) 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 ) @@ -496,27 +342,6 @@ def get_authenticated_email(self) -> str: ) 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.client.create_folder(folder_name, syftbox_folder_id) - - 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 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 create_peer_datasite_folders(self, peer_email: str): """Create inbox + outbox folders for peer's datasite on own drive, share with peer. @@ -525,15 +350,15 @@ def create_peer_datasite_folders(self, peer_email: str): - 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) + inbox_id = self.layout.get_peer_datasite_inbox_id(peer_email) if inbox_id is None: - inbox_id = self._create_peer_datasite_folder(peer_email, "inbox") + inbox_id = self.layout.create_peer_datasite_folder(peer_email, "inbox") self.client.add_permission(inbox_id, peer_email, write=True) # Outbox: peer writes events, I read - outbox_id = self._get_peer_datasite_outbox_id(peer_email) + outbox_id = self.layout.get_peer_datasite_outbox_id(peer_email) if outbox_id is None: - outbox_id = self._create_peer_datasite_folder(peer_email, "outbox") + outbox_id = self.layout.create_peer_datasite_folder(peer_email, "outbox") self.client.add_permission(outbox_id, peer_email, write=True) def add_peer(self, peer_email: str): @@ -542,7 +367,7 @@ def add_peer(self, peer_email: str): def _get_peers_file_id(self) -> str | None: """Find SYFT_peers.json file in /SyftBox folder""" - syftbox_folder_id = self.get_syftbox_folder_id() + syftbox_folder_id = self.layout.get_syftbox_folder_id() query = f"name='{SYFT_PEERS_FILE}' and '{syftbox_folder_id}' in parents and trashed=false" results = execute_with_retries( self.client.drive_service.files().list(q=query, fields="files(id)") @@ -585,7 +410,7 @@ def _write_peers_json(self, peers_data: dict[str, dict[str, Any]]): **peers_data, PEERS_META_KEY: {"version": SYFT_PEERS_VERSION}, } - syftbox_folder_id = self.get_syftbox_folder_id() + syftbox_folder_id = self.layout.get_syftbox_folder_id() file_id = self._get_peers_file_id() # Convert to JSON bytes @@ -680,7 +505,7 @@ def get_peer_requests(self) -> list[str]: 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) + folder_id = self.layout.get_peer_datasite_outbox_id(peer_email) if folder_id is None: return [] @@ -718,7 +543,7 @@ 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) + folder_id = self.layout.get_peer_datasite_outbox_id(peer_email) if folder_id is None: return [] @@ -744,7 +569,7 @@ def watcher_get_outbox_file_metadatas( 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() + personal_syftbox_folder_id = self.layout.get_personal_syftbox_folder_id() file_metadata = { "name": filename, "parents": [personal_syftbox_folder_id], @@ -763,7 +588,7 @@ def owner_download_raw_bytes_by_id(self, file_id: str) -> bytes: 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() + personal_syftbox_folder_id = self.layout.get_personal_syftbox_folder_id() file_metadatas = self.get_file_metadatas_from_folder( personal_syftbox_folder_id, since_timestamp=since_timestamp ) @@ -779,7 +604,7 @@ def owner_get_all_accepted_event_file_ids( @raise_unavailable("Could not read the event log from Drive.") 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() + personal_syftbox_folder_id = self.layout.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) @@ -794,7 +619,7 @@ def owner_download_all_raw_events_from_syftbox(self) -> list[bytes]: 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) + outbox_folder_id = self.layout.get_own_datasite_outbox_id(recipient) if outbox_folder_id is None: raise ValueError(f"Outbox folder for {recipient} not found") @@ -857,7 +682,7 @@ def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: self.client.drive_service.files().get(fileId=file_id, fields="parents") ) previous_parents = ",".join(file_info.get("parents", [])) - archive_folder_id = self.owner_get_archive_folder_id(sender_email) + archive_folder_id = self.layout.get_archive_folder_id(sender_email) execute_with_retries( self.client.drive_service.files().update( fileId=file_id, @@ -868,82 +693,6 @@ def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: ) ) - 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 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 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.client.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 in self.archive_folder_id_cache: - return self.archive_folder_id_cache[sender_email] - else: - archive_folder_id = self.get_archive_folder_id_from_drive(sender_email) - if archive_folder_id: - self.archive_folder_id_cache[sender_email] = archive_folder_id - return archive_folder_id - else: - return self.create_archive_folder(sender_email) - def get_file_metadatas_from_folder( self, folder_id: str, @@ -1061,7 +810,7 @@ def _get_valid_messages_from_file_metadatas( def owner_download_next_raw_proposed_message_from_inbox( self, sender_email: str ) -> tuple[bytes, str] | None: - inbox_folder_id = self._get_own_datasite_inbox_id(sender_email) + inbox_folder_id = self.layout.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) @@ -1088,111 +837,18 @@ def owner_get_next_proposed_filechange_message( msg.platform_id = file_id return msg - 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: + def owner_list_outbox_files(self, recipient_email: str) -> list[dict] | None: + """The files in my outbox for ``recipient_email``, or None when that + outbox does not exist.""" + folder_id = self.layout.get_own_datasite_outbox_id(recipient_email) + if folder_id is None: 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 + return self.get_file_metadatas_from_folder(folder_id, since_timestamp=None) def watcher_send_raw_bytes_to_inbox( self, recipient: str, filename: str, data: bytes ) -> None: - inbox_id = self._get_peer_datasite_inbox_id(recipient) + inbox_id = self.layout.get_peer_datasite_inbox_id(recipient) if inbox_id is None: raise Exception(f"Inbox folder for {recipient}'s datasite not found") @@ -1229,13 +885,7 @@ def watcher_send_proposed_file_changes_message( 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.layout.reset_caches() self.personal_syftbox_event_id_cache.clear() self.collection_folder_id_cache.clear() self._rolling_state_folder_id = None @@ -1244,7 +894,7 @@ def reset_caches(self): self._peers_json_cache = None def gather_all_file_and_folder_ids(self) -> list[str]: - syftbox_folder_id = self.get_syftbox_folder_id() + syftbox_folder_id = self.layout.get_syftbox_folder_id() return gather_all_file_and_folder_ids_recursive( self.client.drive_service, syftbox_folder_id ) @@ -1283,84 +933,10 @@ def find_orphaned_syft_files(self) -> list[str]: ) return [f["id"] for f in self.client.list_all_files(query, fields="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 _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 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 get_inbox_proposed_event_id_from_name( self, sender_email: str, name: str ) -> str | None: - inbox_folder_id = self._get_own_datasite_inbox_id(sender_email) + inbox_folder_id = self.layout.get_own_datasite_inbox_id(sender_email) query = f"name='{name}' and '{inbox_folder_id}' in parents and trashed=false" results = execute_with_retries( self.client.drive_service.files().list(q=query, fields="files(id, name)") @@ -1388,7 +964,7 @@ def owner_create_collection_folder( 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() + syftbox_folder_id = self.layout.get_syftbox_folder_id() # Check if exists folder_id = self.client.find_folder_by_name( @@ -1442,7 +1018,7 @@ def owner_list_collections(self, prefix: str) -> list[str]: A tag appears once however many layouts hold it. """ - syftbox_folder_id = self.get_syftbox_folder_id() + syftbox_folder_id = self.layout.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}'" @@ -1462,7 +1038,7 @@ def owner_list_all_collections_with_permissions( prefix: str, ) -> list[FileCollection]: """List all DO's collections with permissions info, one per layout.""" - syftbox_folder_id = self.get_syftbox_folder_id() + syftbox_folder_id = self.layout.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}'" @@ -1584,7 +1160,7 @@ def _get_collection_folder_id( 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() + syftbox_folder_id = self.layout.get_syftbox_folder_id() folder_id = self.client.find_folder_by_name( folder_name, parent_id=syftbox_folder_id ) @@ -1599,7 +1175,7 @@ def _get_collection_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() + syftbox_folder_id = self.layout.get_syftbox_folder_id() query = f"name='{SYFT_VERSION_FILE}' and '{syftbox_folder_id}' in parents and trashed=false" results = execute_with_retries( self.client.drive_service.files().list(q=query, fields="files(id)") @@ -1610,7 +1186,7 @@ def _get_own_peer_info_file_id(self) -> Optional[str]: 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() + syftbox_folder_id = self.layout.get_syftbox_folder_id() file_id = self._get_own_peer_info_file_id() # Convert to JSON string @@ -1717,9 +1293,9 @@ def _get_checkpoints_folder_id(self) -> str | None: """Find the checkpoints folder ID from Google Drive.""" folders = self.client.find_folders( name_contains=[f"{self.email}-", "-checkpoints"], - parent_id=self.get_syftbox_folder_id(), + parent_id=self.layout.get_syftbox_folder_id(), ) - return self._find_or_adopt_versioned_folder( + return self.layout.find_or_adopt_versioned_folder( folders, current_name=self._get_checkpoints_folder_name() ) @@ -1730,7 +1306,7 @@ def _get_or_create_checkpoints_folder_id(self) -> str: return folder_id # Create the folder folder_name = self._get_checkpoints_folder_name() - syftbox_folder_id = self.get_syftbox_folder_id() + syftbox_folder_id = self.layout.get_syftbox_folder_id() return self.client.create_folder(folder_name, syftbox_folder_id) def upload_raw_checkpoint(self, filename: str, data: bytes) -> str: @@ -1936,7 +1512,7 @@ def get_events_count_since_checkpoint( Returns: Number of events since the checkpoint. """ - personal_folder_id = self.get_personal_syftbox_folder_id() + personal_folder_id = self.layout.get_personal_syftbox_folder_id() file_metadatas = self.get_file_metadatas_from_folder(personal_folder_id) if checkpoint_timestamp is None: @@ -1964,7 +1540,7 @@ 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() + personal_folder_id = self.layout.get_personal_syftbox_folder_id() file_metadatas = self.get_file_metadatas_from_folder( personal_folder_id, since_timestamp=since_timestamp ) @@ -2004,9 +1580,9 @@ def _get_rolling_state_folder_id(self, use_cache: bool = True) -> str | None: folders = self.client.find_folders( name_contains=[f"{self.email}-", "-rolling-state"], - parent_id=self.get_syftbox_folder_id(), + parent_id=self.layout.get_syftbox_folder_id(), ) - folder_id = self._find_or_adopt_versioned_folder( + folder_id = self.layout.find_or_adopt_versioned_folder( folders, current_name=self._get_rolling_state_folder_name() ) if folder_id is not None: @@ -2020,7 +1596,7 @@ def _get_or_create_rolling_state_folder_id(self) -> str: return folder_id # Create the folder folder_name = self._get_rolling_state_folder_name() - syftbox_folder_id = self.get_syftbox_folder_id() + syftbox_folder_id = self.layout.get_syftbox_folder_id() folder_id = self.client.create_folder(folder_name, syftbox_folder_id) self._rolling_state_folder_id = folder_id return folder_id @@ -2138,7 +1714,7 @@ def _find_encryption_bundles_folder_id(self) -> str | None: return self._encryption_bundles_folder_id folder_name = GdriveEncryptionBundlesFolder(email=self.email).as_string() folder_id = self.client.find_folder_by_name( - folder_name, parent_id=self.get_syftbox_folder_id() + folder_name, parent_id=self.layout.get_syftbox_folder_id() ) if folder_id: self._encryption_bundles_folder_id = folder_id @@ -2150,7 +1726,7 @@ def _get_or_create_encryption_bundles_folder_id(self) -> str: if not folder_id: folder_name = GdriveEncryptionBundlesFolder(email=self.email).as_string() folder_id = self.client.create_folder( - folder_name, self.get_syftbox_folder_id() + folder_name, self.layout.get_syftbox_folder_id() ) self._encryption_bundles_folder_id = folder_id return folder_id diff --git a/tests/unit/test_dataset_collection_listing.py b/tests/unit/test_dataset_collection_listing.py index ce4f532bd77..1c421e0894b 100644 --- a/tests/unit/test_dataset_collection_listing.py +++ b/tests/unit/test_dataset_collection_listing.py @@ -9,6 +9,7 @@ import pytest +from syft.sync.connections.drive.gdrive_layout import DriveLayout from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX @@ -19,8 +20,8 @@ def _conn(files): conn = GDriveConnection(email="alice@example.com", verbose=False) - conn.client = GDriveClient(Mock()) - conn._syftbox_folder_id = "syftbox-id" + conn.layout = DriveLayout(GDriveClient(Mock()), conn.email) + conn.layout._syftbox_folder_id = "syftbox-id" conn.client.drive_service.files().list().execute.return_value = {"files": files} return conn diff --git a/tests/unit/test_delete_syftbox.py b/tests/unit/test_delete_syftbox.py index 324d4025074..18d76929913 100644 --- a/tests/unit/test_delete_syftbox.py +++ b/tests/unit/test_delete_syftbox.py @@ -172,9 +172,9 @@ def test_proposal_moved_to_root_by_do_delete_found(self): 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) + archive_id = do_conn.layout.get_archive_folder_id(ds_manager.email) do_conn.client.drive_service.files().delete( - fileId=do_conn.get_syftbox_folder_id() + fileId=do_conn.layout.get_syftbox_folder_id() ).execute() found = _connection(ds_manager).find_orphaned_syft_files() @@ -198,7 +198,7 @@ def test_syftbox_folder_not_found(self): found = conn.find_orphaned_syft_files() - assert conn.get_syftbox_folder_id() not in found + assert conn.layout.get_syftbox_folder_id() not in found def _connection(manager: SyftboxManager): @@ -214,6 +214,6 @@ def _archive_one_proposal(): 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) + archive_id = do_conn.layout.get_archive_folder_id(ds_manager.email) archived = [m["id"] for m in do_conn.get_file_metadatas_from_folder(archive_id)] return ds_manager, do_manager, archived diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index cb2488a6155..a7a28a0d875 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -623,7 +623,7 @@ def test_tampered_inbox_message_logs_error(caplog): # Tamper with raw bytes in the mock backing store do_conn = do_manager._connection_router.connections[0] - inbox_id = do_conn._get_own_datasite_inbox_id(ds_manager.email) + inbox_id = do_conn.layout.get_own_datasite_inbox_id(ds_manager.email) files = do_conn.get_file_metadatas_from_folder(inbox_id) assert len(files) > 0 diff --git a/tests/unit/test_gdrive_client.py b/tests/unit/test_gdrive_client.py index aa241634a8f..d113a27e48c 100644 --- a/tests/unit/test_gdrive_client.py +++ b/tests/unit/test_gdrive_client.py @@ -12,7 +12,9 @@ def test_batch_delete_raises_error_whose_text_contains_404(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - folder_id = do_conn.client.create_folder("folder", do_conn.get_syftbox_folder_id()) + folder_id = do_conn.client.create_folder( + "folder", do_conn.layout.get_syftbox_folder_id() + ) def forbidden(self) -> None: raise http_error(403, "forbidden", "Cannot delete file 1Xy404abc") diff --git a/tests/unit/test_gdrive_pagination.py b/tests/unit/test_gdrive_pagination.py index fdd23cc271d..702246bcebd 100644 --- a/tests/unit/test_gdrive_pagination.py +++ b/tests/unit/test_gdrive_pagination.py @@ -1,6 +1,7 @@ from unittest.mock import Mock, patch from uuid import uuid4 +from syft.sync.connections.drive.gdrive_layout import DriveLayout from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.connections.drive.mock_drive_service import ( @@ -15,7 +16,7 @@ def create_mock_connection(): conn = GDriveConnection(email="test@test.com", verbose=False) - conn.client = GDriveClient(Mock()) + conn.layout = DriveLayout(GDriveClient(Mock()), conn.email) return conn @@ -61,7 +62,7 @@ def test_empty_folder(): @patch( - "syft.sync.connections.drive.gdrive_transport.GDriveConnection._get_peer_datasite_outbox_id" + "syft.sync.connections.drive.gdrive_layout.DriveLayout.get_peer_datasite_outbox_id" ) @patch("syft.sync.connections.drive.gdrive_client.GDriveClient.download_file") def test_events_returned_in_chronological_order(mock_download_file, mock_get_folder_id): @@ -82,7 +83,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 diff --git a/tests/unit/test_mock_drive_service.py b/tests/unit/test_mock_drive_service.py index f7262a15a92..7f95baf9a65 100644 --- a/tests/unit/test_mock_drive_service.py +++ b/tests/unit/test_mock_drive_service.py @@ -448,7 +448,7 @@ def test_create_folder(self): mock_service = MockDriveService(store, "user@example.com") connection = GDriveConnection.from_service("user@example.com", mock_service) - syftbox_id = connection.create_syftbox_folder() + syftbox_id = connection.layout.create_syftbox_folder() folder_id = connection.client.create_folder("TestFolder", syftbox_id) assert store.files[syftbox_id].properties == {"syft": "true"} diff --git a/tests/unit/test_outbox_compaction.py b/tests/unit/test_outbox_compaction.py index cbc867d2663..df54372fde8 100644 --- a/tests/unit/test_outbox_compaction.py +++ b/tests/unit/test_outbox_compaction.py @@ -21,7 +21,7 @@ def test_compact_outboxes_if_needed_collapses_to_single_message(): ) conn_do = do_manager._connection_router.connections[0] - outbox_id = conn_do._get_own_datasite_outbox_id(ds_manager.email) + outbox_id = conn_do.layout.get_own_datasite_outbox_id(ds_manager.email) pre = conn_do.get_file_metadatas_from_folder(outbox_id, since_timestamp=None) assert len(pre) == n_messages diff --git a/tests/unit/test_p2p_folder_lookup.py b/tests/unit/test_p2p_folder_lookup.py index 14468abcd7c..d7da71e935a 100644 --- a/tests/unit/test_p2p_folder_lookup.py +++ b/tests/unit/test_p2p_folder_lookup.py @@ -11,6 +11,7 @@ from unittest.mock import Mock +from syft.sync.connections.drive.gdrive_layout import DriveLayout from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_transport import GDriveConnection @@ -24,13 +25,13 @@ def _name(version: str, datasite: str, folder_type: str, peer: str) -> str: def _conn(found): conn = GDriveConnection(email=ME, verbose=False) - conn.client = GDriveClient(Mock()) + conn.layout = DriveLayout(GDriveClient(Mock()), conn.email) conn.client.find_folders = Mock(return_value=found) return conn def _lookup(conn): - return conn._find_p2p_folder_id( + return conn.layout.find_p2p_folder_id( datasite_email=PEER, folder_type="inbox", peer_email=ME, owner_email=ME ) diff --git a/tests/unit/test_peer_request_owner.py b/tests/unit/test_peer_request_owner.py index d486d9ffaf4..da530fca800 100644 --- a/tests/unit/test_peer_request_owner.py +++ b/tests/unit/test_peer_request_owner.py @@ -6,7 +6,7 @@ import pytest -from syft.sync.connections.drive.gdrive_transport import ( +from syft.sync.connections.drive.gdrive_layout import ( GDRIVE_P2P_FOLDER_DATASITE_PREFIX, GdriveP2PFolder, ) diff --git a/tests/unit/test_peers_json_version.py b/tests/unit/test_peers_json_version.py index 1be702af753..c3ff26f2839 100644 --- a/tests/unit/test_peers_json_version.py +++ b/tests/unit/test_peers_json_version.py @@ -15,6 +15,7 @@ from unittest.mock import Mock, patch from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.gdrive_layout import DriveLayout from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.peers.peer import PeerState from syft.sync.transport.naming import PEERS_META_KEY, SYFT_PEERS_VERSION @@ -24,7 +25,7 @@ def _conn(peers_data): conn = GDriveConnection(email="alice@example.com", verbose=False) - conn.client = GDriveClient(Mock()) + conn.layout = DriveLayout(GDriveClient(Mock()), conn.email) conn._peers_json_cache = dict(peers_data) return conn @@ -41,9 +42,7 @@ def test_a_write_stamps_the_reserved_entry(): conn = _conn({PEER: {"state": "accepted"}}) with ( patch.object(GDriveConnection, "_get_peers_file_id", return_value="file-id"), - patch.object( - GDriveConnection, "get_syftbox_folder_id", return_value="folder-id" - ), + patch.object(DriveLayout, "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"}}) diff --git a/tests/unit/test_versioned_folder_adopt.py b/tests/unit/test_versioned_folder_adopt.py index 3b172710332..b15928f409a 100644 --- a/tests/unit/test_versioned_folder_adopt.py +++ b/tests/unit/test_versioned_folder_adopt.py @@ -18,11 +18,11 @@ import pytest -from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.checkpoints.checkpoint import CHECKPOINT_VERSION, Checkpoint +from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.gdrive_layout import DriveLayout, _partition_by_version from syft.sync.connections.drive.gdrive_transport import ( GDriveConnection, - _partition_by_version, execute_with_retries, ) from syft.sync.syftbox_manager import SyftboxManager @@ -32,7 +32,7 @@ def _conn(): conn = GDriveConnection(email=EMAIL, verbose=False) - conn.client = GDriveClient(Mock()) + conn.layout = DriveLayout(GDriveClient(Mock()), conn.email) return conn @@ -83,7 +83,7 @@ def test_partition_returns_empty_for_a_bad_current_version(): def test_a_compatible_folder_wins_and_nothing_is_renamed(): conn = _conn() folders = [("same", f"0.2.5#{EMAIL}"), ("old", f"0.1.9#{EMAIL}")] - got = conn._find_or_adopt_versioned_folder( + got = conn.layout.find_or_adopt_versioned_folder( folders, current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) assert got == "same" @@ -93,7 +93,7 @@ def test_a_compatible_folder_wins_and_nothing_is_renamed(): def test_an_older_folder_is_adopted_by_rename(): conn = _conn() folders = [("old", f"0.1.9#{EMAIL}")] - got = conn._find_or_adopt_versioned_folder( + got = conn.layout.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" @@ -107,7 +107,7 @@ def test_the_highest_older_folder_is_adopted(): ("v2", f"0.1.20#{EMAIL}"), ("v0", f"0.0.4#{EMAIL}"), ] - got = conn._find_or_adopt_versioned_folder( + got = conn.layout.find_or_adopt_versioned_folder( folders, current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) assert got == "v2" @@ -120,7 +120,7 @@ def test_a_newer_folder_stops_the_client(): conn = _conn() folders = [("new", f"0.3.0#{EMAIL}")] with pytest.raises(RuntimeError, match="0.3.0"): - conn._find_or_adopt_versioned_folder( + conn.layout.find_or_adopt_versioned_folder( folders, current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) assert _renames(conn) == [] @@ -128,7 +128,7 @@ def test_a_newer_folder_stops_the_client(): def test_no_folder_returns_none_so_the_caller_creates_one(): conn = _conn() - got = conn._find_or_adopt_versioned_folder( + got = conn.layout.find_or_adopt_versioned_folder( [], current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) assert got is None @@ -139,7 +139,7 @@ def test_two_compatible_folders_still_raise(): conn = _conn() folders = [("a", f"0.2.1#{EMAIL}"), ("b", f"0.2.2#{EMAIL}")] with pytest.raises(RuntimeError): - conn._find_or_adopt_versioned_folder( + conn.layout.find_or_adopt_versioned_folder( folders, current_name=f"0.2.7#{EMAIL}", current_version="0.2.7" ) @@ -179,7 +179,7 @@ def _checkpoints_folder_names(connection) -> list[str]: name for _, name in connection.client.find_folders( name_contains=[f"{connection.email}-", "-checkpoints"], - parent_id=connection.get_syftbox_folder_id(), + parent_id=connection.layout.get_syftbox_folder_id(), ) ] diff --git a/tests/unit/test_versioned_folder_lookup.py b/tests/unit/test_versioned_folder_lookup.py index b2a991f57fb..ff9c7e5eb19 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.gdrive_layout import ( _extract_version_from_name, _looks_like_version, _partition_by_version, From ca76082568afc4e0fa5548d1806d33c17baedd3d Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 14:34:44 -0300 Subject: [PATCH 11/28] refactor: move the owner state into DriveOwnerStore --- packages/syft-rds/tests/test_sync_manager.py | 4 +- .../tests/test_version_mismatch_flow.py | 2 +- syft/sync/connections/connection_router.py | 12 +- .../connections/drive/drive_connections.py | 59 +- .../connections/drive/drive_owner_store.py | 594 ++++++++++++++++ syft/sync/connections/drive/gdrive_layout.py | 87 +++ .../connections/drive/gdrive_transport.py | 639 +----------------- .../test_gdrive_pagination.py | 6 +- tests/unit/test_delete_syftbox.py | 4 +- tests/unit/test_encryption.py | 2 +- tests/unit/test_gdrive_pagination.py | 18 +- tests/unit/test_outbox_compaction.py | 6 +- tests/unit/test_versioned_folder_adopt.py | 4 +- 13 files changed, 739 insertions(+), 698 deletions(-) create mode 100644 syft/sync/connections/drive/drive_owner_store.py diff --git a/packages/syft-rds/tests/test_sync_manager.py b/packages/syft-rds/tests/test_sync_manager.py index ccd6adba4fc..8e08b06eee7 100644 --- a/packages/syft-rds/tests/test_sync_manager.py +++ b/packages/syft-rds/tests/test_sync_manager.py @@ -284,7 +284,9 @@ def test_jobs(): # queue for longer than a certain time we start pushing) connection_do = do_manager.sync_engine._connection_router.connections[0] inbox_folder_id = connection_do.layout.get_own_datasite_inbox_id(ds_manager.email) - inbox_file_metadatas = connection_do.get_file_metadatas_from_folder(inbox_folder_id) + inbox_file_metadatas = connection_do.layout.get_file_metadatas_from_folder( + inbox_folder_id + ) assert len(inbox_file_metadatas) == 1 do_manager.sync() diff --git a/packages/syft-rds/tests/test_version_mismatch_flow.py b/packages/syft-rds/tests/test_version_mismatch_flow.py index 3a1e23e331e..fb0b2a621ec 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -194,7 +194,7 @@ 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( diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index f6d05760c89..b6625e80353 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -87,15 +87,6 @@ def copy_connection( ) -> SyftboxPlatformConnection: return connection.clone_for_thread() - def connection_for_eventlog( - self, create_new: bool = False - ) -> SyftboxPlatformConnection: - existing_connection = self.connections[0] - if create_new: - return self.copy_connection(existing_connection) - else: - return existing_connection - def connection_for_datasite_watcher(self) -> SyftboxPlatformConnection: return self.connections[0] @@ -268,8 +259,7 @@ def download_events_message_by_id( self, events_message_id: str ) -> FileChangeEventsMessage: """Download and decrypt a single events message (thread-safe).""" - connection = self.connection_for_eventlog(create_new=True) - raw = connection.as_transport().owner_store.fetch(events_message_id) + raw = self.transport_for_thread().owner_store.fetch(events_message_id) raw = self.peer_store.decrypt_and_verify_for_self_if_needed(raw) return FileChangeEventsMessage.from_compressed_data(raw) diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 8cd57a525a7..e514c5d14dc 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -12,10 +12,10 @@ from typing import TYPE_CHECKING +from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform from syft.sync.transport.protocols import ( Collections, - OwnerStore, P2PConnection, Peering, Transport, @@ -73,61 +73,6 @@ def ack(self, message: RawMessage) -> None: ) -class DriveOwnerStore(OwnerStore): - """The data owner's own folders in /SyftBox: event log, checkpoints, - rolling state.""" - - def __init__(self, connection: GDriveConnection): - self._conn = connection - - 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() - - class DriveWatcherStore(WatcherStore): """The data scientist's storage on Drive. Not built yet: today the data scientist keeps what it received in a local cache.""" @@ -258,7 +203,7 @@ def drive_transport(connection: GDriveConnection) -> Transport: """The transport jobs over one Drive connection.""" return Transport( p2p=DriveP2PConnection(connection), - owner_store=DriveOwnerStore(connection), + owner_store=DriveOwnerStore(connection.layout), watcher_store=DriveWatcherStore(connection), collections=DriveCollections(connection), peering=DrivePeering(connection), 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..3a09ad2d8e8 --- /dev/null +++ b/syft/sync/connections/drive/drive_owner_store.py @@ -0,0 +1,594 @@ +"""The data owner'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_layout import ( + DriveLayout, + valid_events_from_file_metadatas, +) +from syft.sync.connections.drive.gdrive_retry import ( + execute_with_retries, + raise_unavailable, +) +from syft.sync.events.file_change_event import FileChangeEventsMessageFileName +from syft.sync.transport.naming import timestamp_from_message_name +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 data owner's own folders in /SyftBox: event log, checkpoints, + rolling state.""" + + def __init__(self, layout: DriveLayout): + self._layout = layout + # 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 copy_caches_to(self, other: "DriveOwnerStore") -> None: + """Seed the caches of ``other`` from this store.""" + other._rolling_state_folder_id = self._rolling_state_folder_id + other._rolling_state_file_id = self._rolling_state_file_id + other.personal_syftbox_event_id_cache = dict( + self.personal_syftbox_event_id_cache + ) + + def reset_caches(self) -> None: + self._rolling_state_folder_id = None + self._rolling_state_file_id = None + self.personal_syftbox_event_id_cache.clear() + + # -- OwnerStore ------------------------------------------------------------ + + def append_event(self, name: str, data: bytes) -> str: + return self.owner_write_raw_bytes_to_syftbox(name, data) + + def list_event_handles(self, since_timestamp: float | None = None) -> list[str]: + return self.owner_get_all_accepted_event_file_ids(since_timestamp) + + def fetch(self, handle: str) -> bytes: + return self.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.download_raw_events_since_timestamp(since_timestamp) + + def fetch_all_events(self) -> list[bytes]: + return self.owner_download_all_raw_events_from_syftbox() + + def put_checkpoint(self, name: str, data: bytes) -> str: + return self.upload_raw_checkpoint(name, data) + + def latest_checkpoint(self) -> bytes | None: + return self.download_raw_latest_checkpoint() + + def put_incremental(self, name: str, data: bytes) -> str: + return self.upload_raw_incremental_checkpoint(name, data) + + def all_incrementals(self) -> list[bytes]: + return self.download_all_raw_incremental_checkpoints() + + def incremental_count(self) -> int: + return self.get_incremental_checkpoint_count() + + def next_incremental_sequence(self) -> int: + return self.get_next_incremental_sequence_number() + + def delete_incrementals(self) -> None: + self.delete_all_incremental_checkpoints() + + def put_rolling_state(self, name: str, data: bytes) -> str: + return self.upload_raw_rolling_state(name, data) + + def get_rolling_state(self) -> bytes | None: + return self.download_raw_rolling_state() + + # -- Drive ----------------------------------------------------------------- + + def owner_write_raw_bytes_to_syftbox(self, filename: str, data: bytes) -> str: + """Write raw bytes to /SyftBox/myemail.""" + personal_syftbox_folder_id = self._layout.get_personal_syftbox_folder_id() + file_metadata = { + "name": filename, + "parents": [personal_syftbox_folder_id], + } + file_payload, _ = self._layout.client.create_file_payload(data) + + res = self._layout.client.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._layout.client.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._layout.get_personal_syftbox_folder_id() + file_metadatas = self._layout.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 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._layout.get_personal_syftbox_folder_id() + file_metadatas = self._layout.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._layout.client.download_file(gdrive_id)) + return result + + @raise_unavailable("Could not read the event log from Drive.") + 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._layout.get_personal_syftbox_folder_id() + file_metadatas = self._layout.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._layout.client.download_file(gdrive_id)) + + return result + + 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._layout.get_personal_syftbox_folder_id() + file_metadatas = self._layout.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 = timestamp_from_message_name(f["name"]) + if event_timestamp is not None and event_timestamp > checkpoint_timestamp: + count += 1 + return count + + def _get_checkpoints_folder_name(self) -> str: + """Get the checkpoints folder name: {email}-{version}-checkpoints""" + return f"{self._layout.email}-{SYFT_VERSION}-checkpoints" + + def _get_checkpoints_folder_id(self) -> str | None: + """Find the checkpoints folder ID from Google Drive.""" + folders = self._layout.client.find_folders( + name_contains=[f"{self._layout.email}-", "-checkpoints"], + parent_id=self._layout.get_syftbox_folder_id(), + ) + return self._layout.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._layout.get_syftbox_folder_id() + return self._layout.client.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._layout.client.create_file_payload(data) + + file_metadata = { + "name": filename, + "parents": [folder_id], + } + + result = self._layout.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 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 = execute_with_retries( + self._layout.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._layout.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._layout.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._layout.client.drive_service.files().delete(fileId=item["id"]) + ) + except Exception as e: + print(f"Warning: Failed to delete checkpoint {item['name']}: {e}") + + 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._layout.client.create_file_payload(data) + + file_metadata = { + "name": filename, + "parents": [folder_id], + } + + result = self._layout.client.create_syft_file(file_metadata, payload) + return result.get("id") + + @raise_unavailable("Could not read the incremental checkpoints from Drive.") + 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 = execute_with_retries( + self._layout.client.drive_service.files().list( + q=query, fields="files(id, name)" + ) + ) + items = results.get("files", []) + + if not items: + return [] + + return [self._layout.client.download_file(item["id"]) for item in items] + + 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 = execute_with_retries( + self._layout.client.drive_service.files().list(q=query, fields="files(id)") + ) + 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 = execute_with_retries( + self._layout.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_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 = execute_with_retries( + self._layout.client.drive_service.files().list( + q=query, fields="files(id, name)" + ) + ) + items = results.get("files", []) + + for item in items: + try: + execute_with_retries( + self._layout.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._layout.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._layout.client.find_folders( + name_contains=[f"{self._layout.email}-", "-rolling-state"], + parent_id=self._layout.get_syftbox_folder_id(), + ) + folder_id = self._layout.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._layout.get_syftbox_folder_id() + folder_id = self._layout.client.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._layout.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._layout.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": filename, + "parents": [folder_id], + } + + result = self._layout.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 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 = execute_with_retries( + self._layout.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._layout.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._layout.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._layout.client.drive_service.files().delete(fileId=item["id"]) + ) + except Exception as e: + print(f"Warning: Failed to delete rolling state {item['name']}: {e}") diff --git a/syft/sync/connections/drive/gdrive_layout.py b/syft/sync/connections/drive/gdrive_layout.py index 6a3a6c97983..f66cb090b76 100644 --- a/syft/sync/connections/drive/gdrive_layout.py +++ b/syft/sync/connections/drive/gdrive_layout.py @@ -14,6 +14,8 @@ GDriveClient, ) 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.version.version_info import _parse_semver from syft.version import SYFT_VERSION @@ -157,6 +159,21 @@ def _ordered(entries: list[_VersionedFolder]) -> list[tuple[str, str]]: 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 DriveLayout: client: GDriveClient @@ -464,3 +481,73 @@ def create_archive_folder(self, sender_email: str) -> str: 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/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index a0f0b11d22f..3dcae5911d2 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -14,21 +14,12 @@ from pydantic import BaseModel 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.drive_connections import drive_transport +from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore from syft.sync.connections.drive.gdrive_client import ( GOOGLE_FOLDER_MIME_TYPE, SYFT_FILE_PROPERTY_KEY, @@ -40,6 +31,7 @@ SYFTBOX_FOLDER, DriveLayout, GdriveP2PFolder, + valid_events_from_file_metadatas, ) from syft.sync.connections.drive.gdrive_retry import ( execute_with_retries, @@ -51,7 +43,6 @@ 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, @@ -61,10 +52,9 @@ from syft.sync.transport.naming import ( PEERS_META_KEY, SYFT_PEERS_VERSION, - timestamp_from_message_name, ) +from syft.sync.transport.protocols import Transport from syft.sync.utils.syftbox_utils import check_env -from syft.version import SYFT_VERSION if TYPE_CHECKING: from syft.sync.connections.drive.grdrive_config import ( @@ -199,17 +189,13 @@ class Config: email: str token_path: Path | None = None _is_setup: bool = False - - # fname -> gdrive id - personal_syftbox_event_id_cache: dict[str, str] = {} + # The transport jobs over this connection, built once so each job keeps + # one set of caches. + _transport: Transport | None = None # 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 @@ -291,9 +277,17 @@ def setup( def client(self) -> GDriveClient: return self.layout.client - def as_transport(self): + def as_transport(self) -> Transport: """This connection as the transport jobs. See ``drive_connections.py``.""" - return drive_transport(self) + if self._transport is None: + self._transport = drive_transport(self) + return self._transport + + @property + def _owner_store(self) -> DriveOwnerStore: + owner_store = self.as_transport().owner_store + assert isinstance(owner_store, DriveOwnerStore) + return owner_store def clone_for_thread(self) -> "GDriveConnection": """A connection safe to use from another thread: the Drive client is not.""" @@ -320,15 +314,11 @@ 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.""" self.layout.copy_caches_to(other.layout) - other._rolling_state_folder_id = self._rolling_state_folder_id - other._rolling_state_file_id = self._rolling_state_file_id + self._owner_store.copy_caches_to(other._owner_store) 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.personal_syftbox_event_id_cache = dict( - self.personal_syftbox_event_id_cache - ) other.collection_folder_id_cache = dict(self.collection_folder_id_cache) @property @@ -509,10 +499,10 @@ def watcher_download_raw_events_from_outbox( if folder_id is None: return [] - file_metadatas = self.get_file_metadatas_from_folder( + file_metadatas = self.layout.get_file_metadatas_from_folder( folder_id, since_timestamp=since_timestamp ) - valid_fname_objs = self._get_valid_events_from_file_metadatas(file_metadatas) + valid_fname_objs = valid_events_from_file_metadatas(file_metadatas) name_to_id = {f["name"]: f["id"] for f in file_metadatas} sorted_fname_objs = [ @@ -547,10 +537,10 @@ def watcher_get_outbox_file_metadatas( if folder_id is None: return [] - file_metadatas = self.get_file_metadatas_from_folder( + file_metadatas = self.layout.get_file_metadatas_from_folder( folder_id, since_timestamp=since_timestamp ) - valid_fname_objs = self._get_valid_events_from_file_metadatas(file_metadatas) + valid_fname_objs = valid_events_from_file_metadatas(file_metadatas) name_to_id = {f["name"]: f["id"] for f in file_metadatas} result = [] @@ -567,55 +557,6 @@ def watcher_get_outbox_file_metadatas( ) 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.layout.get_personal_syftbox_folder_id() - file_metadata = { - "name": filename, - "parents": [personal_syftbox_folder_id], - } - file_payload, _ = self.client.create_file_payload(data) - - res = self.client.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.client.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.layout.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] - - @raise_unavailable("Could not read the event log from Drive.") - 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.layout.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"] - result.append(self.client.download_file(gdrive_id)) - return result - def owner_write_raw_bytes_to_outbox( self, recipient: str, filename: str, data: bytes ) -> None: @@ -693,107 +634,6 @@ def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: ) ) - 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 - - @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], @@ -813,7 +653,7 @@ def owner_download_next_raw_proposed_message_from_inbox( inbox_folder_id = self.layout.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) + file_metadatas = self.layout.get_file_metadatas_from_folder(inbox_folder_id) valid_file_names = self._get_valid_messages_from_file_metadatas(file_metadatas) if len(valid_file_names) == 0: return None @@ -843,7 +683,9 @@ def owner_list_outbox_files(self, recipient_email: str) -> list[dict] | None: folder_id = self.layout.get_own_datasite_outbox_id(recipient_email) if folder_id is None: return None - return self.get_file_metadatas_from_folder(folder_id, since_timestamp=None) + return self.layout.get_file_metadatas_from_folder( + folder_id, since_timestamp=None + ) def watcher_send_raw_bytes_to_inbox( self, recipient: str, filename: str, data: bytes @@ -886,10 +728,8 @@ def watcher_send_proposed_file_changes_message( def reset_caches(self): self.layout.reset_caches() - self.personal_syftbox_event_id_cache.clear() + self._owner_store.reset_caches() 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 @@ -1120,7 +960,7 @@ def watcher_download_collection( 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) + file_metadatas = self.layout.get_file_metadatas_from_folder(folder_id) files = {} for file_meta in file_metadatas: file_id = file_meta["id"] @@ -1143,7 +983,7 @@ def watcher_get_collection_file_metadatas( 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) + file_metadatas = self.layout.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: @@ -1281,429 +1121,6 @@ def share_peer_info_file_with_peer( if file_id: self.client.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.client.find_folders( - name_contains=[f"{self.email}-", "-checkpoints"], - parent_id=self.layout.get_syftbox_folder_id(), - ) - return self.layout.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.layout.get_syftbox_folder_id() - return self.client.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.client.create_file_payload(data) - - file_metadata = { - "name": filename, - "parents": [folder_id], - } - - result = self.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 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 = execute_with_retries( - self.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.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.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.client.drive_service.files().delete(fileId=item["id"]) - ) - 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.client.create_file_payload(data) - - file_metadata = { - "name": filename, - "parents": [folder_id], - } - - result = self.client.create_syft_file(file_metadata, payload) - return result.get("id") - - @raise_unavailable("Could not read the incremental checkpoints from Drive.") - 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 = execute_with_retries( - self.client.drive_service.files().list(q=query, fields="files(id, name)") - ) - items = results.get("files", []) - - if not items: - return [] - - return [self.client.download_file(item["id"]) for item in items] - - 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 = execute_with_retries( - self.client.drive_service.files().list(q=query, fields="files(id)") - ) - 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 = execute_with_retries( - self.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_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 = execute_with_retries( - self.client.drive_service.files().list(q=query, fields="files(id, name)") - ) - items = results.get("files", []) - - for item in items: - try: - execute_with_retries( - self.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_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.layout.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 = timestamp_from_message_name(f["name"]) - if event_timestamp is not None and event_timestamp > checkpoint_timestamp: - count += 1 - return count - - @raise_unavailable("Could not read the event log from Drive.") - 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.layout.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"] - - result.append(self.client.download_file(gdrive_id)) - - 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.client.find_folders( - name_contains=[f"{self.email}-", "-rolling-state"], - parent_id=self.layout.get_syftbox_folder_id(), - ) - folder_id = self.layout.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.layout.get_syftbox_folder_id() - folder_id = self.client.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.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.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": filename, - "parents": [folder_id], - } - - result = self.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 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 = execute_with_retries( - self.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.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.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.client.drive_service.files().delete(fileId=item["id"]) - ) - except Exception as e: - print(f"Warning: Failed to delete rolling state {item['name']}: {e}") - # ========================================================================= # ENCRYPTION BUNDLE METHODS # ========================================================================= diff --git a/tests/integration/with_unit_coverage/test_gdrive_pagination.py b/tests/integration/with_unit_coverage/test_gdrive_pagination.py index 1f2a9ae1046..76430d3a21a 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.gdrive_layout import DriveLayout 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 = DriveLayout.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, + DriveLayout, "get_file_metadatas_from_folder", get_files_page_size_2, ): diff --git a/tests/unit/test_delete_syftbox.py b/tests/unit/test_delete_syftbox.py index 18d76929913..66007d5a579 100644 --- a/tests/unit/test_delete_syftbox.py +++ b/tests/unit/test_delete_syftbox.py @@ -215,5 +215,7 @@ def _archive_one_proposal(): do_manager.sync() do_conn = _connection(do_manager) archive_id = do_conn.layout.get_archive_folder_id(ds_manager.email) - archived = [m["id"] for m in do_conn.get_file_metadatas_from_folder(archive_id)] + archived = [ + m["id"] for m in do_conn.layout.get_file_metadatas_from_folder(archive_id) + ] return ds_manager, do_manager, archived diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index a7a28a0d875..c0350c5f52a 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -624,7 +624,7 @@ def test_tampered_inbox_message_logs_error(caplog): # Tamper with raw bytes in the mock backing store do_conn = do_manager._connection_router.connections[0] inbox_id = do_conn.layout.get_own_datasite_inbox_id(ds_manager.email) - files = do_conn.get_file_metadatas_from_folder(inbox_id) + files = do_conn.layout.get_file_metadatas_from_folder(inbox_id) assert len(files) > 0 file_id = files[0]["id"] diff --git a/tests/unit/test_gdrive_pagination.py b/tests/unit/test_gdrive_pagination.py index 702246bcebd..23c801669b2 100644 --- a/tests/unit/test_gdrive_pagination.py +++ b/tests/unit/test_gdrive_pagination.py @@ -28,7 +28,7 @@ def test_single_page_works(): "files": [{"name": "file1.txt", "id": "id1"}] } - files = conn.get_file_metadatas_from_folder("folder_id") + files = conn.layout.get_file_metadatas_from_folder("folder_id") assert len(files) == 1 assert files[0]["name"] == "file1.txt" @@ -44,7 +44,7 @@ def test_multiple_pages_pagination(): {"files": [{"name": "file3.txt", "id": "id3"}]}, ] - files = conn.get_file_metadatas_from_folder("folder_id") + files = conn.layout.get_file_metadatas_from_folder("folder_id") assert len(files) == 3 assert conn.client.drive_service.files().list().execute.call_count == 3 @@ -56,7 +56,7 @@ def test_empty_folder(): conn.client.drive_service.files().list().execute.return_value = {"files": []} - files = conn.get_file_metadatas_from_folder("folder_id") + files = conn.layout.get_file_metadatas_from_folder("folder_id") assert files == [] @@ -117,7 +117,7 @@ def test_orderby_name_desc_applied(): conn.client.drive_service.files().list().execute.return_value = {"files": []} - conn.get_file_metadatas_from_folder("folder_id") + conn.layout.get_file_metadatas_from_folder("folder_id") # Check that orderBy was passed to Drive API call_kwargs = conn.client.drive_service.files().list.call_args[1] @@ -162,7 +162,7 @@ def test_early_termination_stops_pagination(): }, ] - files = conn.get_file_metadatas_from_folder( + files = conn.layout.get_file_metadatas_from_folder( "folder_id", since_timestamp=1731506400.0 ) @@ -191,7 +191,7 @@ def test_no_early_termination_when_all_files_new(): }, ] - files = conn.get_file_metadatas_from_folder( + files = conn.layout.get_file_metadatas_from_folder( "folder_id", since_timestamp=1731506400.0 ) @@ -215,7 +215,9 @@ def test_no_early_termination_when_since_timestamp_none(): }, ] - files = conn.get_file_metadatas_from_folder("folder_id", since_timestamp=None) + files = conn.layout.get_file_metadatas_from_folder( + "folder_id", since_timestamp=None + ) # Should fetch all pages (no early termination) assert conn.client.drive_service.files().list().execute.call_count == 2 @@ -232,7 +234,7 @@ def test_newer_message_listed_after_older_one_in_same_second_not_skipped(): "files": [{"name": name, "id": name} for name in names] } - files = conn.get_file_metadatas_from_folder( + files = conn.layout.get_file_metadatas_from_folder( "folder_id", since_timestamp=1731506400.92 ) diff --git a/tests/unit/test_outbox_compaction.py b/tests/unit/test_outbox_compaction.py index df54372fde8..54b53b02a48 100644 --- a/tests/unit/test_outbox_compaction.py +++ b/tests/unit/test_outbox_compaction.py @@ -22,11 +22,13 @@ def test_compact_outboxes_if_needed_collapses_to_single_message(): conn_do = do_manager._connection_router.connections[0] outbox_id = conn_do.layout.get_own_datasite_outbox_id(ds_manager.email) - pre = conn_do.get_file_metadatas_from_folder(outbox_id, since_timestamp=None) + pre = conn_do.layout.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 = conn_do.layout.get_file_metadatas_from_folder( + outbox_id, since_timestamp=None + ) assert len(post) == 1 diff --git a/tests/unit/test_versioned_folder_adopt.py b/tests/unit/test_versioned_folder_adopt.py index b15928f409a..74a6e09b1c3 100644 --- a/tests/unit/test_versioned_folder_adopt.py +++ b/tests/unit/test_versioned_folder_adopt.py @@ -164,7 +164,7 @@ def _checkpoints_owner_with_version(checkpoint_version: int): Checkpoint(email=do_manager.email, version=checkpoint_version) ) - folder_id = connection._get_checkpoints_folder_id() + folder_id = connection.as_transport().owner_store._get_checkpoints_folder_id() older_name = f"{do_manager.email}-0.0.1-checkpoints" execute_with_retries( connection.client.drive_service.files().update( @@ -194,7 +194,7 @@ def test_an_adopted_folder_serves_a_checkpoint_from_an_earlier_client(): 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() + connection.as_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) From 5143e7abc44dedb0f895c191c38275d1073627e0 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 14:43:10 -0300 Subject: [PATCH 12/28] refactor: move the collections into DriveCollections --- packages/syft-rds/src/syft_rds/client.py | 4 +- .../tests/test_dataset_upload_private.py | 2 +- syft/sync/connections/base_connection.py | 93 ----- .../connections/drive/drive_collections.py | 373 ++++++++++++++++++ .../connections/drive/drive_connections.py | 59 +-- .../connections/drive/gdrive_transport.py | 316 +-------------- .../p2p/test_dataset_multicopy_delivery.py | 8 +- tests/unit/test_dataset_collection_listing.py | 20 +- tests/unit/test_drive_connections.py | 6 +- 9 files changed, 407 insertions(+), 474 deletions(-) create mode 100644 syft/sync/connections/drive/drive_collections.py diff --git a/packages/syft-rds/src/syft_rds/client.py b/packages/syft-rds/src/syft_rds/client.py index 0066807958d..c50725a2f11 100644 --- a/packages/syft-rds/src/syft_rds/client.py +++ b/packages/syft-rds/src/syft_rds/client.py @@ -811,9 +811,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 61b3fb717dc..236799f18d5 100644 --- a/packages/syft-rds/tests/test_dataset_upload_private.py +++ b/packages/syft-rds/tests/test_dataset_upload_private.py @@ -1,6 +1,6 @@ +from syft.sync.connections.drive.drive_collections import collection_name_query from syft.sync.connections.drive.gdrive_transport import ( GDriveConnection, - collection_name_query, ) from syft_rds import SyftRDSClient from syft_datasets.dataset_manager import ( diff --git a/syft/sync/connections/base_connection.py b/syft/sync/connections/base_connection.py index 54c5cc9eec1..0bb63725e05 100644 --- a/syft/sync/connections/base_connection.py +++ b/syft/sync/connections/base_connection.py @@ -51,63 +51,6 @@ def clone_for_thread(self) -> "SyftboxPlatformConnection": """ 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) # ========================================================================= @@ -132,39 +75,3 @@ 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() diff --git a/syft/sync/connections/drive/drive_collections.py b/syft/sync/connections/drive/drive_collections.py new file mode 100644 index 00000000000..7518e3473e8 --- /dev/null +++ b/syft/sync/connections/drive/drive_collections.py @@ -0,0 +1,373 @@ +"""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.gdrive_layout import DriveLayout +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, layout: DriveLayout): + self._layout = layout + # folder name -> id of my collection folder + self.collection_folder_id_cache: dict[str, str] = {} + + def copy_caches_to(self, other: "DriveCollections") -> None: + """Seed the caches of ``other`` from this store.""" + other.collection_folder_id_cache = dict(self.collection_folder_id_cache) + + def reset_caches(self) -> None: + self.collection_folder_id_cache.clear() + + # -- Collections ----------------------------------------------------------- + + def create(self, prefix: str, tag: str, content_hash: str, owner_email: str) -> str: + return self.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.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.owner_share_collection(prefix, tag, content_hash, users) + + def upload( + self, prefix: str, tag: str, content_hash: str, files: dict[str, bytes] + ) -> None: + self.owner_upload_collection_files(prefix, tag, content_hash, files) + + def list_mine_with_permissions(self, prefix: str) -> list[FileCollection]: + return self.owner_list_all_collections_with_permissions(prefix) + + def list_mine(self, prefix: str) -> list[str]: + return self.owner_list_collections(prefix) + + def delete(self, prefix: str, tag: str) -> None: + self.owner_delete_collection(prefix, tag) + + def list_shared_with_me(self, prefix: str) -> list[dict]: + return self.watcher_list_collections(prefix) + + def list_files( + self, prefix: str, tag: str, content_hash: str, owner_email: str + ) -> list[dict]: + return self.watcher_get_collection_file_metadatas( + prefix, tag, content_hash, owner_email + ) + + def fetch_file(self, handle: str) -> bytes: + return self.watcher_download_collection_file(handle) + + def download( + self, prefix: str, tag: str, content_hash: str, owner_email: str + ) -> dict[str, bytes]: + return self.watcher_download_collection(prefix, tag, content_hash, owner_email) + + # -- Drive ----------------------------------------------------------------- + + 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._layout.get_syftbox_folder_id() + + # Check if exists + folder_id = self._layout.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._layout.client.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._layout.client.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._layout.client.batch_add_permissions(folder_id, users) + + 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._layout.client.create_file_payload(content) + file_name = Path(file_path).name + + file_metadata = {"name": file_name, "parents": [folder_id]} + self._layout.client.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._layout.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._layout.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 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._layout.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._layout.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 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._layout.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 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}'" + ) + result = [] + for folder in self._layout.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 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._layout.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._layout.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._layout.client.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._layout.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._layout.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._layout.client.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._layout.get_syftbox_folder_id() + folder_id = self._layout.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 e514c5d14dc..96066cc841a 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -12,10 +12,10 @@ from typing import TYPE_CHECKING +from syft.sync.connections.drive.drive_collections import DriveCollections from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform from syft.sync.transport.protocols import ( - Collections, P2PConnection, Peering, Transport, @@ -24,7 +24,6 @@ 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 @@ -96,60 +95,6 @@ def reset(self, peer: str) -> None: raise NotImplementedError("Watcher storage on Drive is not built yet") -class DriveCollections(Collections): - """Collection folders in /SyftBox, shared by Drive permission.""" - - def __init__(self, connection: GDriveConnection): - self._conn = connection - - 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 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(Peering): """Identity from the Google account; channels, the peer list, key bundles and version files in /SyftBox.""" @@ -205,7 +150,7 @@ def drive_transport(connection: GDriveConnection) -> Transport: p2p=DriveP2PConnection(connection), owner_store=DriveOwnerStore(connection.layout), watcher_store=DriveWatcherStore(connection), - collections=DriveCollections(connection), + collections=DriveCollections(connection.layout), peering=DrivePeering(connection), platform=GdriveFilesPlatform(), ) diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 3dcae5911d2..e6c740d4e10 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -2,10 +2,8 @@ import json import logging -import re -from functools import lru_cache from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Callable, Optional from google.oauth2.credentials import Credentials as GoogleCredentials from google_auth_httplib2 import AuthorizedHttp @@ -15,9 +13,9 @@ from syft_migration import MigrationError from syft.sync.connections.base_connection import ( - FileCollection, SyftboxPlatformConnection, ) +from syft.sync.connections.drive.drive_collections import DriveCollections from syft.sync.connections.drive.drive_connections import drive_transport from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore from syft.sync.connections.drive.gdrive_client import ( @@ -109,74 +107,6 @@ def as_string(self) -> str: return f"syft_encryption_bundles#{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) - - class GDriveConnection(SyftboxPlatformConnection): """Google Drive Files API transport layer""" @@ -193,9 +123,6 @@ class Config: # one set of caches. _transport: Transport | None = None - # tag -> dataset collection folder id - collection_folder_id_cache: Dict[str, str] = {} - # Encryption bundles folder cache _encryption_bundles_folder_id: str | None = None @@ -289,6 +216,12 @@ def _owner_store(self) -> DriveOwnerStore: assert isinstance(owner_store, DriveOwnerStore) return owner_store + @property + def _collections(self) -> DriveCollections: + collections = self.as_transport().collections + assert isinstance(collections, DriveCollections) + return collections + def clone_for_thread(self) -> "GDriveConnection": """A connection safe to use from another thread: the Drive client is not.""" return self.copy() @@ -315,11 +248,11 @@ def _copy_caches_to(self, other: "GDriveConnection") -> None: copy doesn't re-discover folder ids via Drive search on first use.""" self.layout.copy_caches_to(other.layout) self._owner_store.copy_caches_to(other._owner_store) + self._collections.copy_caches_to(other._collections) 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.collection_folder_id_cache = dict(self.collection_folder_id_cache) @property def environment(self) -> Environment: @@ -729,7 +662,7 @@ def watcher_send_proposed_file_changes_message( def reset_caches(self): self.layout.reset_caches() self._owner_store.reset_caches() - self.collection_folder_id_cache.clear() + self._collections.reset_caches() self._encryption_bundles_folder_id = None self._peers_json_cache = None @@ -784,235 +717,6 @@ def get_inbox_proposed_event_id_from_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.layout.get_syftbox_folder_id() - - # Check if exists - folder_id = self.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.client.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.client.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.client.batch_add_permissions(folder_id, users) - - 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.client.create_file_payload(content) - file_name = Path(file_path).name - - file_metadata = {"name": file_name, "parents": [folder_id]} - self.client.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.layout.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.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 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.layout.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.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 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.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 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}'" - ) - result = [] - for folder in self.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 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.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.layout.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.client.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.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.layout.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.client.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.layout.get_syftbox_folder_id() - folder_id = self.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 - 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.layout.get_syftbox_folder_id() 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/unit/test_dataset_collection_listing.py b/tests/unit/test_dataset_collection_listing.py index 1c421e0894b..bb821a27588 100644 --- a/tests/unit/test_dataset_collection_listing.py +++ b/tests/unit/test_dataset_collection_listing.py @@ -8,11 +8,11 @@ from unittest.mock import Mock import pytest +from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX -from syft.sync.connections.drive.gdrive_layout import DriveLayout from syft.sync.connections.drive.gdrive_client import GDriveClient +from syft.sync.connections.drive.gdrive_layout import DriveLayout from syft.sync.connections.drive.gdrive_transport import GDriveConnection -from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX VALID = f"{DATASET_COLLECTION_PREFIX}_mytag_abc123" UNPARSEABLE = DATASET_COLLECTION_PREFIX @@ -28,7 +28,9 @@ def _conn(files): 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) + got = conn.as_transport().collections.owner_list_all_collections_with_permissions( + DATASET_COLLECTION_PREFIX + ) assert [(c.folder_id, c.tag, c.content_hash) for c in got] == [ ("f1", "mytag", "abc123") ] @@ -45,7 +47,9 @@ def test_the_any_permission_flag_comes_from_app_properties(): } ] ) - got = conn.owner_list_all_collections_with_permissions(DATASET_COLLECTION_PREFIX) + got = conn.as_transport().collections.owner_list_all_collections_with_permissions( + DATASET_COLLECTION_PREFIX + ) assert got[0].has_any_permission is True @@ -56,7 +60,9 @@ def test_a_name_the_client_cannot_parse_is_skipped(): {"id": "f1", "name": VALID, "appProperties": {}}, ] ) - got = conn.owner_list_all_collections_with_permissions(DATASET_COLLECTION_PREFIX) + got = conn.as_transport().collections.owner_list_all_collections_with_permissions( + DATASET_COLLECTION_PREFIX + ) assert [c.folder_id for c in got] == ["f1"] @@ -65,4 +71,6 @@ def test_a_missing_name_field_raises(): # without a message. conn = _conn([{"id": "f1", "appProperties": {}}]) with pytest.raises(KeyError): - conn.owner_list_all_collections_with_permissions(DATASET_COLLECTION_PREFIX) + conn.as_transport().collections.owner_list_all_collections_with_permissions( + DATASET_COLLECTION_PREFIX + ) diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index d406d9979e9..008dd1b285f 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -12,12 +12,9 @@ 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_collections import CollectionFolder from syft.sync.connections.drive.gdrive_client import GDriveClient -from syft.sync.connections.drive.gdrive_transport import ( - CollectionFolder, -) from syft.sync.connections.drive.mock_drive_service import ( http_error, pair_with_mock_service, @@ -28,6 +25,7 @@ from syft.sync.peers.peer_store import PeerStore 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, From 1ca1cd94055d34435f5e53354087eeb2b6054946 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 14:53:44 -0300 Subject: [PATCH 13/28] refactor: move peering into DrivePeering --- .../tests/test_version_mismatch_flow.py | 19 +- syft/sync/connections/base_connection.py | 4 - syft/sync/connections/connection_router.py | 15 +- .../connections/drive/drive_connections.py | 54 +- syft/sync/connections/drive/drive_peering.py | 487 ++++++++++++++++++ .../connections/drive/gdrive_transport.py | 425 +-------------- syft/sync/login_utils.py | 2 +- syft/sync/version/peer_manager.py | 13 +- .../with_unit_coverage/test_sync_manager.py | 4 +- tests/unit/test_drive_connections.py | 4 +- tests/unit/test_encryption_policy.py | 9 +- tests/unit/test_peers_json_version.py | 13 +- tests/unit/test_sync_manager.py | 2 +- tests/unit/test_version_negotiation.py | 2 +- 14 files changed, 540 insertions(+), 513 deletions(-) create mode 100644 syft/sync/connections/drive/drive_peering.py diff --git a/packages/syft-rds/tests/test_version_mismatch_flow.py b/packages/syft-rds/tests/test_version_mismatch_flow.py index fb0b2a621ec..42fe40d1100 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -97,7 +97,7 @@ def _simulate_continue_on_mismatch(manager, backing_store): mock_conn = GDriveConnection.from_service(email, mock_service) def read_remote(e, t): - return mock_conn.read_own_peer_info_file().version_info + return mock_conn.as_transport().peering.read_own_peer_info().version_info with ( patch( @@ -333,7 +333,10 @@ def test_login_writes_the_remote_version_file_too(): ) conn = do_manager.peer_manager.connection_router.connections[0] assert ( - conn.read_own_peer_info_file().version_info.syft_client_version == SYFT_VERSION + conn.as_transport() + .peering.read_own_peer_info() + .version_info.syft_client_version + == SYFT_VERSION ) with ( @@ -344,14 +347,18 @@ def test_login_writes_the_remote_version_file_too(): new_conn = upgraded.peer_manager.connection_router.connections[0] # Still the pre-upgrade version: nothing has refreshed it yet. assert ( - new_conn.read_own_peer_info_file().version_info.syft_client_version + new_conn.as_transport() + .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_conn.as_transport() + .peering.read_own_peer_info() + .version_info.syft_client_version == NEW_VERSION ) local = read_local_version(upgraded.syftbox_folder) @@ -395,7 +402,9 @@ def test_the_mismatch_prompt_does_not_return_after_a_login(): conn = do_manager.peer_manager.connection_router.connections[0] 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: conn.as_transport() + .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/syft/sync/connections/base_connection.py b/syft/sync/connections/base_connection.py index 0bb63725e05..016634a2f5d 100644 --- a/syft/sync/connections/base_connection.py +++ b/syft/sync/connections/base_connection.py @@ -26,10 +26,6 @@ class ConnectionConfig(BaseModel): 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 ): diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index b6625e80353..f1d70b85c56 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -74,7 +74,7 @@ def add_connection(self, connection: SyftboxPlatformConnection): self.connections.append(connection) def get_authenticated_email(self) -> str: - return self.connections[0].get_authenticated_email() + return self.transport.peering.whoami() def connection_for_send_message(self) -> SyftboxPlatformConnection: return self.connections[0] @@ -420,19 +420,6 @@ def watcher_download_collection( } return files - def connection_for_peer_file_read( - self, create_new: bool = False - ) -> SyftboxPlatformConnection: - """Connection for reading files a peer publishes: peer info and key bundles. - - Can create a new one for thread safety. - """ - existing_connection = self.connections[0] - if create_new: - return self.copy_connection(existing_connection) - else: - return existing_connection - def write_peer_info_file(self, peer_info: "PeerInfo") -> None: """Write the peer info file to own SyftBox folder.""" self.transport.peering.publish_peer_info(peer_info) diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 96066cc841a..5c7f3b00e41 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -14,10 +14,10 @@ from syft.sync.connections.drive.drive_collections import DriveCollections from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore +from syft.sync.connections.drive.drive_peering import DrivePeering from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform from syft.sync.transport.protocols import ( P2PConnection, - Peering, Transport, WatcherStore, ) @@ -25,7 +25,6 @@ if TYPE_CHECKING: from syft.sync.connections.drive.gdrive_transport import GDriveConnection - from syft.sync.version.peer_info import PeerInfo class DriveP2PConnection(P2PConnection): @@ -95,55 +94,6 @@ def reset(self, peer: str) -> None: raise NotImplementedError("Watcher storage on Drive is not built yet") -class DrivePeering(Peering): - """Identity from the Google account; channels, the peer list, key bundles - and version files in /SyftBox.""" - - def __init__(self, connection: GDriveConnection): - self._conn = connection - - 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 drive_transport(connection: GDriveConnection) -> Transport: """The transport jobs over one Drive connection.""" return Transport( @@ -151,6 +101,6 @@ def drive_transport(connection: GDriveConnection) -> Transport: owner_store=DriveOwnerStore(connection.layout), watcher_store=DriveWatcherStore(connection), collections=DriveCollections(connection.layout), - peering=DrivePeering(connection), + peering=DrivePeering(connection.layout), platform=GdriveFilesPlatform(), ) diff --git a/syft/sync/connections/drive/drive_peering.py b/syft/sync/connections/drive/drive_peering.py new file mode 100644 index 00000000000..a573c2a83ad --- /dev/null +++ b/syft/sync/connections/drive/drive_peering.py @@ -0,0 +1,487 @@ +"""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.gdrive_layout import ( + GDRIVE_P2P_FOLDER_DATASITE_PREFIX, + DriveLayout, + 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, layout: DriveLayout): + self._layout = layout + # Encryption bundles folder cache + self._encryption_bundles_folder_id: str | None = None + # Cached SYFT_peers.json contents (None = not loaded yet). + self._peers_json_cache: dict[str, dict[str, Any]] | None = None + + def copy_caches_to(self, other: "DrivePeering") -> None: + """Seed the caches of ``other`` from this store.""" + 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 + ) + + def reset_caches(self) -> None: + self._encryption_bundles_folder_id = None + self._peers_json_cache = None + + # -- Peering --------------------------------------------------------------- + + def whoami(self) -> str: + return self.get_authenticated_email() + + def open_channels(self, peer: str) -> None: + self.add_peer(peer_email=peer) + + def incoming_requests(self) -> list[str]: + return self.get_peer_requests() + + def read_peer_registry(self, force_download: bool = False) -> dict: + return self._get_peers_json(force_download=force_download) + + def update_peer_registry( + self, peer: str, state: str, public_encryption_bundle: dict | None = None + ) -> None: + self._update_peer_state(peer, state, public_encryption_bundle) + + def publish_key_bundle(self, peer: str, bundle_json: str) -> None: + self.write_encryption_bundle(peer, bundle_json) + + def grant_key_bundle_read(self, peer: str) -> None: + self.share_encryption_bundles_folder(peer) + + def read_own_key_bundle(self, peer: str) -> str | None: + return self.read_own_encryption_bundle(peer) + + def read_peer_key_bundle(self, peer: str) -> str | None: + return self.read_peer_encryption_bundle(peer) + + def publish_peer_info(self, peer_info: PeerInfo) -> None: + self.write_peer_info_file(peer_info) + + def read_own_peer_info(self) -> PeerInfo | None: + return self.read_own_peer_info_file() + + def read_peer_info(self, peer: str) -> PeerInfo | None: + return self.read_peer_info_file(peer) + + def grant_peer_info_read(self, peer: str, peer_info: PeerInfo) -> None: + self.share_peer_info_file_with_peer(peer, peer_info) + + # -- Drive ----------------------------------------------------------------- + + def get_authenticated_email(self) -> str: + """Return the email of the Google account behind the Drive service.""" + about = execute_with_retries( + self._layout.client.drive_service.about().get(fields="user(emailAddress)") + ) + return about["user"]["emailAddress"] + + 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._layout.get_peer_datasite_inbox_id(peer_email) + if inbox_id is None: + inbox_id = self._layout.create_peer_datasite_folder(peer_email, "inbox") + self._layout.client.add_permission(inbox_id, peer_email, write=True) + + # Outbox: peer writes events, I read + outbox_id = self._layout.get_peer_datasite_outbox_id(peer_email) + if outbox_id is None: + outbox_id = self._layout.create_peer_datasite_folder(peer_email, "outbox") + self._layout.client.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._layout.get_syftbox_folder_id() + query = f"name='{SYFT_PEERS_FILE}' and '{syftbox_folder_id}' in parents and trashed=false" + results = execute_with_retries( + self._layout.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._layout.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 _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._layout.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._layout.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._layout.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._layout.client.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. + """ + folders = self._layout.client.list_all_files( + f"name contains '{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#' " + f"and name contains '#{self._layout.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._layout.email: + all_folder_peers.add(folder.peer_email) + + peers_data = self._get_peers_json() + 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._layout.get_syftbox_folder_id() + query = f"name='{SYFT_VERSION_FILE}' and '{syftbox_folder_id}' in parents and trashed=false" + results = execute_with_retries( + self._layout.client.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._layout.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._layout.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._layout.client.create_syft_file(file_metadata, file_payload) + else: + # Update existing file + execute_with_retries( + self._layout.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._layout.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_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._layout.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 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._layout.client.add_permission(file_id, peer_email, 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._layout.email + ).as_string() + folder_id = self._layout.client.find_folder_by_name( + folder_name, parent_id=self._layout.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._layout.email + ).as_string() + folder_id = self._layout.client.create_folder( + folder_name, self._layout.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._layout.email, peer_email) + file_payload, _ = self._layout.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._layout.client.drive_service.files().list(q=query, fields="files(id)") + ) + items = results.get("files", []) + + if items: + execute_with_retries( + self._layout.client.drive_service.files().update( + fileId=items[0]["id"], media_body=file_payload + ) + ) + else: + file_metadata = {"name": filename, "parents": [folder_id]} + self._layout.client.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._layout.client.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._layout.email, peer_email) + query = f"name='{filename}' and '{folder_id}' in parents and trashed=false" + results = execute_with_retries( + self._layout.client.drive_service.files().list(q=query, fields="files(id)") + ) + items = results.get("files", []) + if not items: + return None + return self._layout.client.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._layout.email}.json + in peer's bundles folder. + """ + filename = self._encryption_bundle_filename(peer_email, self._layout.email) + query = f"name='{filename}' and trashed=false and '{peer_email}' in owners" + results = execute_with_retries( + self._layout.client.drive_service.files().list(q=query, fields="files(id)") + ) + items = results.get("files", []) + if not items: + return None + try: + data = self._layout.client.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/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index e6c740d4e10..275286d5ff4 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -1,16 +1,13 @@ """Google Drive Files transport layer implementation""" -import json import logging from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Optional +from typing import TYPE_CHECKING, Any from google.oauth2.credentials import Credentials as GoogleCredentials from google_auth_httplib2 import AuthorizedHttp from googleapiclient.discovery import build from googleapiclient.http import build_http -from pydantic import BaseModel -from syft_migration import MigrationError from syft.sync.connections.base_connection import ( SyftboxPlatformConnection, @@ -18,22 +15,19 @@ from syft.sync.connections.drive.drive_collections import DriveCollections from syft.sync.connections.drive.drive_connections import drive_transport from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore +from syft.sync.connections.drive.drive_peering import DrivePeering from syft.sync.connections.drive.gdrive_client import ( - GOOGLE_FOLDER_MIME_TYPE, SYFT_FILE_PROPERTY_KEY, SYFT_FILE_PROPERTY_VALUE, GDriveClient, ) from syft.sync.connections.drive.gdrive_layout import ( - GDRIVE_P2P_FOLDER_DATASITE_PREFIX, SYFTBOX_FOLDER, DriveLayout, - GdriveP2PFolder, 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, @@ -47,10 +41,6 @@ MessageFileName, ProposedFileChangesMessage, ) -from syft.sync.transport.naming import ( - PEERS_META_KEY, - SYFT_PEERS_VERSION, -) from syft.sync.transport.protocols import Transport from syft.sync.utils.syftbox_utils import check_env @@ -58,7 +48,6 @@ 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 @@ -93,20 +82,6 @@ def build_drive_service( return build("drive", "v3", http=authorized_http) -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 GDriveConnection(SyftboxPlatformConnection): """Google Drive Files API transport layer""" @@ -123,12 +98,6 @@ class Config: # one set of caches. _transport: Transport | 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) @@ -222,6 +191,12 @@ def _collections(self) -> DriveCollections: assert isinstance(collections, DriveCollections) return collections + @property + def _peering(self) -> DrivePeering: + peering = self.as_transport().peering + assert isinstance(peering, DrivePeering) + return peering + def clone_for_thread(self) -> "GDriveConnection": """A connection safe to use from another thread: the Drive client is not.""" return self.copy() @@ -249,182 +224,12 @@ def _copy_caches_to(self, other: "GDriveConnection") -> None: self.layout.copy_caches_to(other.layout) self._owner_store.copy_caches_to(other._owner_store) self._collections.copy_caches_to(other._collections) - 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 - ) + self._peering.copy_caches_to(other._peering) @property def environment(self) -> Environment: return check_env() - def get_authenticated_email(self) -> str: - """Return the email of the Google account behind the Drive service.""" - about = execute_with_retries( - self.client.drive_service.about().get(fields="user(emailAddress)") - ) - return about["user"]["emailAddress"] - - 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.layout.get_peer_datasite_inbox_id(peer_email) - if inbox_id is None: - inbox_id = self.layout.create_peer_datasite_folder(peer_email, "inbox") - self.client.add_permission(inbox_id, peer_email, write=True) - - # Outbox: peer writes events, I read - outbox_id = self.layout.get_peer_datasite_outbox_id(peer_email) - if outbox_id is None: - outbox_id = self.layout.create_peer_datasite_folder(peer_email, "outbox") - self.client.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.layout.get_syftbox_folder_id() - query = f"name='{SYFT_PEERS_FILE}' and '{syftbox_folder_id}' in parents 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 _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.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 _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.layout.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.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.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.client.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. - """ - folders = self.client.list_all_files( - 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="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.email: - all_folder_peers.add(folder.peer_email) - - peers_data = self._get_peers_json() - 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 watcher_download_raw_events_from_outbox( self, peer_email: str, since_timestamp: float | None ) -> list[bytes]: @@ -663,8 +468,7 @@ def reset_caches(self): self.layout.reset_caches() self._owner_store.reset_caches() self._collections.reset_caches() - self._encryption_bundles_folder_id = None - self._peers_json_cache = None + self._peering.reset_caches() def gather_all_file_and_folder_ids(self) -> list[str]: syftbox_folder_id = self.layout.get_syftbox_folder_id() @@ -716,212 +520,3 @@ def get_inbox_proposed_event_id_from_name( ) items = results.get("files", []) return items[0]["id"] if items else None - - 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.layout.get_syftbox_folder_id() - query = f"name='{SYFT_VERSION_FILE}' and '{syftbox_folder_id}' in parents 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 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.layout.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.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.client.create_syft_file(file_metadata, file_payload) - else: - # Update existing file - execute_with_retries( - self.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.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_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.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 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.client.add_permission(file_id, peer_email, write=False) - - # ========================================================================= - # 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.client.find_folder_by_name( - folder_name, parent_id=self.layout.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.client.create_folder( - folder_name, self.layout.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.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.client.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - - if items: - execute_with_retries( - self.client.drive_service.files().update( - fileId=items[0]["id"], media_body=file_payload - ) - ) - else: - file_metadata = {"name": filename, "parents": [folder_id]} - self.client.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.client.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.client.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - if not items: - return None - return self.client.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.client.drive_service.files().list(q=query, fields="files(id)") - ) - items = results.get("files", []) - if not items: - return None - try: - data = self.client.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/login_utils.py b/syft/sync/login_utils.py index ee757e0bea5..1bb473c7939 100644 --- a/syft/sync/login_utils.py +++ b/syft/sync/login_utils.py @@ -22,7 +22,7 @@ def _read_remote_version( from syft.sync.connections.drive.gdrive_transport import GDriveConnection conn = GDriveConnection.from_token_path(email=email, token_path=token_path) - peer_info = conn.read_own_peer_info_file() + peer_info = conn.as_transport().peering.read_own_peer_info() return peer_info.version_info if peer_info else None diff --git a/syft/sync/version/peer_manager.py b/syft/sync/version/peer_manager.py index b06676aae9f..6fb06097e7f 100644 --- a/syft/sync/version/peer_manager.py +++ b/syft/sync/version/peer_manager.py @@ -274,11 +274,8 @@ def _load_single_peer_info(self, peer_email: str) -> tuple[str, Optional[PeerInf TransportUnavailableError: when the backend cannot be reached. ValueError: when the document is not one this client reads. """ - # Create new connection for thread safety - connection = self.connection_router.connection_for_peer_file_read( - create_new=True - ) - return (peer_email, connection.read_peer_info_file(peer_email)) + peering = self.connection_router.transport_for_thread().peering + return (peer_email, peering.read_peer_info(peer_email)) def load_peer_infos_parallel( self, peer_emails: List[str], force: bool = False @@ -640,10 +637,8 @@ def _check_published_bundles(self, peer_emails: list[str]) -> None: def _read_own_published_bundle(self, peer_email: str) -> dict | None: """The bundle this client published for ``peer_email``, or None (thread-safe).""" - connection = self.connection_router.connection_for_peer_file_read( - create_new=True - ) - bundle_json = connection.read_own_encryption_bundle(peer_email) + peering = self.connection_router.transport_for_thread().peering + bundle_json = peering.read_own_key_bundle(peer_email) if bundle_json is None: return None return json.loads(bundle_json)["public_encryption_bundle"] diff --git a/tests/integration/with_unit_coverage/test_sync_manager.py b/tests/integration/with_unit_coverage/test_sync_manager.py index ab6fb0cd59a..300f102d867 100644 --- a/tests/integration/with_unit_coverage/test_sync_manager.py +++ b/tests/integration/with_unit_coverage/test_sync_manager.py @@ -565,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_connection.as_transport().peering.publish_peer_info( own_info.model_copy(update={"version_info": new_version}) ) @@ -587,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_connection.as_transport().peering.publish_peer_info(own_info) diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 008dd1b285f..2d37a321526 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -283,7 +283,7 @@ def test_peer_registry_update_while_drive_is_down_keeps_other_peers(pair, monkey do_conn, _ = pair peering = do_conn.as_transport().peering peering.update_peer_registry(DS, "accepted") - do_conn._peers_json_cache = None # a new process has no cache + do_conn._peering._peers_json_cache = None # a new process has no cache monkeypatch.setattr(GDriveClient, "download_file", _drive_down) with pytest.raises(TransportUnavailableError): @@ -296,7 +296,7 @@ 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 = pair[0].as_transport() assert owner.peering.read_own_key_bundle(DS) is None - assert pair[0]._find_encryption_bundles_folder_id() is None + assert pair[0]._peering._find_encryption_bundles_folder_id() is None def test_peering_publishes_a_key_bundle_both_sides_can_read(pair): diff --git a/tests/unit/test_encryption_policy.py b/tests/unit/test_encryption_policy.py index dad4165ab32..ee8cd4ca6ec 100644 --- a/tests/unit/test_encryption_policy.py +++ b/tests/unit/test_encryption_policy.py @@ -173,8 +173,13 @@ def test_missing_version_file_created_with_own_policy(): a_conn, b_conn = 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 + a_conn.as_transport().peering.grant_peer_info_read( + "b@example.com", _peer_info(REQUIRED) + ) + assert ( + b_conn.as_transport().peering.read_peer_info("a@example.com").encryption_policy + == REQUIRED + ) def test_malformed_peer_bundle_raises_on_peer_load(): diff --git a/tests/unit/test_peers_json_version.py b/tests/unit/test_peers_json_version.py index c3ff26f2839..f9e0fe282c5 100644 --- a/tests/unit/test_peers_json_version.py +++ b/tests/unit/test_peers_json_version.py @@ -16,6 +16,7 @@ from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.gdrive_layout import DriveLayout +from syft.sync.connections.drive.drive_peering import DrivePeering from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.peers.peer import PeerState from syft.sync.transport.naming import PEERS_META_KEY, SYFT_PEERS_VERSION @@ -26,7 +27,7 @@ def _conn(peers_data): conn = GDriveConnection(email="alice@example.com", verbose=False) conn.layout = DriveLayout(GDriveClient(Mock()), conn.email) - conn._peers_json_cache = dict(peers_data) + conn._peering._peers_json_cache = dict(peers_data) return conn @@ -41,14 +42,16 @@ def _router(conn): def test_a_write_stamps_the_reserved_entry(): conn = _conn({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(DriveLayout, "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"}}) + conn._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 conn._peering._peers_json_cache[PEERS_META_KEY] == { + "version": SYFT_PEERS_VERSION + } + assert conn._peering._peers_json_cache[PEER] == {"state": "accepted"} def test_the_reserved_entry_is_not_a_peer(): diff --git a/tests/unit/test_sync_manager.py b/tests/unit/test_sync_manager.py index aaaee0fcd3f..c3ed90cf2e1 100644 --- a/tests/unit/test_sync_manager.py +++ b/tests/unit/test_sync_manager.py @@ -489,7 +489,7 @@ def _peer_router(ds_manager, do_manager, email: str) -> ConnectionRouter: conn = GDriveConnection.from_service( email, ds_manager._connection_router.connections[0].client.drive_service ) - conn.add_peer(do_manager.email) + conn.as_transport().peering.open_channels(do_manager.email) _add_known_peer(do_manager, email) return ConnectionRouter( connections=[conn], 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 a6af55d91d1..004de5e13bb 100644 --- a/tests/unit/test_version_negotiation.py +++ b/tests/unit/test_version_negotiation.py @@ -227,7 +227,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 ( From eb44e5e94366401e22b4a2fc1c1977b81016833b Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 16:03:53 -0300 Subject: [PATCH 14/28] refactor: move sending into DriveP2PConnection --- syft/sync/connections/base_connection.py | 16 --- .../connections/drive/drive_connections.py | 49 +------- syft/sync/connections/drive/drive_p2p.py | 112 ++++++++++++++++++ .../connections/drive/gdrive_transport.py | 76 ------------ tests/unit/test_sync_manager.py | 3 +- 5 files changed, 115 insertions(+), 141 deletions(-) create mode 100644 syft/sync/connections/drive/drive_p2p.py diff --git a/syft/sync/connections/base_connection.py b/syft/sync/connections/base_connection.py index 016634a2f5d..f0616f2a3e6 100644 --- a/syft/sync/connections/base_connection.py +++ b/syft/sync/connections/base_connection.py @@ -1,6 +1,5 @@ from typing import TYPE_CHECKING, ClassVar, Type from pydantic import BaseModel -from syft.sync.messages.proposed_filechange import ProposedFileChangesMessage if TYPE_CHECKING: from syft.sync.transport.protocols import Transport @@ -26,11 +25,6 @@ class ConnectionConfig(BaseModel): class SyftboxPlatformConnection(BaseModel): config: ConnectionConfig | None = None - 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) @@ -51,22 +45,12 @@ def clone_for_thread(self) -> "SyftboxPlatformConnection": # 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_download_next_raw_proposed_message_from_inbox( self, sender_email: str ) -> tuple[bytes, str] | None: """Download next message from inbox as raw bytes. Returns (data, file_id) or None.""" 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]: diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 5c7f3b00e41..52b698290ef 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -14,63 +14,18 @@ from syft.sync.connections.drive.drive_collections import DriveCollections 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.platforms.gdrive_files_platform import GdriveFilesPlatform from syft.sync.transport.protocols import ( - P2PConnection, Transport, WatcherStore, ) -from syft.sync.transport.types import MessageKind, RawMessage, SendStatus if TYPE_CHECKING: from syft.sync.connections.drive.gdrive_transport import GDriveConnection -class DriveP2PConnection(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. - - Only sending goes through this class so far. The router still reads and - archives messages with the Drive methods directly, because doing it through - ``receive`` and ``ack`` changes how they behave (one rule for both folders, - delete on ack, no read cursor); that lands with send-and-forget. - """ - - def __init__(self, connection: GDriveConnection): - self._conn = connection - - 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, max_messages: int | None = None) -> list[RawMessage]: - raise NotImplementedError( - "Messages are still read through the Drive methods; receiving " - "through P2PConnection lands with send-and-forget." - ) - - def ack(self, message: RawMessage) -> None: - raise NotImplementedError( - "Messages are still archived through the Drive methods; " - "acknowledging through P2PConnection lands with send-and-forget." - ) - - class DriveWatcherStore(WatcherStore): """The data scientist's storage on Drive. Not built yet: today the data scientist keeps what it received in a local cache.""" @@ -97,7 +52,7 @@ def reset(self, peer: str) -> None: def drive_transport(connection: GDriveConnection) -> Transport: """The transport jobs over one Drive connection.""" return Transport( - p2p=DriveP2PConnection(connection), + p2p=DriveP2PConnection(connection.layout), owner_store=DriveOwnerStore(connection.layout), watcher_store=DriveWatcherStore(connection), collections=DriveCollections(connection.layout), diff --git a/syft/sync/connections/drive/drive_p2p.py b/syft/sync/connections/drive/drive_p2p.py new file mode 100644 index 00000000000..ee2eb9f5fb6 --- /dev/null +++ b/syft/sync/connections/drive/drive_p2p.py @@ -0,0 +1,112 @@ +"""Messages between peers on Google Drive: the inbox and outbox folders.""" + +from syft.sync.connections.drive.gdrive_layout import DriveLayout +from syft.sync.connections.drive.gdrive_retry import execute_with_retries +from syft.sync.transport.protocols import P2PConnection +from syft.sync.transport.types import MessageKind, RawMessage, SendStatus + + +class DriveP2PConnection(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. + + Only sending goes through this class so far. The router still reads and + archives messages with the Drive methods directly, because doing it through + ``receive`` and ``ack`` changes how they behave (one rule for both folders, + delete on ack, no read cursor); that lands with send-and-forget. + """ + + def __init__(self, layout: DriveLayout): + self._layout = layout + + 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, max_messages: int | None = None) -> list[RawMessage]: + raise NotImplementedError( + "Messages are still read through the Drive methods; receiving " + "through P2PConnection lands with send-and-forget." + ) + + def ack(self, message: RawMessage) -> None: + raise NotImplementedError( + "Messages are still archived through the Drive methods; " + "acknowledging through P2PConnection lands with send-and-forget." + ) + + # -- Drive ----------------------------------------------------------------- + + def watcher_send_raw_bytes_to_inbox( + self, recipient: str, filename: str, data: bytes + ) -> None: + inbox_id = self._layout.get_peer_datasite_inbox_id(recipient) + if inbox_id is None: + raise Exception(f"Inbox folder for {recipient}'s datasite not found") + + payload, _ = self._layout.client.create_file_payload(data) + file_metadata = {"name": filename, "parents": [inbox_id]} + + result = self._layout.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._layout.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._layout.get_own_datasite_outbox_id(recipient) + if outbox_folder_id is None: + raise ValueError(f"Outbox folder for {recipient} not found") + + file_payload, _ = self._layout.client.create_file_payload(data) + file_metadata = {"name": filename, "parents": [outbox_folder_id]} + + result = self._layout.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._layout.client.drive_service.files().update( + fileId=file_id, + addParents=outbox_folder_id, + fields="id, parents", + ) + ) diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 275286d5ff4..1d792df3fd9 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -295,43 +295,6 @@ def watcher_get_outbox_file_metadatas( ) return result - def owner_write_raw_bytes_to_outbox( - self, recipient: str, filename: str, data: bytes - ) -> None: - outbox_folder_id = self.layout.get_own_datasite_outbox_id(recipient) - if outbox_folder_id is None: - raise ValueError(f"Outbox folder for {recipient} not found") - - file_payload, _ = self.client.create_file_payload(data) - file_metadata = {"name": filename, "parents": [outbox_folder_id]} - - result = self.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.client.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_remove_proposed_filechange_message_from_inbox( self, proposed_filechange_message: ProposedFileChangesMessage, @@ -425,45 +388,6 @@ def owner_list_outbox_files(self, recipient_email: str) -> list[dict] | None: folder_id, since_timestamp=None ) - def watcher_send_raw_bytes_to_inbox( - self, recipient: str, filename: str, data: bytes - ) -> None: - inbox_id = self.layout.get_peer_datasite_inbox_id(recipient) - if inbox_id is None: - raise Exception(f"Inbox folder for {recipient}'s datasite not found") - - payload, _ = self.client.create_file_payload(data) - file_metadata = {"name": filename, "parents": [inbox_id]} - - result = self.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.client.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.layout.reset_caches() self._owner_store.reset_caches() diff --git a/tests/unit/test_sync_manager.py b/tests/unit/test_sync_manager.py index c3ed90cf2e1..72b47086d87 100644 --- a/tests/unit/test_sync_manager.py +++ b/tests/unit/test_sync_manager.py @@ -186,13 +186,12 @@ def test_sync_existing_inbox_state_do(): do_manager.datasite_owner_syncer.perm_context.open(".").grant_write_access( ds_manager.email ) - connection_ds = ds_manager._connection_router.connections[0] 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 ) From 30c771e27cf25155668527c3fb0b3937fae37da9 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 28 Oct 2026 16:06:55 -0300 Subject: [PATCH 15/28] docs: describe the Drive connection and transport jobs after the move --- syft/sync/connections/drive/drive_connections.py | 10 ++++------ syft/sync/connections/drive/gdrive_transport.py | 10 ++++++++-- tests/unit/test_drive_connections.py | 4 ++-- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 52b698290ef..99d37c7165e 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -1,11 +1,9 @@ """The transport jobs 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. +``drive_transport`` builds one of each job over the ``DriveLayout`` of a +connection. Each job has its own module: ``drive_p2p``, ``drive_owner_store``, +``drive_collections`` and ``drive_peering``. Storage for the data scientist is +not built on Drive yet. """ from __future__ import annotations diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 1d792df3fd9..185b01b626a 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -1,4 +1,4 @@ -"""Google Drive Files transport layer implementation""" +"""A Google Drive account as a syft connection.""" import logging from pathlib import Path @@ -83,7 +83,13 @@ def build_drive_service( class GDriveConnection(SyftboxPlatformConnection): - """Google Drive Files API transport layer""" + """A Google Drive account as a syft connection. + + It builds the Drive service and the ``DriveLayout``, and gives out the + transport jobs (``as_transport``) and copies for other threads. Two jobs + stay here until they have a contract: reading and archiving messages, which + waits for send-and-forget, and the account cleanup of ``delete_syftbox``. + """ class Config: arbitrary_types_allowed = True diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 2d37a321526..b2df5e7364a 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 From 0bdff7a2a4aa0e9b65183e7b668d68189aa93133 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 10:53:50 -0300 Subject: [PATCH 16/28] fix: skip incremental checkpoints the full checkpoint already holds --- syft/sync/sync/datasite_owner_syncer.py | 22 +++++++++-- tests/unit/test_checkpoints.py | 49 +++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 19fd3ba1a36..39b1ad91920 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -118,6 +118,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: @@ -381,7 +394,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 @@ -1239,10 +1254,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/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index b33bc12adfc..0a06188d263 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -779,3 +779,52 @@ def forbidden(self, file_id: str) -> bytes: 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" From 47ed9aeb4b53116bccf1ba257d0a40cd85a54813 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 11:54:53 -0300 Subject: [PATCH 17/28] fix: raise when a collection folder delete fails --- .../sync/connections/drive/drive_collections.py | 2 +- tests/unit/test_drive_connections.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/syft/sync/connections/drive/drive_collections.py b/syft/sync/connections/drive/drive_collections.py index 7518e3473e8..5ee10466ac9 100644 --- a/syft/sync/connections/drive/drive_collections.py +++ b/syft/sync/connections/drive/drive_collections.py @@ -268,7 +268,7 @@ def owner_delete_collection(self, prefix: str, tag: str) -> None: collections = self.owner_list_all_collections_with_permissions(prefix) for c in collections: if c.tag == tag: - self._layout.client.delete_file_by_id(c.folder_id) + self._layout.client.delete_file_by_id(c.folder_id, raise_on_error=True) folder_name = CollectionFolder( prefix=prefix, tag=c.tag, diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index b2df5e7364a..4f06dd5434a 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -16,6 +16,7 @@ from syft.sync.connections.drive.drive_collections import CollectionFolder 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, ) @@ -262,6 +263,22 @@ def test_a_published_collection_is_listed_and_downloaded_by_its_reader(pair): assert owner.collections.list_mine(DATASETS) == [] +def test_a_failed_collection_delete_raises(pair, monkeypatch): + do_conn, _ = pair + owner = do_conn.as_transport().collections + 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 ---------------------------------------------------------------------- From 26e852b929f71f27c016d72685e7e94e3845fa5a Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 13:34:36 -0300 Subject: [PATCH 18/28] fix: point the benchmark and version comment at the moved Drive code --- benchmarks/download_files_parallel.py | 4 ++-- syft/version.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) 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/syft/version.py b/syft/version.py index 9309193be21..8f1f47a0782 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 gdrive_layout.py and +# drive_owner_store.py in sync/connections/drive/). SYFT_VERSION = "0.10.1" # Minimum client version we support communicating with. From 35f0487f88fd8fb1207826ef0008649c8d60ad38 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 17:17:18 -0300 Subject: [PATCH 19/28] fix: raise every permission error in batch shares The batch callback ignored any error whose text contained "alreadyShared". Tested live on real Drive: sharing a folder with a user who already has access returns that permission with status 200, in single and batch requests. Drive never sends an "alreadyShared" error, so the text check could only hide real errors. --- packages/syft-rds/src/syft_rds/client.py | 5 ++--- syft/sync/connections/drive/gdrive_client.py | 4 ++-- tests/unit/test_gdrive_client.py | 15 +++++++++++++++ 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/packages/syft-rds/src/syft_rds/client.py b/packages/syft-rds/src/syft_rds/client.py index c50725a2f11..6e5cdd90642 100644 --- a/packages/syft-rds/src/syft_rds/client.py +++ b/packages/syft-rds/src/syft_rds/client.py @@ -387,9 +387,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) # ------------------------------------------------------------------ # diff --git a/syft/sync/connections/drive/gdrive_client.py b/syft/sync/connections/drive/gdrive_client.py index 7963fcd5d82..b4523b6fe0c 100644 --- a/syft/sync/connections/drive/gdrive_client.py +++ b/syft/sync/connections/drive/gdrive_client.py @@ -174,8 +174,8 @@ 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): + # Drive answers a grant the user already has with that permission. + if exception is not None: raise exception BATCH_SIZE = 100 diff --git a/tests/unit/test_gdrive_client.py b/tests/unit/test_gdrive_client.py index d113a27e48c..252a17882a5 100644 --- a/tests/unit/test_gdrive_client.py +++ b/tests/unit/test_gdrive_client.py @@ -5,6 +5,7 @@ from syft.sync.connections.drive.mock_drive_service import ( MockDeleteRequest, + MockPermissionCreateRequest, http_error, pair_with_mock_service, ) @@ -22,3 +23,17 @@ def forbidden(self) -> None: monkeypatch.setattr(MockDeleteRequest, "execute", forbidden) with pytest.raises(HttpError): do_conn.client.delete_multiple_files_by_ids([folder_id]) + + +def test_batch_share_raises_error_whose_text_says_already_shared(monkeypatch): + do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") + folder_id = do_conn.client.create_folder( + "folder", do_conn.layout.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_conn.client.batch_add_permissions(folder_id, ["ds@example.com"]) From 5397be3ce645a22c5ddde61950ecc59d66d38e9c Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 19:38:07 -0300 Subject: [PATCH 20/28] docs: state what the Drive classes do and use DO/DS --- syft/sync/connections/drive/drive_connections.py | 7 +++---- syft/sync/connections/drive/drive_owner_store.py | 5 ++--- syft/sync/connections/drive/drive_p2p.py | 4 ++-- syft/sync/connections/drive/gdrive_transport.py | 6 +++--- 4 files changed, 10 insertions(+), 12 deletions(-) diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 99d37c7165e..c642bb78352 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -2,8 +2,7 @@ ``drive_transport`` builds one of each job over the ``DriveLayout`` of a connection. Each job has its own module: ``drive_p2p``, ``drive_owner_store``, -``drive_collections`` and ``drive_peering``. Storage for the data scientist is -not built on Drive yet. +``drive_collections`` and ``drive_peering``. """ from __future__ import annotations @@ -25,8 +24,8 @@ class DriveWatcherStore(WatcherStore): - """The data scientist's storage on Drive. Not built yet: today the data - scientist keeps what it received in a local cache.""" + """The DS's storage on Drive. Not built yet: today the DS keeps what it + received in a local cache.""" def __init__(self, connection: GDriveConnection): self._conn = connection diff --git a/syft/sync/connections/drive/drive_owner_store.py b/syft/sync/connections/drive/drive_owner_store.py index 3a09ad2d8e8..81414846522 100644 --- a/syft/sync/connections/drive/drive_owner_store.py +++ b/syft/sync/connections/drive/drive_owner_store.py @@ -1,4 +1,4 @@ -"""The data owner's own storage on Google Drive. +"""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. @@ -45,8 +45,7 @@ def _filter_valid_file_metadatas( class DriveOwnerStore(OwnerStore): - """The data owner's own folders in /SyftBox: event log, checkpoints, - rolling state.""" + """The DO's own folders in /SyftBox: event log, checkpoints, rolling state.""" def __init__(self, layout: DriveLayout): self._layout = layout diff --git a/syft/sync/connections/drive/drive_p2p.py b/syft/sync/connections/drive/drive_p2p.py index ee2eb9f5fb6..6debfd8ee31 100644 --- a/syft/sync/connections/drive/drive_p2p.py +++ b/syft/sync/connections/drive/drive_p2p.py @@ -9,8 +9,8 @@ class DriveP2PConnection(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 + 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. Only sending goes through this class so far. The router still reads and diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 185b01b626a..2d754fc8427 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -86,9 +86,9 @@ class GDriveConnection(SyftboxPlatformConnection): """A Google Drive account as a syft connection. It builds the Drive service and the ``DriveLayout``, and gives out the - transport jobs (``as_transport``) and copies for other threads. Two jobs - stay here until they have a contract: reading and archiving messages, which - waits for send-and-forget, and the account cleanup of ``delete_syftbox``. + transport jobs (``as_transport``) and copies for other threads. It also + reads and archives messages, and cleans up the account for + ``delete_syftbox``. """ class Config: From 257fa8af7953b6da41e8d61963a340d763c5ea50 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 19:50:11 -0300 Subject: [PATCH 21/28] refactor: cache the transport only on the connection and rebuild it on setup --- syft/sync/connections/connection_router.py | 14 +++----------- syft/sync/connections/drive/gdrive_transport.py | 2 ++ tests/unit/test_drive_connections.py | 13 +++++++++++-- 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index f1d70b85c56..311a1d5c211 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -2,7 +2,7 @@ import warnings from typing import TYPE_CHECKING, List, Optional -from pydantic import BaseModel, PrivateAttr +from pydantic import BaseModel from syft_migration import MigratableObject from syft.migrations import ( @@ -47,18 +47,10 @@ class ConnectionRouter(BaseModel): peer_store: PeerStore - # The transport jobs over connections[0], rebuilt if that connection is - # replaced (tests swap in a mock connection with add_connection). - _transport: Optional[Transport] = PrivateAttr(default=None) - _transport_source: Optional[SyftboxPlatformConnection] = PrivateAttr(default=None) - @property def transport(self) -> Transport: - connection = self.connections[0] - if self._transport is None or self._transport_source is not connection: - self._transport = connection.as_transport() - self._transport_source = connection - return self._transport + """The transport jobs over connections[0].""" + return self.connections[0].as_transport() @classmethod def from_configs(cls, email: str, connection_configs: List[ConnectionConfig]): diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 2d754fc8427..8b8d069180b 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -170,6 +170,8 @@ def setup( self.credentials, environment=self.environment ) self.layout = DriveLayout(GDriveClient(drive_service), self.email) + # The jobs hold the layout, so a new layout needs new jobs. + self._transport = None if warm_syftbox_folder_id_cache: self.layout.get_personal_syftbox_folder_id() diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 4f06dd5434a..d61680d5873 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -352,8 +352,7 @@ def _router_with_mock_p2p(): router = ConnectionRouter(connections=[do_conn], peer_store=_store_knowing(DO, DS)) drive = do_conn.as_transport() p2p = Mock(spec=P2PConnection) - router._transport = drive.model_copy(update={"p2p": p2p}) - router._transport_source = do_conn + do_conn._transport = drive.model_copy(update={"p2p": p2p}) return router, p2p @@ -383,6 +382,16 @@ def test_the_router_sends_proposed_changes_through_p2p(): ) +def test_setup_gives_the_jobs_the_new_layout(): + do_conn, _ = pair_with_mock_service(DO, DS) + do_conn.as_transport() + do_conn.setup( + drive_service=do_conn.client.drive_service, + warm_syftbox_folder_id_cache=False, + ) + assert do_conn.as_transport().owner_store._layout is do_conn.layout + + def test_the_router_reuses_its_transport_until_its_connection_changes(): do_conn, other_conn = pair_with_mock_service(DO, DS) router = ConnectionRouter(connections=[do_conn], peer_store=PeerStore(email=DO)) From c90cbe48a58bf368d50bff3287f16950ae39013f Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 20:25:23 -0300 Subject: [PATCH 22/28] fix: raise when deleting a Drive file fails --- syft/sync/connections/drive/drive_collections.py | 2 +- syft/sync/connections/drive/gdrive_client.py | 9 ++------- syft/sync/connections/drive/gdrive_transport.py | 4 ++-- 3 files changed, 5 insertions(+), 10 deletions(-) diff --git a/syft/sync/connections/drive/drive_collections.py b/syft/sync/connections/drive/drive_collections.py index 5ee10466ac9..7518e3473e8 100644 --- a/syft/sync/connections/drive/drive_collections.py +++ b/syft/sync/connections/drive/drive_collections.py @@ -268,7 +268,7 @@ def owner_delete_collection(self, prefix: str, tag: str) -> None: collections = self.owner_list_all_collections_with_permissions(prefix) for c in collections: if c.tag == tag: - self._layout.client.delete_file_by_id(c.folder_id, raise_on_error=True) + self._layout.client.delete_file_by_id(c.folder_id) folder_name = CollectionFolder( prefix=prefix, tag=c.tag, diff --git a/syft/sync/connections/drive/gdrive_client.py b/syft/sync/connections/drive/gdrive_client.py index b4523b6fe0c..b26efeaef0b 100644 --- a/syft/sync/connections/drive/gdrive_client.py +++ b/syft/sync/connections/drive/gdrive_client.py @@ -105,13 +105,8 @@ def create_syft_file( ) ) - 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 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, diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 8b8d069180b..1190ac08494 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -411,8 +411,8 @@ def gather_all_file_and_folder_ids(self) -> list[str]: def download_file(self, file_id: str) -> bytes: return self.client.download_file(file_id) - def delete_file_by_id(self, file_id: str, raise_on_error: bool = False): - self.client.delete_file_by_id(file_id, raise_on_error=raise_on_error) + def delete_file_by_id(self, file_id: str) -> None: + self.client.delete_file_by_id(file_id) def delete_multiple_files_by_ids( self, From 3c77a52277907c2b43f9f2a6f30425fcb31fc215 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 21:06:33 -0300 Subject: [PATCH 23/28] refactor: remove the unused DriveOwnerStore.get_events_count_since_checkpoint Nothing called it: ConnectionRouter.get_events_count_since_checkpoint counts through OwnerStore.count_events_since, which parses each event name. This method matched the name prefix by hand instead, so the two could disagree. --- .../connections/drive/drive_owner_store.py | 36 ------------------- syft/sync/connections/drive/drive_peering.py | 3 +- 2 files changed, 2 insertions(+), 37 deletions(-) diff --git a/syft/sync/connections/drive/drive_owner_store.py b/syft/sync/connections/drive/drive_owner_store.py index 81414846522..3f64a3cdda3 100644 --- a/syft/sync/connections/drive/drive_owner_store.py +++ b/syft/sync/connections/drive/drive_owner_store.py @@ -23,7 +23,6 @@ raise_unavailable, ) from syft.sync.events.file_change_event import FileChangeEventsMessageFileName -from syft.sync.transport.naming import timestamp_from_message_name from syft.sync.transport.protocols import OwnerStore from syft.version import SYFT_VERSION @@ -192,41 +191,6 @@ def download_raw_events_since_timestamp( return result - 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._layout.get_personal_syftbox_folder_id() - file_metadatas = self._layout.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 = timestamp_from_message_name(f["name"]) - if event_timestamp is not None and event_timestamp > checkpoint_timestamp: - count += 1 - return count - def _get_checkpoints_folder_name(self) -> str: """Get the checkpoints folder name: {email}-{version}-checkpoints""" return f"{self._layout.email}-{SYFT_VERSION}-checkpoints" diff --git a/syft/sync/connections/drive/drive_peering.py b/syft/sync/connections/drive/drive_peering.py index a573c2a83ad..ea3fd8f985c 100644 --- a/syft/sync/connections/drive/drive_peering.py +++ b/syft/sync/connections/drive/drive_peering.py @@ -51,7 +51,8 @@ def __init__(self, layout: DriveLayout): self._layout = layout # Encryption bundles folder cache self._encryption_bundles_folder_id: str | None = None - # Cached SYFT_peers.json contents (None = not loaded yet). + # 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 copy_caches_to(self, other: "DrivePeering") -> None: From 101ee1b6c4f16853b28e91c400310b130188ba5f Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 21:15:50 -0300 Subject: [PATCH 24/28] refactor: give each DriveOwnerStore operation only its contract name --- .../connections/drive/drive_owner_store.py | 96 +++++-------------- 1 file changed, 23 insertions(+), 73 deletions(-) diff --git a/syft/sync/connections/drive/drive_owner_store.py b/syft/sync/connections/drive/drive_owner_store.py index 3f64a3cdda3..8867c668f6e 100644 --- a/syft/sync/connections/drive/drive_owner_store.py +++ b/syft/sync/connections/drive/drive_owner_store.py @@ -67,75 +67,25 @@ def reset_caches(self) -> None: self._rolling_state_file_id = None self.personal_syftbox_event_id_cache.clear() - # -- OwnerStore ------------------------------------------------------------ - def append_event(self, name: str, data: bytes) -> str: - return self.owner_write_raw_bytes_to_syftbox(name, data) - - def list_event_handles(self, since_timestamp: float | None = None) -> list[str]: - return self.owner_get_all_accepted_event_file_ids(since_timestamp) - - def fetch(self, handle: str) -> bytes: - return self.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.download_raw_events_since_timestamp(since_timestamp) - - def fetch_all_events(self) -> list[bytes]: - return self.owner_download_all_raw_events_from_syftbox() - - def put_checkpoint(self, name: str, data: bytes) -> str: - return self.upload_raw_checkpoint(name, data) - - def latest_checkpoint(self) -> bytes | None: - return self.download_raw_latest_checkpoint() - - def put_incremental(self, name: str, data: bytes) -> str: - return self.upload_raw_incremental_checkpoint(name, data) - - def all_incrementals(self) -> list[bytes]: - return self.download_all_raw_incremental_checkpoints() - - def incremental_count(self) -> int: - return self.get_incremental_checkpoint_count() - - def next_incremental_sequence(self) -> int: - return self.get_next_incremental_sequence_number() - - def delete_incrementals(self) -> None: - self.delete_all_incremental_checkpoints() - - def put_rolling_state(self, name: str, data: bytes) -> str: - return self.upload_raw_rolling_state(name, data) - - def get_rolling_state(self) -> bytes | None: - return self.download_raw_rolling_state() - - # -- Drive ----------------------------------------------------------------- - - def owner_write_raw_bytes_to_syftbox(self, filename: str, data: bytes) -> str: """Write raw bytes to /SyftBox/myemail.""" personal_syftbox_folder_id = self._layout.get_personal_syftbox_folder_id() file_metadata = { - "name": filename, + "name": name, "parents": [personal_syftbox_folder_id], } file_payload, _ = self._layout.client.create_file_payload(data) res = self._layout.client.create_syft_file(file_metadata, file_payload) gdrive_id = res.get("id") - self.personal_syftbox_event_id_cache[filename] = gdrive_id + self.personal_syftbox_event_id_cache[name] = gdrive_id return gdrive_id - def owner_download_raw_bytes_by_id(self, file_id: str) -> bytes: + def fetch(self, handle: str) -> bytes: """Download raw bytes by file ID.""" - return self._layout.client.download_file(file_id) + return self._layout.client.download_file(handle) - def owner_get_all_accepted_event_file_ids( - self, since_timestamp: float | None = None - ) -> list[str]: + def list_event_handles(self, since_timestamp: float | None = None) -> list[str]: personal_syftbox_folder_id = self._layout.get_personal_syftbox_folder_id() file_metadatas = self._layout.get_file_metadatas_from_folder( personal_syftbox_folder_id, since_timestamp=since_timestamp @@ -150,7 +100,7 @@ def owner_get_all_accepted_event_file_ids( return [f["id"] for f in oldest_first] @raise_unavailable("Could not read the event log from Drive.") - def owner_download_all_raw_events_from_syftbox(self) -> list[bytes]: + def fetch_all_events(self) -> list[bytes]: """Download all event files from /SyftBox/myemail as raw bytes.""" personal_syftbox_folder_id = self._layout.get_personal_syftbox_folder_id() file_metadatas = self._layout.get_file_metadatas_from_folder( @@ -167,10 +117,10 @@ def owner_download_all_raw_events_from_syftbox(self) -> list[bytes]: return result @raise_unavailable("Could not read the event log from Drive.") - def download_raw_events_since_timestamp( - self, since_timestamp: float - ) -> list[bytes]: + 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._layout.get_personal_syftbox_folder_id() file_metadatas = self._layout.get_file_metadatas_from_folder( personal_folder_id, since_timestamp=since_timestamp @@ -215,7 +165,7 @@ def _get_or_create_checkpoints_folder_id(self) -> str: syftbox_folder_id = self._layout.get_syftbox_folder_id() return self._layout.client.create_folder(folder_name, syftbox_folder_id) - def upload_raw_checkpoint(self, filename: str, data: bytes) -> str: + 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. @@ -224,19 +174,19 @@ def upload_raw_checkpoint(self, filename: str, data: bytes) -> str: payload, _ = self._layout.client.create_file_payload(data) file_metadata = { - "name": filename, + "name": name, "parents": [folder_id], } result = self._layout.client.create_syft_file(file_metadata, payload) # Only delete old checkpoints after successful upload - self.delete_all_checkpoints(exclude_file_id=result.get("id")) + 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 download_raw_latest_checkpoint(self) -> bytes | None: + 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: @@ -269,7 +219,7 @@ def download_raw_latest_checkpoint(self) -> bytes | None: return self._layout.client.download_file(latest_file["id"]) - def delete_all_checkpoints(self, exclude_file_id: str | None = None): + def _delete_all_checkpoints(self, exclude_file_id: str | None = None): """Delete all existing full checkpoints (not incremental ones). Args: @@ -305,13 +255,13 @@ def delete_all_checkpoints(self, exclude_file_id: str | None = None): except Exception as e: print(f"Warning: Failed to delete checkpoint {item['name']}: {e}") - def upload_raw_incremental_checkpoint(self, filename: str, data: bytes) -> str: + 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._layout.client.create_file_payload(data) file_metadata = { - "name": filename, + "name": name, "parents": [folder_id], } @@ -319,7 +269,7 @@ def upload_raw_incremental_checkpoint(self, filename: str, data: bytes) -> str: return result.get("id") @raise_unavailable("Could not read the incremental checkpoints from Drive.") - def download_all_raw_incremental_checkpoints(self) -> list[bytes]: + 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: @@ -341,7 +291,7 @@ def download_all_raw_incremental_checkpoints(self) -> list[bytes]: return [self._layout.client.download_file(item["id"]) for item in items] - def get_incremental_checkpoint_count(self) -> int: + 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: @@ -356,7 +306,7 @@ def get_incremental_checkpoint_count(self) -> int: ) return len(results.get("files", [])) - def get_next_incremental_sequence_number(self) -> int: + 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: @@ -385,7 +335,7 @@ def get_next_incremental_sequence_number(self) -> int: return max_seq + 1 - def delete_all_incremental_checkpoints(self) -> None: + def delete_incrementals(self) -> None: """Delete all incremental checkpoints (called after compacting).""" folder_id = self._get_checkpoints_folder_id() if folder_id is None: @@ -449,7 +399,7 @@ def _get_or_create_rolling_state_folder_id(self) -> str: self._rolling_state_folder_id = folder_id return folder_id - def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: + 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. @@ -478,7 +428,7 @@ def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: folder_id = self._get_or_create_rolling_state_folder_id() file_metadata = { - "name": filename, + "name": name, "parents": [folder_id], } @@ -488,7 +438,7 @@ def upload_raw_rolling_state(self, filename: str, data: bytes) -> str: return file_id @raise_unavailable("Could not read the rolling state from Drive.") - def download_raw_rolling_state(self) -> bytes | None: + 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. From 1fbd4665a59c3af79359c838f236327c4eb00e2a Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 21:18:24 -0300 Subject: [PATCH 25/28] refactor: give each DriveCollections operation only its contract name --- .../connections/drive/drive_collections.py | 77 +++---------------- tests/unit/test_dataset_collection_listing.py | 10 +-- 2 files changed, 17 insertions(+), 70 deletions(-) diff --git a/syft/sync/connections/drive/drive_collections.py b/syft/sync/connections/drive/drive_collections.py index 7518e3473e8..bdfd3d7d7c4 100644 --- a/syft/sync/connections/drive/drive_collections.py +++ b/syft/sync/connections/drive/drive_collections.py @@ -100,58 +100,7 @@ def copy_caches_to(self, other: "DriveCollections") -> None: def reset_caches(self) -> None: self.collection_folder_id_cache.clear() - # -- Collections ----------------------------------------------------------- - def create(self, prefix: str, tag: str, content_hash: str, owner_email: str) -> str: - return self.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.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.owner_share_collection(prefix, tag, content_hash, users) - - def upload( - self, prefix: str, tag: str, content_hash: str, files: dict[str, bytes] - ) -> None: - self.owner_upload_collection_files(prefix, tag, content_hash, files) - - def list_mine_with_permissions(self, prefix: str) -> list[FileCollection]: - return self.owner_list_all_collections_with_permissions(prefix) - - def list_mine(self, prefix: str) -> list[str]: - return self.owner_list_collections(prefix) - - def delete(self, prefix: str, tag: str) -> None: - self.owner_delete_collection(prefix, tag) - - def list_shared_with_me(self, prefix: str) -> list[dict]: - return self.watcher_list_collections(prefix) - - def list_files( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> list[dict]: - return self.watcher_get_collection_file_metadatas( - prefix, tag, content_hash, owner_email - ) - - def fetch_file(self, handle: str) -> bytes: - return self.watcher_download_collection_file(handle) - - def download( - self, prefix: str, tag: str, content_hash: str, owner_email: str - ) -> dict[str, bytes]: - return self.watcher_download_collection(prefix, tag, content_hash, owner_email) - - # -- Drive ----------------------------------------------------------------- - - 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 @@ -176,9 +125,7 @@ def owner_create_collection_folder( 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: + 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( @@ -188,7 +135,7 @@ def owner_tag_collection_as_any( ) ) - def owner_share_collection( + def grant_read( self, prefix: str, tag: str, content_hash: str, users: list[str] ) -> None: """Share collection folder with specific users via batch API.""" @@ -197,7 +144,7 @@ def owner_share_collection( folder_id = self._get_collection_folder_id(prefix, tag, content_hash) self._layout.client.batch_add_permissions(folder_id, users) - def owner_upload_collection_files( + def upload( self, prefix: str, tag: str, content_hash: str, files: dict[str, bytes] ) -> None: """Upload files to a collection folder.""" @@ -210,7 +157,7 @@ def owner_upload_collection_files( file_metadata = {"name": file_name, "parents": [folder_id]} self._layout.client.create_syft_file(file_metadata, file_payload) - def owner_list_collections(self, prefix: str) -> list[str]: + 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. @@ -230,7 +177,7 @@ def owner_list_collections(self, prefix: str) -> list[str]: result.append(cf.tag) return result - def owner_list_all_collections_with_permissions( + def list_mine_with_permissions( self, prefix: str, ) -> list[FileCollection]: @@ -263,9 +210,9 @@ def owner_list_all_collections_with_permissions( return collections - def owner_delete_collection(self, prefix: str, tag: str) -> None: + def delete(self, prefix: str, tag: str) -> None: """Delete every layout of ``tag`` published under ``prefix``.""" - collections = self.owner_list_all_collections_with_permissions(prefix) + collections = self.list_mine_with_permissions(prefix) for c in collections: if c.tag == tag: self._layout.client.delete_file_by_id(c.folder_id) @@ -277,7 +224,7 @@ def owner_delete_collection(self, prefix: str, tag: str) -> None: ).folder_name self.collection_folder_id_cache.pop(folder_name, None) - def watcher_list_collections(self, prefix: str) -> list[dict]: + 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 @@ -304,7 +251,7 @@ def watcher_list_collections(self, prefix: str) -> list[dict]: ) return result - def watcher_download_collection( + def download( self, prefix: str, tag: str, content_hash: str, owner_email: str ) -> dict[str, bytes]: """Download all files from a collection.""" @@ -328,7 +275,7 @@ def watcher_download_collection( return files - def watcher_get_collection_file_metadatas( + def list_files( self, prefix: str, tag: str, content_hash: str, owner_email: str ) -> List[Dict]: """Get file metadata from a collection without downloading.""" @@ -345,9 +292,9 @@ def watcher_get_collection_file_metadatas( file_metadatas = self._layout.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: + def fetch_file(self, handle: str) -> bytes: """Download a single file from a collection.""" - return self._layout.client.download_file(file_id) + return self._layout.client.download_file(handle) def _get_collection_folder_id( self, prefix: str, tag: str, content_hash: str diff --git a/tests/unit/test_dataset_collection_listing.py b/tests/unit/test_dataset_collection_listing.py index bb821a27588..f61d34da44d 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 @@ -28,7 +28,7 @@ def _conn(files): def test_a_valid_collection_is_returned(): conn = _conn([{"id": "f1", "name": VALID, "appProperties": {}}]) - got = conn.as_transport().collections.owner_list_all_collections_with_permissions( + got = conn.as_transport().collections.list_mine_with_permissions( DATASET_COLLECTION_PREFIX ) assert [(c.folder_id, c.tag, c.content_hash) for c in got] == [ @@ -47,7 +47,7 @@ def test_the_any_permission_flag_comes_from_app_properties(): } ] ) - got = conn.as_transport().collections.owner_list_all_collections_with_permissions( + got = conn.as_transport().collections.list_mine_with_permissions( DATASET_COLLECTION_PREFIX ) assert got[0].has_any_permission is True @@ -60,7 +60,7 @@ def test_a_name_the_client_cannot_parse_is_skipped(): {"id": "f1", "name": VALID, "appProperties": {}}, ] ) - got = conn.as_transport().collections.owner_list_all_collections_with_permissions( + got = conn.as_transport().collections.list_mine_with_permissions( DATASET_COLLECTION_PREFIX ) assert [c.folder_id for c in got] == ["f1"] @@ -71,6 +71,6 @@ def test_a_missing_name_field_raises(): # without a message. conn = _conn([{"id": "f1", "appProperties": {}}]) with pytest.raises(KeyError): - conn.as_transport().collections.owner_list_all_collections_with_permissions( + conn.as_transport().collections.list_mine_with_permissions( DATASET_COLLECTION_PREFIX ) From 6ce25769154ee0065cfe7e042a61e2e1f54233e3 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Thu, 29 Oct 2026 21:49:28 -0300 Subject: [PATCH 26/28] refactor: give each DrivePeering operation only its contract name --- syft/sync/connections/drive/drive_peering.py | 123 ++++++------------- tests/unit/test_peers_json_version.py | 2 +- 2 files changed, 37 insertions(+), 88 deletions(-) diff --git a/syft/sync/connections/drive/drive_peering.py b/syft/sync/connections/drive/drive_peering.py index ea3fd8f985c..00d6ee48643 100644 --- a/syft/sync/connections/drive/drive_peering.py +++ b/syft/sync/connections/drive/drive_peering.py @@ -66,59 +66,14 @@ def reset_caches(self) -> None: self._encryption_bundles_folder_id = None self._peers_json_cache = None - # -- Peering --------------------------------------------------------------- - def whoami(self) -> str: - return self.get_authenticated_email() - - def open_channels(self, peer: str) -> None: - self.add_peer(peer_email=peer) - - def incoming_requests(self) -> list[str]: - return self.get_peer_requests() - - def read_peer_registry(self, force_download: bool = False) -> dict: - return self._get_peers_json(force_download=force_download) - - def update_peer_registry( - self, peer: str, state: str, public_encryption_bundle: dict | None = None - ) -> None: - self._update_peer_state(peer, state, public_encryption_bundle) - - def publish_key_bundle(self, peer: str, bundle_json: str) -> None: - self.write_encryption_bundle(peer, bundle_json) - - def grant_key_bundle_read(self, peer: str) -> None: - self.share_encryption_bundles_folder(peer) - - def read_own_key_bundle(self, peer: str) -> str | None: - return self.read_own_encryption_bundle(peer) - - def read_peer_key_bundle(self, peer: str) -> str | None: - return self.read_peer_encryption_bundle(peer) - - def publish_peer_info(self, peer_info: PeerInfo) -> None: - self.write_peer_info_file(peer_info) - - def read_own_peer_info(self) -> PeerInfo | None: - return self.read_own_peer_info_file() - - def read_peer_info(self, peer: str) -> PeerInfo | None: - return self.read_peer_info_file(peer) - - def grant_peer_info_read(self, peer: str, peer_info: PeerInfo) -> None: - self.share_peer_info_file_with_peer(peer, peer_info) - - # -- Drive ----------------------------------------------------------------- - - def get_authenticated_email(self) -> str: """Return the email of the Google account behind the Drive service.""" about = execute_with_retries( self._layout.client.drive_service.about().get(fields="user(emailAddress)") ) return about["user"]["emailAddress"] - def create_peer_datasite_folders(self, peer_email: str): + def open_channels(self, peer: str): """Create inbox + outbox folders for peer's datasite on own drive, share with peer. Creates: @@ -126,20 +81,16 @@ def create_peer_datasite_folders(self, peer_email: str): - syft_datasite_{peer}_outbox_{self} — peer pushes events to me here """ # Inbox: I write proposals, peer reads - inbox_id = self._layout.get_peer_datasite_inbox_id(peer_email) + inbox_id = self._layout.get_peer_datasite_inbox_id(peer) if inbox_id is None: - inbox_id = self._layout.create_peer_datasite_folder(peer_email, "inbox") - self._layout.client.add_permission(inbox_id, peer_email, write=True) + inbox_id = self._layout.create_peer_datasite_folder(peer, "inbox") + self._layout.client.add_permission(inbox_id, peer, write=True) # Outbox: peer writes events, I read - outbox_id = self._layout.get_peer_datasite_outbox_id(peer_email) + outbox_id = self._layout.get_peer_datasite_outbox_id(peer) if outbox_id is None: - outbox_id = self._layout.create_peer_datasite_folder(peer_email, "outbox") - self._layout.client.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) + outbox_id = self._layout.create_peer_datasite_folder(peer, "outbox") + self._layout.client.add_permission(outbox_id, peer, write=True) def _get_peers_file_id(self) -> str | None: """Find SYFT_peers.json file in /SyftBox folder""" @@ -171,7 +122,7 @@ def _download_peers_json(self) -> dict[str, dict[str, Any]]: except ValueError as e: raise ValueError(f"Could not read the peers file: {e}") from e - def _get_peers_json( + def read_peer_registry( self, force_download: bool = False ) -> dict[str, dict[str, Any]]: """Return peers JSON, using the in-memory cache when available.""" @@ -212,9 +163,9 @@ def _write_peers_json(self, peers_data: dict[str, dict[str, Any]]): self._peers_json_cache = peers_data return file_id - def _update_peer_state( + def update_peer_registry( self, - peer_email: str, + peer: str, state: str, public_encryption_bundle: dict | None = None, ): @@ -222,15 +173,15 @@ def _update_peer_state( Preserves existing fields when updating state. """ - peers_data = self._get_peers_json() - existing = peers_data.get(peer_email, {}) + 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_email] = existing + peers_data[peer] = existing self._write_peers_json(peers_data) - def get_peer_requests(self) -> list[str]: + 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 @@ -266,7 +217,7 @@ def get_peer_requests(self) -> list[str]: if folder.datasite_email == self._layout.email: all_folder_peers.add(folder.peer_email) - peers_data = self._get_peers_json() + 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 [ @@ -288,7 +239,7 @@ def _get_own_peer_info_file_id(self) -> Optional[str]: items = results.get("files", []) return items[0]["id"] if items else None - def write_peer_info_file(self, peer_info: "PeerInfo") -> 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._layout.get_syftbox_folder_id() @@ -324,7 +275,7 @@ def _get_peer_info_file_id_of(self, peer_email: str) -> Optional[str]: items = results.get("files", []) return items[0]["id"] if items else None - def read_own_peer_info_file(self) -> Optional["PeerInfo"]: + def read_own_peer_info(self) -> Optional["PeerInfo"]: """Read the own peer info file from /SyftBox. Returns None when there is no file. @@ -337,7 +288,7 @@ def read_own_peer_info_file(self) -> Optional["PeerInfo"]: self._get_own_peer_info_file_id, "the own peer info file" ) - def read_peer_info_file(self, peer_email: str) -> Optional["PeerInfo"]: + 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. @@ -347,8 +298,8 @@ def read_peer_info_file(self, peer_email: str) -> Optional["PeerInfo"]: 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}", + lambda: self._get_peer_info_file_id_of(peer), + f"the peer info file of {peer}", ) def _read_peer_info_file( @@ -371,20 +322,18 @@ def _read_peer_info_file( 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: + 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.write_peer_info_file(peer_info) + self.publish_peer_info(peer_info) file_id = self._get_own_peer_info_file_id() if file_id: - self._layout.client.add_permission(file_id, peer_email, write=False) + self._layout.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.""" @@ -416,10 +365,10 @@ def _get_or_create_encryption_bundles_folder_id(self) -> str: 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: + 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._layout.email, peer_email) + filename = self._encryption_bundle_filename(self._layout.email, peer) file_payload, _ = self._layout.client.create_file_payload(bundle_json) # Check if file already exists @@ -439,20 +388,20 @@ def write_encryption_bundle(self, peer_email: str, bundle_json: str) -> None: file_metadata = {"name": filename, "parents": [folder_id]} self._layout.client.create_syft_file(file_metadata, file_payload) - def share_encryption_bundles_folder(self, peer_email: str) -> None: + 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._layout.client.add_permission(folder_id, peer_email, write=False) + self._layout.client.add_permission(folder_id, peer, write=False) - def read_own_encryption_bundle(self, peer_email: str) -> str | None: - """Read the bundle this client published for ``peer_email``, or None. + 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._layout.email, peer_email) + filename = self._encryption_bundle_filename(self._layout.email, peer) query = f"name='{filename}' and '{folder_id}' in parents and trashed=false" results = execute_with_retries( self._layout.client.drive_service.files().list(q=query, fields="files(id)") @@ -462,14 +411,14 @@ def read_own_encryption_bundle(self, peer_email: str) -> str | None: return None return self._layout.client.download_file(items[0]["id"]).decode("utf-8") - def read_peer_encryption_bundle(self, peer_email: str) -> str | None: + def read_peer_key_bundle(self, peer: str) -> str | None: """Read encryption bundle that peer wrote for us. - Searches for: encryption_bundle_{peer_email}_for_{self._layout.email}.json + Searches for: encryption_bundle_{peer}_for_{self._layout.email}.json in peer's bundles folder. """ - filename = self._encryption_bundle_filename(peer_email, self._layout.email) - query = f"name='{filename}' and trashed=false and '{peer_email}' in owners" + filename = self._encryption_bundle_filename(peer, self._layout.email) + query = f"name='{filename}' and trashed=false and '{peer}' in owners" results = execute_with_retries( self._layout.client.drive_service.files().list(q=query, fields="files(id)") ) @@ -479,10 +428,10 @@ def read_peer_encryption_bundle(self, peer_email: str) -> str | None: try: data = self._layout.client.download_file(items[0]["id"]) except Exception as e: - print(f"Warning: could not download the bundle of {peer_email}: {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_email}: {e}") + print(f"Warning: could not read the bundle of {peer}: {e}") return None diff --git a/tests/unit/test_peers_json_version.py b/tests/unit/test_peers_json_version.py index f9e0fe282c5..87eb67bc55c 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. """ From 4b72270f00c98b1e5d3704922a0e1dac41e39bd9 Mon Sep 17 00:00:00 2001 From: Koen van der Veen Date: Fri, 30 Oct 2026 11:49:16 +0100 Subject: [PATCH 27/28] refactor: rename DriveLayout to DriveCachedClient --- .../tests/test_dataset_upload_private.py | 14 +- packages/syft-rds/tests/test_sync_manager.py | 28 +-- .../tests/test_version_mismatch_flow.py | 58 +++--- syft/gdrive_utils.py | 4 +- ...drive_layout.py => drive_cached_client.py} | 24 ++- .../connections/drive/drive_collections.py | 58 +++--- .../connections/drive/drive_connections.py | 12 +- .../connections/drive/drive_owner_store.py | 133 ++++++++------ syft/sync/connections/drive/drive_p2p.py | 60 ++++--- syft/sync/connections/drive/drive_peering.py | 118 ++++++------ .../connections/drive/gdrive_transport.py | 17 +- syft/sync/login_utils.py | 2 +- syft/sync/syftbox_manager.py | 12 +- syft/version.py | 2 +- .../test_gdrive_pagination.py | 6 +- tests/unit/test_checkpoints.py | 6 +- tests/unit/test_dataset_collection_listing.py | 12 +- tests/unit/test_delete_syftbox.py | 36 ++-- tests/unit/test_drive_connections.py | 169 +++++++++--------- tests/unit/test_encryption.py | 12 +- tests/unit/test_encryption_policy.py | 6 +- tests/unit/test_gdrive_client.py | 14 +- tests/unit/test_gdrive_pagination.py | 97 +++++----- tests/unit/test_login_email_check.py | 6 +- tests/unit/test_mock_drive_service.py | 8 +- tests/unit/test_outbox_compaction.py | 14 +- tests/unit/test_p2p_folder_lookup.py | 32 ++-- tests/unit/test_peer_request_owner.py | 6 +- tests/unit/test_peers_json_version.py | 10 +- tests/unit/test_rolling_state.py | 8 +- tests/unit/test_sender_binding.py | 11 +- tests/unit/test_shared_router.py | 26 +-- tests/unit/test_sync_manager.py | 8 +- tests/unit/test_version_negotiation.py | 4 +- tests/unit/test_versioned_folder_adopt.py | 75 ++++---- tests/unit/test_versioned_folder_lookup.py | 2 +- 36 files changed, 613 insertions(+), 497 deletions(-) rename syft/sync/connections/drive/{gdrive_layout.py => drive_cached_client.py} (96%) diff --git a/packages/syft-rds/tests/test_dataset_upload_private.py b/packages/syft-rds/tests/test_dataset_upload_private.py index d321b87ca18..a015c9bab6d 100644 --- a/packages/syft-rds/tests/test_dataset_upload_private.py +++ b/packages/syft-rds/tests/test_dataset_upload_private.py @@ -5,7 +5,7 @@ PRIVATE_DATASET_COLLECTION_PREFIX, ) from dataset_test_utils import create_tmp_dataset_files -from syft.sync.connections.drive.drive_connections import drive_layout_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of class TestDatasetUploadPrivate: @@ -143,9 +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_layout = drive_layout_of(ds_manager.sync_engine._connection_router.transport) + ds_cached_client = drive_cached_client_of( + ds_manager.sync_engine._connection_router.transport + ) results = ( - ds_layout.client.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)", @@ -155,9 +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_layout = drive_layout_of(do_manager.sync_engine._connection_router.transport) + do_cached_client = drive_cached_client_of( + do_manager.sync_engine._connection_router.transport + ) do_results = ( - do_layout.client.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 db29d8f245a..f7e496368e5 100644 --- a/packages/syft-rds/tests/test_sync_manager.py +++ b/packages/syft-rds/tests/test_sync_manager.py @@ -22,7 +22,7 @@ create_tmp_dataset_files, create_tmp_dataset_files_with_parquet, ) -from syft.sync.connections.drive.drive_connections import drive_layout_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,7 +233,7 @@ def test_dataset_only_mock_data_uploaded(): # Sync so DS receives the dataset ds_manager.sync() - files = drive_layout_of( + files = drive_cached_client_of( ds_manager.sync_engine._connection_router.transport ).client.drive_service._backing_store.files list(files.keys()) @@ -284,9 +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) - do_layout = drive_layout_of(do_manager.sync_engine._connection_router.transport) - inbox_folder_id = do_layout.get_own_datasite_inbox_id(ds_manager.email) - inbox_file_metadatas = do_layout.get_file_metadatas_from_folder(inbox_folder_id) + do_cached_client = drive_cached_client_of( + do_manager.sync_engine._connection_router.transport + ) + 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() @@ -1038,7 +1042,7 @@ 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 = drive_layout_of( + mock_backing_store = drive_cached_client_of( ds_manager.sync_engine._connection_router.transport ).client.drive_service._backing_store ds_folder = ds_manager.syftbox_folder @@ -1057,10 +1061,10 @@ def test_ds_dataset_cache_aware_sync(): ) # Replace mock backing store to share dataset collections - drive_layout_of( + drive_cached_client_of( ds_manager2.sync_engine._connection_router.transport ).client.drive_service._backing_store = mock_backing_store - drive_layout_of( + drive_cached_client_of( do_manager2.sync_engine._connection_router.transport ).client.drive_service._backing_store = mock_backing_store @@ -1138,7 +1142,7 @@ 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 = drive_layout_of( + mock_backing_store = drive_cached_client_of( ds_manager.sync_engine._connection_router.transport ).client.drive_service._backing_store ds_folder = ds_manager.syftbox_folder @@ -1157,10 +1161,10 @@ def test_do_dataset_cache_aware_sync(): ) # Replace mock backing store to share dataset collections - drive_layout_of( + drive_cached_client_of( ds_manager2.sync_engine._connection_router.transport ).client.drive_service._backing_store = mock_backing_store - drive_layout_of( + drive_cached_client_of( do_manager2.sync_engine._connection_router.transport ).client.drive_service._backing_store = mock_backing_store @@ -1205,7 +1209,7 @@ def test_in_memory_connection_load_state(): ) # Get the backing store that will persist across "restarts" - backing_store = drive_layout_of( + backing_store = drive_cached_client_of( do_manager1.sync_engine._connection_router.transport ).client.drive_service._backing_store ds_folder = ds_manager1.syftbox_folder diff --git a/packages/syft-rds/tests/test_version_mismatch_flow.py b/packages/syft-rds/tests/test_version_mismatch_flow.py index 83d16c140b0..f8be407a889 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -9,7 +9,9 @@ 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.gdrive_layout import GDRIVE_P2P_FOLDER_DATASITE_PREFIX +from syft.sync.connections.drive.drive_cached_client import ( + GDRIVE_P2P_FOLDER_DATASITE_PREFIX, +) from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.connections.drive.mock_drive_service import ( MockDriveService, @@ -19,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 drive_layout_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(layout, 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}'" @@ -34,23 +36,25 @@ def _find_p2p_folders(layout, peer_email): " and trashed=false" ) results = ( - layout.client.drive_service.files() + cached_client.client.drive_service.files() .list(q=q, fields="files(id, name)") .execute() ) return results.get("files", []) -def _find_versioned_p2p_folders(layout, 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(layout, 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.""" - layout = drive_layout_of(manager.peer_manager.connection_router.transport) - return layout.client.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( @@ -98,7 +102,7 @@ def _simulate_continue_on_mismatch(manager, backing_store): mock_conn = GDriveConnection.from_service(email, mock_service) def read_remote(e, t): - return DrivePeering(mock_conn.layout).read_own_peer_info().version_info + return DrivePeering(mock_conn.cached_client).read_own_peer_info().version_info with ( patch( @@ -173,15 +177,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_layout = drive_layout_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_layout, 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_layout.client.find_folder_by_name( + old_personal_id = do_cached_client.client.find_folder_by_name( old_personal_name, - parent_id=do_layout.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 @@ -199,7 +205,7 @@ def test_version_mismatch_continues_and_repairs(): NEW_VERSION, ), patch( - "syft.sync.connections.drive.gdrive_layout.SYFT_VERSION", + "syft.sync.connections.drive.drive_cached_client.SYFT_VERSION", NEW_VERSION, ), patch("syft.sync.login_utils.SYFT_VERSION", NEW_VERSION), @@ -216,21 +222,21 @@ def test_version_mismatch_continues_and_repairs(): ) # Personal folder is adopted (same Drive id, new name), not recreated. - do_layout_new = drive_layout_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_layout_new.client.find_folder_by_name( + new_personal_id = do_cached_client_new.client.find_folder_by_name( new_personal_name, - parent_id=do_layout_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_layout_new.client.find_folder_by_name( + do_cached_client_new.client.find_folder_by_name( old_personal_name, - parent_id=do_layout_new.get_syftbox_folder_id(), + parent_id=do_cached_client_new.get_syftbox_folder_id(), owner_email=do_email, ) is None @@ -254,13 +260,13 @@ def test_version_mismatch_continues_and_repairs(): syftbox_folder=ds_syftbox, ) - ds_layout_new = drive_layout_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_layout_new.client.find_folder_by_name( + ds_personal_id = ds_cached_client_new.client.find_folder_by_name( ds_personal_name, - parent_id=ds_layout_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 @@ -270,9 +276,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_layout_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_layout_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) diff --git a/syft/gdrive_utils.py b/syft/gdrive_utils.py index a304016272e..fce56384bdd 100644 --- a/syft/gdrive_utils.py +++ b/syft/gdrive_utils.py @@ -317,7 +317,7 @@ def delete_remote_syftbox( email = _resolve_email(email) token_path = _resolve_token_path(token_path) conn = GDriveConnection.from_token_path(email=email, token_path=token_path) - store = DriveOwnerStore(conn.layout) + store = DriveOwnerStore(conn.cached_client) # Gather file IDs via folder hierarchy folder_file_ids = set(store.gather_all_file_and_folder_ids()) @@ -328,7 +328,7 @@ def delete_remote_syftbox( all_file_ids = list(folder_file_ids | orphaned_file_ids) start = time.time() - conn.layout.client.delete_multiple_files_by_ids(all_file_ids) + conn.cached_client.client.delete_multiple_files_by_ids(all_file_ids) elapsed = time.time() - start if verbose: diff --git a/syft/sync/connections/drive/gdrive_layout.py b/syft/sync/connections/drive/drive_cached_client.py similarity index 96% rename from syft/sync/connections/drive/gdrive_layout.py rename to syft/sync/connections/drive/drive_cached_client.py index 8e64ee79d31..785db426de9 100644 --- a/syft/sync/connections/drive/gdrive_layout.py +++ b/syft/sync/connections/drive/drive_cached_client.py @@ -1,7 +1,8 @@ """Where syft keeps its folders in one account's Google Drive. -The layout 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. +``DriveCachedClient`` 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 @@ -175,7 +176,14 @@ def valid_events_from_file_metadatas( @dataclass -class DriveLayout: +class DriveCachedClient: + """The Drive client of one account, and where syft's folders are in it. + + It 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. + """ + client: GDriveClient email: str # /SyftBox: the top folder, with the inboxes, outboxes and personal folder. @@ -191,15 +199,15 @@ class DriveLayout: # sender email -> id of my archive folder for that sender. archive_folder_id_cache: dict[str, str] = field(default_factory=dict, init=False) - def clone_for_thread(self) -> "DriveLayout": - """A layout over a new Drive client, for use from another thread, that + 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 = DriveLayout(self.client.clone_for_thread(), self.email) + clone = DriveCachedClient(self.client.clone_for_thread(), self.email) self.copy_caches_to(clone) return clone - def copy_caches_to(self, other: "DriveLayout") -> None: - """Seed the folder ids of ``other`` from this layout.""" + 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) diff --git a/syft/sync/connections/drive/drive_collections.py b/syft/sync/connections/drive/drive_collections.py index 0004bdcd379..2919b2f00ef 100644 --- a/syft/sync/connections/drive/drive_collections.py +++ b/syft/sync/connections/drive/drive_collections.py @@ -12,7 +12,7 @@ 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.gdrive_layout import DriveLayout +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 @@ -88,20 +88,20 @@ def compute_hash(files: dict[str, bytes]) -> str: class DriveCollections(Collections): """Collection folders in /SyftBox, shared by Drive permission.""" - def __init__(self, layout: DriveLayout): - self._layout = layout + 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._layout.clone_for_thread()) + 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._layout.reset_caches() + 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: @@ -114,10 +114,10 @@ def create(self, prefix: str, tag: str, content_hash: str, owner_email: str) -> if folder_name in self.collection_folder_id_cache: return self.collection_folder_id_cache[folder_name] - syftbox_folder_id = self._layout.get_syftbox_folder_id() + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() # Check if exists - folder_id = self._layout.client.find_folder_by_name( + folder_id = self._cached_client.client.find_folder_by_name( folder_name, parent_id=syftbox_folder_id ) if folder_id: @@ -125,7 +125,9 @@ def create(self, prefix: str, tag: str, content_hash: str, owner_email: str) -> return folder_id # Create new folder - folder_id = self._layout.client.create_folder(folder_name, syftbox_folder_id) + 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 @@ -133,7 +135,7 @@ 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._layout.client.drive_service.files().update( + self._cached_client.client.drive_service.files().update( fileId=folder_id, body={"appProperties": {"syft_shared_with_any": "true"}}, ) @@ -146,7 +148,7 @@ def grant_read( if not users: return folder_id = self._get_collection_folder_id(prefix, tag, content_hash) - self._layout.client.batch_add_permissions(folder_id, users) + self._cached_client.client.batch_add_permissions(folder_id, users) def upload( self, prefix: str, tag: str, content_hash: str, files: dict[str, bytes] @@ -155,24 +157,24 @@ def upload( folder_id = self._get_collection_folder_id(prefix, tag, content_hash) for file_path, content in files.items(): - file_payload, _ = self._layout.client.create_file_payload(content) + 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._layout.client.create_syft_file(file_metadata, file_payload) + 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._layout.get_syftbox_folder_id() + 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._layout.client.list_all_files(query, fields="name"): + for folder in self._cached_client.client.list_all_files(query, fields="name"): try: cf = CollectionFolder.from_name(prefix, folder["name"]) except ValueError: @@ -186,13 +188,13 @@ def list_mine_with_permissions( prefix: str, ) -> list[FileCollection]: """List all DO's collections with permissions info, one per layout.""" - syftbox_folder_id = self._layout.get_syftbox_folder_id() + 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._layout.client.list_all_files( + for folder in self._cached_client.client.list_all_files( query, fields="id,name,appProperties" ): try: @@ -219,7 +221,7 @@ def delete(self, prefix: str, tag: str) -> None: collections = self.list_mine_with_permissions(prefix) for c in collections: if c.tag == tag: - self._layout.client.delete_file_by_id(c.folder_id) + self._cached_client.client.delete_file_by_id(c.folder_id) folder_name = CollectionFolder( prefix=prefix, tag=c.tag, @@ -229,7 +231,7 @@ def delete(self, prefix: str, tag: str) -> None: self.collection_folder_id_cache.pop(folder_name, None) def delete_by_handle(self, handle: str) -> None: - self._layout.client.delete_file_by_id(handle) + 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() @@ -246,7 +248,9 @@ def list_shared_with_me(self, prefix: str) -> list[dict]: f"and trashed=false and mimeType='{GOOGLE_FOLDER_MIME_TYPE}'" ) result = [] - for folder in self._layout.client.list_all_files(query, fields="name, owners"): + for folder in self._cached_client.client.list_all_files( + query, fields="name, owners" + ): try: cf = CollectionFolder.from_name(prefix, folder["name"]) except ValueError: @@ -271,19 +275,19 @@ def download( 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._layout.client.find_folder_by_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._layout.get_file_metadatas_from_folder(folder_id) + 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._layout.client.download_file(file_id) + files[file_name] = self._cached_client.client.download_file(file_id) return files @@ -294,19 +298,19 @@ def list_files( folder_name = CollectionFolder( prefix=prefix, tag=tag, content_hash=content_hash ).folder_name - folder_id = self._layout.client.find_folder_by_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._layout.get_file_metadatas_from_folder(folder_id) + 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._layout.client.download_file(handle) + return self._cached_client.client.download_file(handle) def _get_collection_folder_id( self, prefix: str, tag: str, content_hash: str @@ -318,8 +322,8 @@ def _get_collection_folder_id( if folder_name in self.collection_folder_id_cache: return self.collection_folder_id_cache[folder_name] - syftbox_folder_id = self._layout.get_syftbox_folder_id() - folder_id = self._layout.client.find_folder_by_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 ) diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index dc3e358d7fa..93ec9b1051e 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -1,12 +1,12 @@ """The transport components on Google Drive. Each component has its own module: ``drive_p2p``, ``drive_owner_store``, -``drive_collections`` and ``drive_peering``. They share one ``DriveLayout``, +``drive_collections`` and ``drive_peering``. They share one ``DriveCachedClient``, which holds the Drive client. ``SyftboxManager._drive_transport`` builds them. """ from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore -from syft.sync.connections.drive.gdrive_layout import DriveLayout +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.transport.protocols import Transport, WatcherStore @@ -30,11 +30,11 @@ def reset(self, peer: str) -> None: raise NotImplementedError("Watcher storage on Drive is not built yet") -def drive_layout_of(transport: Transport) -> DriveLayout: - """The Drive layout under ``transport``, for tests and tools that need - Drive itself. The sync code reaches Drive only through the components. +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, DriveOwnerStore): raise TypeError("This transport is not built on Google Drive") - return store._layout + return store._cached_client diff --git a/syft/sync/connections/drive/drive_owner_store.py b/syft/sync/connections/drive/drive_owner_store.py index 88f5a12fd73..b0e20170951 100644 --- a/syft/sync/connections/drive/drive_owner_store.py +++ b/syft/sync/connections/drive/drive_owner_store.py @@ -18,9 +18,9 @@ SYFT_FILE_PROPERTY_KEY, SYFT_FILE_PROPERTY_VALUE, ) -from syft.sync.connections.drive.gdrive_layout import ( +from syft.sync.connections.drive.drive_cached_client import ( SYFTBOX_FOLDER, - DriveLayout, + DriveCachedClient, valid_events_from_file_metadatas, ) from syft.sync.connections.drive.gdrive_retry import ( @@ -54,8 +54,8 @@ def _filter_valid_file_metadatas( class DriveOwnerStore(OwnerStore): """The DO's own folders in /SyftBox: event log, checkpoints, rolling state.""" - def __init__(self, layout: DriveLayout): - self._layout = layout + 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 @@ -65,7 +65,7 @@ def __init__(self, layout: DriveLayout): 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._layout.clone_for_thread()) + 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( @@ -74,32 +74,36 @@ def clone_for_thread(self) -> "DriveOwnerStore": return clone def reset_caches(self) -> None: - self._layout.reset_caches() + 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._layout.get_personal_syftbox_folder_id() + personal_syftbox_folder_id = ( + self._cached_client.get_personal_syftbox_folder_id() + ) file_metadata = { "name": name, "parents": [personal_syftbox_folder_id], } - file_payload, _ = self._layout.client.create_file_payload(data) + file_payload, _ = self._cached_client.client.create_file_payload(data) - res = self._layout.client.create_syft_file(file_metadata, file_payload) + 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._layout.client.download_file(handle) + 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._layout.get_personal_syftbox_folder_id() - file_metadatas = self._layout.get_file_metadatas_from_folder( + 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) @@ -114,8 +118,10 @@ def list_event_handles(self, since_timestamp: float | None = None) -> list[str]: @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._layout.get_personal_syftbox_folder_id() - file_metadatas = self._layout.get_file_metadatas_from_folder( + 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) @@ -125,7 +131,7 @@ def fetch_all_events(self) -> list[bytes]: gdrive_id = [ f for f in file_metadatas if f["name"] == fname_obj.as_string() ][0]["id"] - result.append(self._layout.client.download_file(gdrive_id)) + result.append(self._cached_client.client.download_file(gdrive_id)) return result @raise_unavailable("Could not read the event log from Drive.") @@ -133,8 +139,8 @@ 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._layout.get_personal_syftbox_folder_id() - file_metadatas = self._layout.get_file_metadatas_from_folder( + 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 ) @@ -149,21 +155,21 @@ def fetch_events_since(self, since_timestamp: float | None) -> list[bytes]: f for f in file_metadatas if f["name"] == fname_obj.as_string() ][0]["id"] - result.append(self._layout.client.download_file(gdrive_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._layout.email}-{SYFT_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._layout.client.find_folders( - name_contains=[f"{self._layout.email}-", "-checkpoints"], - parent_id=self._layout.get_syftbox_folder_id(), + 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._layout.find_or_adopt_versioned_folder( + return self._cached_client.find_or_adopt_versioned_folder( folders, current_name=self._get_checkpoints_folder_name() ) @@ -174,8 +180,8 @@ def _get_or_create_checkpoints_folder_id(self) -> str: return folder_id # Create the folder folder_name = self._get_checkpoints_folder_name() - syftbox_folder_id = self._layout.get_syftbox_folder_id() - return self._layout.client.create_folder(folder_name, syftbox_folder_id) + 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. @@ -183,14 +189,14 @@ def put_checkpoint(self, name: str, data: bytes) -> str: Uploads the new checkpoint first, then deletes old ones. """ folder_id = self._get_or_create_checkpoints_folder_id() - payload, _ = self._layout.client.create_file_payload(data) + payload, _ = self._cached_client.client.create_file_payload(data) file_metadata = { "name": name, "parents": [folder_id], } - result = self._layout.client.create_syft_file(file_metadata, payload) + 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")) @@ -209,7 +215,7 @@ def latest_checkpoint(self) -> bytes | None: f"and name contains '{CHECKPOINT_FILENAME_PREFIX}'" ) results = execute_with_retries( - self._layout.client.drive_service.files().list( + self._cached_client.client.drive_service.files().list( q=query, fields="files(id, name)" ) ) @@ -229,7 +235,7 @@ def latest_checkpoint(self) -> bytes | None: if latest_file is None: return None - return self._layout.client.download_file(latest_file["id"]) + 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). @@ -248,7 +254,7 @@ def _delete_all_checkpoints(self, exclude_file_id: str | None = None): f"and name contains '{CHECKPOINT_FILENAME_PREFIX}'" ) results = execute_with_retries( - self._layout.client.drive_service.files().list( + self._cached_client.client.drive_service.files().list( q=query, fields="files(id, name)" ) ) @@ -262,7 +268,9 @@ def _delete_all_checkpoints(self, exclude_file_id: str | None = None): continue # Skip the newly uploaded checkpoint try: execute_with_retries( - self._layout.client.drive_service.files().delete(fileId=item["id"]) + 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}") @@ -270,14 +278,14 @@ def _delete_all_checkpoints(self, exclude_file_id: str | None = None): 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._layout.client.create_file_payload(data) + payload, _ = self._cached_client.client.create_file_payload(data) file_metadata = { "name": name, "parents": [folder_id], } - result = self._layout.client.create_syft_file(file_metadata, payload) + 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.") @@ -292,7 +300,7 @@ def all_incrementals(self) -> list[bytes]: f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" ) results = execute_with_retries( - self._layout.client.drive_service.files().list( + self._cached_client.client.drive_service.files().list( q=query, fields="files(id, name)" ) ) @@ -301,7 +309,7 @@ def all_incrementals(self) -> list[bytes]: if not items: return [] - return [self._layout.client.download_file(item["id"]) for item in items] + 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.""" @@ -314,7 +322,9 @@ def incremental_count(self) -> int: f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" ) results = execute_with_retries( - self._layout.client.drive_service.files().list(q=query, fields="files(id)") + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) ) return len(results.get("files", [])) @@ -329,7 +339,7 @@ def next_incremental_sequence(self) -> int: f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" ) results = execute_with_retries( - self._layout.client.drive_service.files().list( + self._cached_client.client.drive_service.files().list( q=query, fields="files(name)" ) ) @@ -358,7 +368,7 @@ def delete_incrementals(self) -> None: f"and name contains '{INCREMENTAL_CHECKPOINT_PREFIX}'" ) results = execute_with_retries( - self._layout.client.drive_service.files().list( + self._cached_client.client.drive_service.files().list( q=query, fields="files(id, name)" ) ) @@ -367,7 +377,9 @@ def delete_incrementals(self) -> None: for item in items: try: execute_with_retries( - self._layout.client.drive_service.files().delete(fileId=item["id"]) + self._cached_client.client.drive_service.files().delete( + fileId=item["id"] + ) ) except Exception as e: print( @@ -376,7 +388,7 @@ def delete_incrementals(self) -> None: def _get_rolling_state_folder_name(self) -> str: """Get the rolling state folder name: {email}-{version}-rolling-state""" - return f"{self._layout.email}-{SYFT_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: """ @@ -388,11 +400,11 @@ def _get_rolling_state_folder_id(self, use_cache: bool = True) -> str | None: if use_cache and self._rolling_state_folder_id is not None: return self._rolling_state_folder_id - folders = self._layout.client.find_folders( - name_contains=[f"{self._layout.email}-", "-rolling-state"], - parent_id=self._layout.get_syftbox_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._layout.find_or_adopt_versioned_folder( + 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: @@ -406,8 +418,10 @@ def _get_or_create_rolling_state_folder_id(self) -> str: return folder_id # Create the folder folder_name = self._get_rolling_state_folder_name() - syftbox_folder_id = self._layout.get_syftbox_folder_id() - folder_id = self._layout.client.create_folder(folder_name, syftbox_folder_id) + 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 @@ -416,13 +430,13 @@ def put_rolling_state(self, name: str, data: bytes) -> str: Optimized to use update() if file ID is cached. """ - payload, _ = self._layout.client.create_file_payload(data) + 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._layout.client.drive_service.files().update( + self._cached_client.client.drive_service.files().update( fileId=self._rolling_state_file_id, media_body=payload, ) @@ -444,7 +458,7 @@ def put_rolling_state(self, name: str, data: bytes) -> str: "parents": [folder_id], } - result = self._layout.client.create_syft_file(file_metadata, payload) + 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 @@ -464,7 +478,7 @@ def get_rolling_state(self) -> bytes | None: f"and name contains '{ROLLING_STATE_FILENAME_PREFIX}'" ) results = execute_with_retries( - self._layout.client.drive_service.files().list( + self._cached_client.client.drive_service.files().list( q=query, fields="files(id, name)" ) ) @@ -484,7 +498,7 @@ def get_rolling_state(self) -> bytes | None: if latest_file is None: return None - data = self._layout.client.download_file(latest_file["id"]) + data = self._cached_client.client.download_file(latest_file["id"]) self._rolling_state_file_id = latest_file["id"] return data @@ -503,7 +517,7 @@ def delete_rolling_state(self) -> None: f"and name contains '{ROLLING_STATE_FILENAME_PREFIX}'" ) results = execute_with_retries( - self._layout.client.drive_service.files().list( + self._cached_client.client.drive_service.files().list( q=query, fields="files(id, name)" ) ) @@ -513,7 +527,9 @@ def delete_rolling_state(self) -> None: for item in items: try: execute_with_retries( - self._layout.client.drive_service.files().delete(fileId=item["id"]) + 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}") @@ -525,14 +541,14 @@ def delete_all(self) -> int: 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._layout.client.delete_multiple_files_by_ids(list(file_ids)) + 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._layout.get_syftbox_folder_id() + syftbox_folder_id = self._cached_client.get_syftbox_folder_id() return gather_all_file_and_folder_ids_recursive( - self._layout.client.drive_service, syftbox_folder_id + self._cached_client.client.drive_service, syftbox_folder_id ) def find_orphaned_syft_files(self) -> list[str]: @@ -549,4 +565,7 @@ def find_orphaned_syft_files(self) -> list[str]: f"properties has {{ key='{SYFT_FILE_PROPERTY_KEY}' " f"and value='{SYFT_FILE_PROPERTY_VALUE}' }}" ) - return [f["id"] for f in self._layout.client.list_all_files(query, fields="id")] + 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 index b0ef0a697d7..9132e94cbbc 100644 --- a/syft/sync/connections/drive/drive_p2p.py +++ b/syft/sync/connections/drive/drive_p2p.py @@ -5,8 +5,8 @@ from concurrent.futures import ThreadPoolExecutor from syft.sync.connections.drive.gdrive_client import DRIVE_BATCH_SIZE, GDriveClient -from syft.sync.connections.drive.gdrive_layout import ( - DriveLayout, +from syft.sync.connections.drive.drive_cached_client import ( + DriveCachedClient, valid_events_from_file_metadatas, ) from syft.sync.connections.drive.gdrive_retry import ( @@ -74,15 +74,15 @@ class DriveP2PConnection(P2PConnection): it received in its own storage. """ - def __init__(self, layout: DriveLayout): - self._layout = layout + 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._layout.clone_for_thread()) + return DriveP2PConnection(self._cached_client.clone_for_thread()) def reset_caches(self) -> None: - self._layout.reset_caches() + self._cached_client.reset_caches() def send( self, peer: str, msg_id: str, data: bytes, kind: MessageKind @@ -131,12 +131,12 @@ def _download_all(self, listed: list[dict]) -> list[bytes | None]: 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._layout.client, m) for m in listed] + 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._layout.client.clone_for_thread() + per_thread.client = self._cached_client.client.clone_for_thread() return _download_or_skip(per_thread.client, m) workers = min(DOWNLOAD_WORKERS, len(listed)) @@ -163,7 +163,7 @@ def callback(request_id, response, exception): f"Could not archive inbox file {request_id}: {exception!r}" ) - drive_service = self._layout.client.drive_service + 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]: @@ -175,7 +175,7 @@ def callback(request_id, response, exception): 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._layout.client.download_file(m["file_id"]) + self._cached_client.client.download_file(m["file_id"]) for m in self.outbox_message_metadatas(peer, since_timestamp) ] @@ -184,11 +184,11 @@ def outbox_message_metadatas( ) -> list[dict]: """Like ``outbox_messages``, without downloading: each entry's ``file_id`` is the handle to ``fetch_message`` it with.""" - folder_id = self._layout.get_peer_datasite_outbox_id(peer) + folder_id = self._cached_client.get_peer_datasite_outbox_id(peer) if folder_id is None: return [] - file_metadatas = self._layout.get_file_metadatas_from_folder( + 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) @@ -209,30 +209,32 @@ def outbox_message_metadatas( return result def fetch_message(self, handle: str) -> bytes: - return self._layout.client.download_file(handle) + 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._layout.get_own_datasite_outbox_id(peer) + folder_id = self._cached_client.get_own_datasite_outbox_id(peer) if folder_id is None: return [] - return self._layout.get_file_metadatas_from_folder( + return self._cached_client.get_file_metadatas_from_folder( folder_id, since_timestamp=None ) def delete_messages(self, handles: list[str]) -> None: - self._layout.client.delete_multiple_files_by_ids(handles) + 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._layout.get_own_datasite_inbox_id(peer) + 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._layout.get_file_metadatas_from_folder(inbox_folder_id) + 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()} @@ -247,10 +249,10 @@ def _archive_request(self, file_id: str, sender: str): 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._layout.client.drive_service.files().update( + return self._cached_client.client.drive_service.files().update( fileId=file_id, - addParents=self._layout.get_archive_folder_id(sender), - removeParents=self._layout.get_own_datasite_inbox_id(sender), + addParents=self._cached_client.get_archive_folder_id(sender), + removeParents=self._cached_client.get_own_datasite_inbox_id(sender), fields="id", supportsAllDrives=True, ) @@ -258,14 +260,14 @@ def _archive_request(self, file_id: str, sender: str): def watcher_send_raw_bytes_to_inbox( self, recipient: str, filename: str, data: bytes ) -> None: - inbox_id = self._layout.get_peer_datasite_inbox_id(recipient) + 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._layout.client.create_file_payload(data) + payload, _ = self._cached_client.client.create_file_payload(data) file_metadata = {"name": filename, "parents": [inbox_id]} - result = self._layout.client.create_syft_file( + result = self._cached_client.client.create_syft_file( file_metadata, payload, fields="id, parents" ) @@ -278,7 +280,7 @@ def watcher_send_raw_bytes_to_inbox( f"Moving file to correct folder..." ) execute_with_retries( - self._layout.client.drive_service.files().update( + self._cached_client.client.drive_service.files().update( fileId=file_id, addParents=inbox_id, fields="id, parents", @@ -288,14 +290,14 @@ def watcher_send_raw_bytes_to_inbox( def owner_write_raw_bytes_to_outbox( self, recipient: str, filename: str, data: bytes ) -> None: - outbox_folder_id = self._layout.get_own_datasite_outbox_id(recipient) + 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._layout.client.create_file_payload(data) + file_payload, _ = self._cached_client.client.create_file_payload(data) file_metadata = {"name": filename, "parents": [outbox_folder_id]} - result = self._layout.client.create_syft_file( + result = self._cached_client.client.create_syft_file( file_metadata, file_payload, fields="id, parents" ) @@ -308,7 +310,7 @@ def owner_write_raw_bytes_to_outbox( f"Moving file to correct folder..." ) execute_with_retries( - self._layout.client.drive_service.files().update( + 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 index 80a04dca551..7884de0f838 100644 --- a/syft/sync/connections/drive/drive_peering.py +++ b/syft/sync/connections/drive/drive_peering.py @@ -13,9 +13,9 @@ from syft_migration import MigrationError from syft.sync.connections.drive.gdrive_client import GOOGLE_FOLDER_MIME_TYPE -from syft.sync.connections.drive.gdrive_layout import ( +from syft.sync.connections.drive.drive_cached_client import ( GDRIVE_P2P_FOLDER_DATASITE_PREFIX, - DriveLayout, + DriveCachedClient, GdriveP2PFolder, ) from syft.sync.connections.drive.gdrive_retry import ( @@ -47,8 +47,8 @@ class DrivePeering(Peering): """Identity from the Google account; channels, the peer list, key bundles and version files in /SyftBox.""" - def __init__(self, layout: DriveLayout): - self._layout = layout + 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"}, @@ -58,7 +58,7 @@ def __init__(self, layout: DriveLayout): def clone_for_thread(self) -> "DrivePeering": """Peering over a new Drive client, for use from another thread, with these caches.""" - clone = DrivePeering(self._layout.clone_for_thread()) + 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 @@ -66,14 +66,16 @@ def clone_for_thread(self) -> "DrivePeering": return clone def reset_caches(self) -> None: - self._layout.reset_caches() + 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._layout.client.drive_service.about().get(fields="user(emailAddress)") + self._cached_client.client.drive_service.about().get( + fields="user(emailAddress)" + ) ) return about["user"]["emailAddress"] @@ -85,23 +87,25 @@ def open_channels(self, peer: str): - syft_datasite_{peer}_outbox_{self} — peer pushes events to me here """ # Inbox: I write proposals, peer reads - inbox_id = self._layout.get_peer_datasite_inbox_id(peer) + inbox_id = self._cached_client.get_peer_datasite_inbox_id(peer) if inbox_id is None: - inbox_id = self._layout.create_peer_datasite_folder(peer, "inbox") - self._layout.client.add_permission(inbox_id, peer, write=True) + 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._layout.get_peer_datasite_outbox_id(peer) + outbox_id = self._cached_client.get_peer_datasite_outbox_id(peer) if outbox_id is None: - outbox_id = self._layout.create_peer_datasite_folder(peer, "outbox") - self._layout.client.add_permission(outbox_id, peer, write=True) + 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._layout.get_syftbox_folder_id() + 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._layout.client.drive_service.files().list(q=query, fields="files(id)") + 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 @@ -120,7 +124,7 @@ def _download_peers_json(self) -> dict[str, dict[str, Any]]: file_id = self._get_peers_file_id() if file_id is None: return {} - file_data = self._layout.client.download_file(file_id) + file_data = self._cached_client.client.download_file(file_id) try: return json.loads(file_data.decode("utf-8")) except ValueError as e: @@ -141,12 +145,12 @@ def _write_peers_json(self, peers_data: dict[str, dict[str, Any]]): **peers_data, PEERS_META_KEY: {"version": SYFT_PEERS_VERSION}, } - syftbox_folder_id = self._layout.get_syftbox_folder_id() + 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._layout.client.create_file_payload(json_data) + file_payload, _ = self._cached_client.client.create_file_payload(json_data) if file_id is None: # Create new file @@ -154,13 +158,15 @@ def _write_peers_json(self, peers_data: dict[str, dict[str, Any]]): "name": SYFT_PEERS_FILE, "parents": [syftbox_folder_id], } - result = self._layout.client.create_syft_file(file_metadata, file_payload) + 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._layout.client.drive_service.files().update( + self._cached_client.client.drive_service.files().update( fileId=file_id, media_body=file_payload ) ) @@ -192,9 +198,9 @@ def incoming_requests(self) -> list[str]: peers who created folders for our datasite. Filters out already-accepted or rejected peers from SYFT_peers.json. """ - folders = self._layout.client.list_all_files( + folders = self._cached_client.client.list_all_files( f"name contains '{GDRIVE_P2P_FOLDER_DATASITE_PREFIX}#' " - f"and name contains '#{self._layout.email}#' " + 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", @@ -218,7 +224,7 @@ def incoming_requests(self) -> list[str]: != folder.peer_email.casefold() ): continue - if folder.datasite_email == self._layout.email: + if folder.datasite_email == self._cached_client.email: all_folder_peers.add(folder.peer_email) peers_data = self.read_peer_registry() @@ -235,10 +241,12 @@ def incoming_requests(self) -> list[str]: 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._layout.get_syftbox_folder_id() + 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._layout.client.drive_service.files().list(q=query, fields="files(id)") + 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 @@ -246,12 +254,12 @@ def _get_own_peer_info_file_id(self) -> Optional[str]: 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._layout.get_syftbox_folder_id() + 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._layout.client.create_file_payload(json_data) + file_payload, _ = self._cached_client.client.create_file_payload(json_data) if file_id is None: # Create new file @@ -259,11 +267,11 @@ def publish_peer_info(self, peer_info: "PeerInfo") -> None: "name": SYFT_VERSION_FILE, "parents": [syftbox_folder_id], } - self._layout.client.create_syft_file(file_metadata, file_payload) + self._cached_client.client.create_syft_file(file_metadata, file_payload) else: # Update existing file execute_with_retries( - self._layout.client.drive_service.files().update( + self._cached_client.client.drive_service.files().update( fileId=file_id, media_body=file_payload ) ) @@ -274,7 +282,9 @@ def _get_peer_info_file_id_of(self, peer_email: str) -> Optional[str]: f"name='{SYFT_VERSION_FILE}' and trashed=false and '{peer_email}' in owners" ) results = execute_with_retries( - self._layout.client.drive_service.files().list(q=query, fields="files(id)") + 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 @@ -320,7 +330,7 @@ def _read_peer_info_file( file_id = find_file_id() if file_id is None: return None - file_data = self._layout.client.download_file(file_id) + 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: @@ -337,17 +347,17 @@ def grant_peer_info_read(self, peer: str, peer_info: "PeerInfo") -> None: file_id = self._get_own_peer_info_file_id() if file_id: - self._layout.client.add_permission(file_id, peer, write=False) + 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._layout.email + email=self._cached_client.email ).as_string() - folder_id = self._layout.client.find_folder_by_name( - folder_name, parent_id=self._layout.get_syftbox_folder_id() + 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 @@ -358,10 +368,10 @@ def _get_or_create_encryption_bundles_folder_id(self) -> str: folder_id = self._find_encryption_bundles_folder_id() if not folder_id: folder_name = GdriveEncryptionBundlesFolder( - email=self._layout.email + email=self._cached_client.email ).as_string() - folder_id = self._layout.client.create_folder( - folder_name, self._layout.get_syftbox_folder_id() + 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 @@ -372,30 +382,32 @@ def _encryption_bundle_filename(self, owner_email: str, peer_email: str) -> str: 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._layout.email, peer) - file_payload, _ = self._layout.client.create_file_payload(bundle_json) + 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._layout.client.drive_service.files().list(q=query, fields="files(id)") + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) ) items = results.get("files", []) if items: execute_with_retries( - self._layout.client.drive_service.files().update( + 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._layout.client.create_syft_file(file_metadata, file_payload) + 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._layout.client.add_permission(folder_id, peer, write=False) + 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. @@ -405,32 +417,36 @@ def read_own_key_bundle(self, peer: str) -> str | None: folder_id = self._find_encryption_bundles_folder_id() if folder_id is None: return None - filename = self._encryption_bundle_filename(self._layout.email, peer) + 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._layout.client.drive_service.files().list(q=query, fields="files(id)") + self._cached_client.client.drive_service.files().list( + q=query, fields="files(id)" + ) ) items = results.get("files", []) if not items: return None - return self._layout.client.download_file(items[0]["id"]).decode("utf-8") + 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._layout.email}.json + Searches for: encryption_bundle_{peer}_for_{self._cached_client.email}.json in peer's bundles folder. """ - filename = self._encryption_bundle_filename(peer, self._layout.email) + 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._layout.client.drive_service.files().list(q=query, fields="files(id)") + 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._layout.client.download_file(items[0]["id"]) + 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 diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 7bb73df70dc..b1c13a5d227 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -14,7 +14,7 @@ SyftboxPlatformConnection, ) from syft.sync.connections.drive.gdrive_client import GDriveClient -from syft.sync.connections.drive.gdrive_layout import DriveLayout +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.environments.environment import Environment from syft.sync.utils.syftbox_utils import check_env @@ -60,14 +60,14 @@ def build_drive_service( class GDriveConnection(SyftboxPlatformConnection): """A Google Drive account as a syft connection. - It logs in and builds the ``DriveLayout`` that the transport components + It logs in and builds the ``DriveCachedClient`` that the transport components share: ``SyftboxManager._drive_transport`` builds them over it. """ class Config: arbitrary_types_allowed = True - layout: DriveLayout | None = None + cached_client: DriveCachedClient | None = None verbose: bool = True email: str token_path: Path | None = None @@ -130,8 +130,9 @@ def setup( new_service: Callable[[], Any], warm_syftbox_folder_id_cache: bool = True, ): - """Build the layout over a Drive service from ``new_service``, which - also builds the service of each copy for another thread. + """Build the ``DriveCachedClient`` over a Drive service from + ``new_service``, which also builds the service of each copy for another + thread. Args: new_service: Builds a Drive service, real or mock. @@ -139,9 +140,11 @@ def setup( SyftBox folder id (2-3 API calls). Pass False when the caller seeds the caches. """ - self.layout = DriveLayout(GDriveClient(new_service(), new_service), self.email) + self.cached_client = DriveCachedClient( + GDriveClient(new_service(), new_service), self.email + ) if warm_syftbox_folder_id_cache: - self.layout.get_personal_syftbox_folder_id() + self.cached_client.get_personal_syftbox_folder_id() self._is_setup = True @property diff --git a/syft/sync/login_utils.py b/syft/sync/login_utils.py index a9ddbf4c744..5025fefc9a1 100644 --- a/syft/sync/login_utils.py +++ b/syft/sync/login_utils.py @@ -23,7 +23,7 @@ def _read_remote_version( from syft.sync.connections.drive.gdrive_transport import GDriveConnection conn = GDriveConnection.from_token_path(email=email, token_path=token_path) - peer_info = DrivePeering(conn.layout).read_own_peer_info() + peer_info = DrivePeering(conn.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 c9e6a3e2978..a4367fbeb6b 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -16,7 +16,7 @@ from syft.sync.connections.drive.drive_collections import DriveCollections from syft.sync.connections.drive.drive_connections import ( DriveWatcherStore, - drive_layout_of, + drive_cached_client_of, ) from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore from syft.sync.connections.drive.drive_p2p import DriveP2PConnection @@ -610,11 +610,11 @@ def _transport_from_config(config: SyftboxManagerConfig) -> Transport: def _drive_transport(drive: GDriveConnection) -> Transport: """Every transport component over one Google Drive connection.""" return Transport( - p2p=DriveP2PConnection(drive.layout), - owner_store=DriveOwnerStore(drive.layout), + p2p=DriveP2PConnection(drive.cached_client), + owner_store=DriveOwnerStore(drive.cached_client), watcher_store=DriveWatcherStore(), - collections=DriveCollections(drive.layout), - peering=DrivePeering(drive.layout), + collections=DriveCollections(drive.cached_client), + peering=DrivePeering(drive.cached_client), platform=GdriveFilesPlatform(), ) @@ -1325,7 +1325,7 @@ def _copy(self): from copy import deepcopy new_config = deepcopy(self.config) - drive_service = drive_layout_of( + drive_service = drive_cached_client_of( self._connection_router.transport ).client.drive_service if isinstance(drive_service, mock_drive_service.MockDriveService): diff --git a/syft/version.py b/syft/version.py index 8f1f47a0782..b8a788bd066 100644 --- a/syft/version.py +++ b/syft/version.py @@ -7,7 +7,7 @@ # 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_layout.py and +# 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" diff --git a/tests/integration/with_unit_coverage/test_gdrive_pagination.py b/tests/integration/with_unit_coverage/test_gdrive_pagination.py index 76430d3a21a..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_layout import DriveLayout + 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 = DriveLayout.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( - DriveLayout, + DriveCachedClient, "get_file_metadatas_from_folder", get_files_page_size_2, ): diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index 2b7a5a1493c..c888eeec713 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -753,7 +753,7 @@ def test_owner_state_is_written_by_default(): def test_latest_checkpoint_read_retries_drive_error(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.layout) + store = DriveOwnerStore(do_conn.cached_client) store.put_checkpoint("checkpoint_1.0.tar.gz", b"full") fail_next_listing(monkeypatch, CHECKPOINT_FILENAME_PREFIX) @@ -762,7 +762,7 @@ def test_latest_checkpoint_read_retries_drive_error(monkeypatch): def test_incremental_count_retries_drive_error(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.layout) + store = DriveOwnerStore(do_conn.cached_client) store.put_incremental("incremental_checkpoint_1_1.0.tar.gz", b"inc") fail_next_listing(monkeypatch, INCREMENTAL_CHECKPOINT_PREFIX) @@ -771,7 +771,7 @@ def test_incremental_count_retries_drive_error(monkeypatch): def test_latest_checkpoint_download_error_raises(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.layout) + store = DriveOwnerStore(do_conn.cached_client) store.put_checkpoint("checkpoint_1.0.tar.gz", b"full") def forbidden(self, file_id: str) -> bytes: diff --git a/tests/unit/test_dataset_collection_listing.py b/tests/unit/test_dataset_collection_listing.py index 026b33cc61d..5aa34119d11 100644 --- a/tests/unit/test_dataset_collection_listing.py +++ b/tests/unit/test_dataset_collection_listing.py @@ -12,17 +12,19 @@ from syft.sync.connections.drive.drive_collections import DriveCollections from syft.sync.connections.drive.gdrive_client import GDriveClient -from syft.sync.connections.drive.gdrive_layout import DriveLayout +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient VALID = f"{DATASET_COLLECTION_PREFIX}_mytag_abc123" UNPARSEABLE = DATASET_COLLECTION_PREFIX def _collections(files): - layout = DriveLayout(GDriveClient(Mock(), Mock), "alice@example.com") - layout._syftbox_folder_id = "syftbox-id" - layout.client.drive_service.files().list().execute.return_value = {"files": files} - return DriveCollections(layout) + 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(): diff --git a/tests/unit/test_delete_syftbox.py b/tests/unit/test_delete_syftbox.py index dcae0919d1d..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 drive_layout_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of EMAIL = "test@example.com" @@ -128,11 +128,11 @@ def test_a_terminal_still_asks(self, mock_stdin): assert _prompt_mismatch(_old_version_info(), _old_version_info()) == "2" -def _query_files(layout, 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 = ( - layout.client.drive_service.files() + cached_client.client.drive_service.files() .list(q=q, fields="files(id, name)") .execute() ) @@ -172,21 +172,23 @@ def test_proposal_moved_to_root_by_do_delete_found(self): def test_proposal_moved_to_root_when_do_deletes_syftbox_folder_only(self): ds_manager, do_manager, archived = _archive_one_proposal() - do_layout = _layout(do_manager) - archive_id = do_layout.get_archive_folder_id(ds_manager.email) - do_layout.client.drive_service.files().delete( - fileId=do_layout.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 = _owner_store(ds_manager).find_orphaned_syft_files() - assert archive_id not in do_layout.client.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 = _layout(ds_manager).client.drive_service + drive = _cached_client(ds_manager).client.drive_service user_file = drive.files().create(body={"name": "checkpoint_final.pt"}).execute() found = _owner_store(ds_manager).find_orphaned_syft_files() @@ -195,15 +197,15 @@ def test_untagged_file_in_root_not_found(self): def test_syftbox_folder_not_found(self): ds_manager, _, _ = _archive_one_proposal() - ds_layout = _layout(ds_manager) + ds_cached_client = _cached_client(ds_manager) found = _owner_store(ds_manager).find_orphaned_syft_files() - assert ds_layout.get_syftbox_folder_id() not in found + assert ds_cached_client.get_syftbox_folder_id() not in found -def _layout(manager: SyftboxManager): - return drive_layout_of(manager._connection_router.transport) +def _cached_client(manager: SyftboxManager): + return drive_cached_client_of(manager._connection_router.transport) def _owner_store(manager: SyftboxManager): @@ -218,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_layout = _layout(do_manager) - archive_id = do_layout.get_archive_folder_id(ds_manager.email) - archived = [m["id"] for m in do_layout.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 481e6784be9..ef38cdffacc 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -72,10 +72,10 @@ def _events_name(ts: float) -> str: @pytest.fixture def pair(): - """DO and DS layouts over one mock Drive, with the DS's channels open.""" + """DO and DS cached clients over one mock Drive, with the DS's channels open.""" do_conn, ds_conn = pair_with_mock_service(DO, DS) - DrivePeering(ds_conn.layout).open_channels(DO) - return do_conn.layout, ds_conn.layout + DrivePeering(ds_conn.cached_client).open_channels(DO) + return do_conn.cached_client, ds_conn.cached_client # -- the transport ---------------------------------------------------------------- @@ -97,12 +97,14 @@ def test_drive_transport_does_every_job(): def test_a_proposed_change_sent_through_p2p_reaches_the_owner_inbox(pair): - do_layout, ds_layout = pair - status = DriveP2PConnection(ds_layout).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_layout).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, @@ -111,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_layout, ds_layout = 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_layout).send( + DriveP2PConnection(ds_cached_client).send( DO, name, f"job {submitted}".encode(), MessageKind.PROPOSED_CHANGES ) - p2p = DriveP2PConnection(do_layout) + 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] == [ @@ -161,20 +163,20 @@ def counted_get_media(self, *args, **kwargs): return calls -def _send_proposals(ds_layout, 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_layout).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_layout, ds_layout = pair - p2p = DriveP2PConnection(do_layout) + 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_layout, 3) + names = _send_proposals(ds_cached_client, 3) calls = _count_drive_calls(monkeypatch) received = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) @@ -184,8 +186,8 @@ def test_receiving_inbox_lists_once_and_downloads_each_message(pair, monkeypatch def test_inbox_messages_download_at_the_same_time(pair, monkeypatch): - do_layout, ds_layout = pair - _send_proposals(ds_layout, 3) + do_cached_client, ds_cached_client = pair + _send_proposals(ds_cached_client, 3) running, most_running = [0], [0] lock = threading.Lock() original_download = GDriveClient.download_file @@ -200,14 +202,14 @@ def slow_download(self, file_id: str) -> bytes: return original_download(self, file_id) monkeypatch.setattr(GDriveClient, "download_file", slow_download) - DriveP2PConnection(do_layout).receive(DS, MessageKind.PROPOSED_CHANGES) + 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_layout, ds_layout = pair - _send_proposals(ds_layout, 2) - p2p = DriveP2PConnection(do_layout) + 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 @@ -219,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_layout, ds_layout = pair - _send_proposals(ds_layout, 4) - p2p = DriveP2PConnection(do_layout) + 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 @@ -233,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_layout, ds_layout = pair - _send_proposals(ds_layout, 2) - p2p = DriveP2PConnection(do_layout) + 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"}) @@ -250,11 +252,11 @@ 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_layout, ds_layout = pair - names = _send_proposals(ds_layout, 3) + do_cached_client, ds_cached_client = pair + names = _send_proposals(ds_cached_client, 3) [broken_id] = [ m["file_id"] - for m in DriveP2PConnection(do_layout)._list_inbox(DS) + for m in DriveP2PConnection(do_cached_client)._list_inbox(DS) if m["file_name"] == names[1] ] original_download = GDriveClient.download_file @@ -266,7 +268,7 @@ def download_all_but_broken(self, file_id: str) -> bytes: monkeypatch.setattr(GDriveClient, "download_file", download_all_but_broken) with pytest.warns(UserWarning, match=f"Skipped inbox message {names[1]}"): - received = DriveP2PConnection(do_layout).receive( + received = DriveP2PConnection(do_cached_client).receive( DS, MessageKind.PROPOSED_CHANGES ) @@ -274,19 +276,19 @@ def download_all_but_broken(self, file_id: str) -> bytes: def test_receiving_inbox_while_drive_is_down_raises(pair, monkeypatch): - do_layout, ds_layout = pair - _send_proposals(ds_layout, 2) + 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_layout).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_layout, ds_layout = pair - DriveP2PConnection(ds_layout).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_layout) + p2p = DriveP2PConnection(do_cached_client) [message] = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) p2p.mark_processed(message) @@ -294,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_layout, ds_layout = pair + do_cached_client, ds_cached_client = pair name = _events_name(1.0) - status = DriveP2PConnection(do_layout).send( + status = DriveP2PConnection(do_cached_client).send( DS, name, b"results", MessageKind.FILE_EVENTS ) assert status is SendStatus.DELIVERED - p2p = DriveP2PConnection(ds_layout) + 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" @@ -308,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_layout, ds_layout = pair - DriveP2PConnection(do_layout).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_layout) + p2p = DriveP2PConnection(ds_cached_client) [listed] = p2p.outbox_message_metadatas(DO, since_timestamp=None) copy = p2p.clone_for_thread() - assert copy._layout is not ds_layout + 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_layout, _ = pair + do_cached_client, _ = pair with pytest.raises(ValueError): - DriveP2PConnection(do_layout).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_layout, _ = pair - assert DriveP2PConnection(do_layout).delivered() == [] + do_cached_client, _ = pair + assert DriveP2PConnection(do_cached_client).delivered() == [] def test_receiving_file_events_through_p2p_is_not_built_yet(pair): - _, ds_layout = pair - p2p = DriveP2PConnection(ds_layout) + _, 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) @@ -345,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_layout, _ = pair - store = DriveOwnerStore(do_layout) + 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") @@ -361,8 +365,8 @@ def test_owner_event_log_is_read_oldest_first(pair): def test_owner_event_count_matches_the_listed_handles(pair): - do_layout, _ = pair - store = DriveOwnerStore(do_layout) + 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") @@ -370,8 +374,8 @@ def test_owner_event_count_matches_the_listed_handles(pair): def test_owner_event_log_read_while_drive_is_down_raises(pair, monkeypatch): - do_layout, _ = pair - store = DriveOwnerStore(do_layout) + 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) @@ -380,8 +384,8 @@ def test_owner_event_log_read_while_drive_is_down_raises(pair, monkeypatch): def test_owner_event_log_download_error_raises(pair, monkeypatch): - do_layout, _ = pair - store = DriveOwnerStore(do_layout) + 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: @@ -393,16 +397,16 @@ def forbidden(self, file_id: str) -> bytes: def test_owner_store_keeps_the_newest_full_checkpoint(pair): - do_layout, _ = pair - store = DriveOwnerStore(do_layout) + 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_layout, _ = pair - store = DriveOwnerStore(do_layout) + 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 @@ -412,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_layout, _ = pair - store = DriveOwnerStore(do_layout) + 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" @@ -422,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_layout, _ = pair - store = DriveOwnerStore(do_layout) + 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_layout.client.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"}] @@ -442,7 +446,7 @@ def test_owner_store_files_carry_the_syft_tag(pair): def test_watcher_storage_on_drive_is_not_built_yet(pair): - _, ds_layout = pair + _, ds_cached_client = pair with pytest.raises(NotImplementedError): DriveWatcherStore().latest_checkpoint(DO) @@ -451,8 +455,11 @@ def test_watcher_storage_on_drive_is_not_built_yet(pair): def test_a_published_collection_is_listed_and_downloaded_by_its_reader(pair): - do_layout, ds_layout = pair - owner, reader = DriveCollections(do_layout), DriveCollections(ds_layout) + 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) @@ -475,8 +482,8 @@ def test_a_published_collection_is_listed_and_downloaded_by_its_reader(pair): def test_a_failed_collection_delete_raises(pair, monkeypatch): - do_layout, _ = pair - owner = DriveCollections(do_layout) + do_cached_client, _ = pair + owner = DriveCollections(do_cached_client) owner.create(DATASETS, "census", "hash", DO) def forbidden(self) -> None: @@ -494,22 +501,22 @@ def forbidden(self) -> None: def test_peering_finds_the_peer_that_opened_channels(pair): - do_layout, _ = pair - peering = DrivePeering(do_layout) + 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_layout, _ = pair - peering = DrivePeering(do_layout) + 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_layout, _ = pair - peering = DrivePeering(do_layout) + do_cached_client, _ = pair + peering = DrivePeering(do_cached_client) peering.update_peer_registry(DS, "accepted") peering._peers_json_cache = None # a new process has no cache @@ -528,8 +535,8 @@ def test_reading_own_key_bundle_creates_no_folder(pair): def test_peering_publishes_a_key_bundle_both_sides_can_read(pair): - do_layout, ds_layout = pair - owner, scientist = DrivePeering(do_layout), DrivePeering(ds_layout) + 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) @@ -541,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_layout, ds_layout = pair - owner, scientist = DrivePeering(do_layout), DrivePeering(ds_layout) + 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, diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index 155fe2726ef..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 drive_layout_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_layout = drive_layout_of(do_manager._connection_router.transport) - inbox_id = do_layout.get_own_datasite_inbox_id(ds_manager.email) - files = do_layout.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_layout.client.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_layout.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 5031a5ad38b..588d7833823 100644 --- a/tests/unit/test_encryption_policy.py +++ b/tests/unit/test_encryption_policy.py @@ -173,11 +173,13 @@ def test_missing_version_file_created_with_own_policy(): a_conn, b_conn = mock_drive_service.pair_with_mock_service( "a@example.com", "b@example.com" ) - DrivePeering(a_conn.layout).grant_peer_info_read( + DrivePeering(a_conn.cached_client).grant_peer_info_read( "b@example.com", _peer_info(REQUIRED) ) assert ( - DrivePeering(b_conn.layout).read_peer_info("a@example.com").encryption_policy + DrivePeering(b_conn.cached_client) + .read_peer_info("a@example.com") + .encryption_policy == REQUIRED ) diff --git a/tests/unit/test_gdrive_client.py b/tests/unit/test_gdrive_client.py index f7a6e39d600..957261e2ef8 100644 --- a/tests/unit/test_gdrive_client.py +++ b/tests/unit/test_gdrive_client.py @@ -13,8 +13,8 @@ def test_batch_delete_raises_error_whose_text_contains_404(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - folder_id = do_conn.layout.client.create_folder( - "folder", do_conn.layout.get_syftbox_folder_id() + folder_id = do_conn.cached_client.client.create_folder( + "folder", do_conn.cached_client.get_syftbox_folder_id() ) def forbidden(self) -> None: @@ -22,13 +22,13 @@ def forbidden(self) -> None: monkeypatch.setattr(MockDeleteRequest, "execute", forbidden) with pytest.raises(HttpError): - do_conn.layout.client.delete_multiple_files_by_ids([folder_id]) + do_conn.cached_client.client.delete_multiple_files_by_ids([folder_id]) def test_batch_share_raises_error_whose_text_says_already_shared(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - folder_id = do_conn.layout.client.create_folder( - "folder", do_conn.layout.get_syftbox_folder_id() + folder_id = do_conn.cached_client.client.create_folder( + "folder", do_conn.cached_client.get_syftbox_folder_id() ) def forbidden(self) -> None: @@ -36,4 +36,6 @@ def forbidden(self) -> None: monkeypatch.setattr(MockPermissionCreateRequest, "execute", forbidden) with pytest.raises(HttpError): - do_conn.layout.client.batch_add_permissions(folder_id, ["ds@example.com"]) + do_conn.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 51ede913ff4..726903c4611 100644 --- a/tests/unit/test_gdrive_pagination.py +++ b/tests/unit/test_gdrive_pagination.py @@ -5,7 +5,7 @@ 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.gdrive_layout import DriveLayout +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.connections.drive.mock_drive_service import ( MockDriveBackingStore, @@ -17,19 +17,19 @@ ) -def create_mock_layout() -> DriveLayout: - return DriveLayout(GDriveClient(Mock(), Mock), "test@test.com") +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)""" - layout = create_mock_layout() + cached_client = create_mock_cached_client() - layout.client.drive_service.files().list().execute.return_value = { + cached_client.client.drive_service.files().list().execute.return_value = { "files": [{"name": "file1.txt", "id": "id1"}] } - files = layout.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" @@ -37,38 +37,40 @@ def test_single_page_works(): def test_multiple_pages_pagination(): """When there are >100 files, should fetch all pages""" - layout = create_mock_layout() + cached_client = create_mock_cached_client() - layout.client.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 = layout.get_file_metadatas_from_folder("folder_id") + files = cached_client.get_file_metadatas_from_folder("folder_id") assert len(files) == 3 - assert layout.client.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""" - layout = create_mock_layout() + cached_client = create_mock_cached_client() - layout.client.drive_service.files().list().execute.return_value = {"files": []} + cached_client.client.drive_service.files().list().execute.return_value = { + "files": [] + } - files = layout.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_layout.DriveLayout.get_peer_datasite_outbox_id" + "syft.sync.connections.drive.drive_cached_client.DriveCachedClient.get_peer_datasite_outbox_id" ) @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""" - layout = create_mock_layout() + cached_client = create_mock_cached_client() # Create valid filenames with proper UUIDs uuid1 = uuid4() @@ -76,7 +78,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) - layout.client.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"}, @@ -103,7 +105,9 @@ def mock_download(file_id): events = [ FileChangeEventsMessage.from_compressed_data(data) - for data in DriveP2PConnection(layout).outbox_messages("peer@test.com", None) + for data in DriveP2PConnection(cached_client).outbox_messages( + "peer@test.com", None + ) ] # Should return in chronological order (1.0, 2.0, 3.0) @@ -117,20 +121,22 @@ def mock_download(file_id): def test_orderby_name_desc_applied(): """Drive API should be called with orderBy='name desc'""" - layout = create_mock_layout() + cached_client = create_mock_cached_client() - layout.client.drive_service.files().list().execute.return_value = {"files": []} + cached_client.client.drive_service.files().list().execute.return_value = { + "files": [] + } - layout.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 = layout.client.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_early_termination_stops_pagination(): """Should stop fetching pages when encountering old file""" - layout = create_mock_layout() + cached_client = create_mock_cached_client() # Create filenames using FileChangeEventsMessageFileName fname1 = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506450.0) @@ -142,7 +148,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 - layout.client.drive_service.files().list().execute.side_effect = [ + cached_client.client.drive_service.files().list().execute.side_effect = [ { "files": [ {"name": fname1.as_string(), "id": "id1"}, @@ -166,12 +172,12 @@ def test_early_termination_stops_pagination(): }, ] - files = layout.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 layout.client.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 @@ -179,12 +185,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""" - layout = create_mock_layout() + cached_client = create_mock_cached_client() fname1 = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506450.0) fname2 = FileChangeEventsMessageFileName(id=uuid4(), timestamp=1731506440.0) - layout.client.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", @@ -195,20 +201,20 @@ def test_no_early_termination_when_all_files_new(): }, ] - files = layout.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 layout.client.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""" - layout = create_mock_layout() + cached_client = create_mock_cached_client() - layout.client.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", @@ -219,24 +225,26 @@ def test_no_early_termination_when_since_timestamp_none(): }, ] - files = layout.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 layout.client.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(): - layout = create_mock_layout() + 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) - layout.client.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 = layout.get_file_metadatas_from_folder( + files = cached_client.get_file_metadatas_from_folder( "folder_id", since_timestamp=1731506400.92 ) @@ -247,17 +255,22 @@ def test_newer_message_listed_after_older_one_in_same_second_not_skipped(): FILE_COUNT = 250 -def _mock_layouts(*emails: str) -> list[DriveLayout]: +def _mock_cached_clients(*emails: str) -> list[DriveCachedClient]: store = MockDriveBackingStore() return [ - GDriveConnection.from_service(e, MockDriveService(store, e)).layout + GDriveConnection.from_service(e, MockDriveService(store, e)).cached_client for e in emails ] def test_collection_listings_read_every_page(): - do_layout, ds_layout = _mock_layouts("do@example.com", "ds@example.com") - owner, reader = DriveCollections(do_layout), DriveCollections(ds_layout) + 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") @@ -272,9 +285,9 @@ def test_collection_listings_read_every_page(): def test_peer_requests_read_every_page(): peer_emails = [f"peer{i}@example.com" for i in range(FILE_COUNT)] - do_layout, *peers = _mock_layouts("do@example.com", *peer_emails) + 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_layout).incoming_requests() + 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 2273a24a9d3..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 drive_layout_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 drive_layout_of(manager._connection_router.transport).client.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 9684c183566..8a3b97d21b2 100644 --- a/tests/unit/test_mock_drive_service.py +++ b/tests/unit/test_mock_drive_service.py @@ -448,8 +448,10 @@ def test_create_folder(self): mock_service = MockDriveService(store, "user@example.com") connection = GDriveConnection.from_service("user@example.com", mock_service) - syftbox_id = connection.layout.create_syftbox_folder() - folder_id = connection.layout.client.create_folder("TestFolder", syftbox_id) + syftbox_id = connection.cached_client.create_syftbox_folder() + folder_id = connection.cached_client.client.create_folder( + "TestFolder", syftbox_id + ) assert store.files[syftbox_id].properties == {"syft": "true"} assert store.files[folder_id].properties == {"syft": "true"} @@ -468,7 +470,7 @@ def test_download_file(self): mock_service = MockDriveService(store, "user@example.com") connection = GDriveConnection.from_service("user@example.com", mock_service) - content = connection.layout.client.download_file(file.id) + content = connection.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 43b22f0831c..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 drive_layout_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 ) - do_layout = drive_layout_of(do_manager._connection_router.transport) - outbox_id = do_layout.get_own_datasite_outbox_id(ds_manager.email) - pre = do_layout.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 = do_layout.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 4c53b05cf81..afa8c7d2c2b 100644 --- a/tests/unit/test_p2p_folder_lookup.py +++ b/tests/unit/test_p2p_folder_lookup.py @@ -11,7 +11,7 @@ from unittest.mock import Mock -from syft.sync.connections.drive.gdrive_layout import DriveLayout +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.gdrive_client import GDriveClient ME = "alice@example.com" @@ -22,14 +22,14 @@ def _name(version: str, datasite: str, folder_type: str, peer: str) -> str: return f"syft_datasite#{version}#{datasite}#{folder_type}#{peer}" -def _layout(found): - layout = DriveLayout(GDriveClient(Mock(), Mock), ME) - layout.client.find_folders = Mock(return_value=found) - return layout +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(layout): - return layout.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(_layout([("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(_layout([("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(_layout(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(_layout(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(_layout(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(_layout([("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(_layout([("outbox", outbox)])) is None + assert _lookup(_cached_client([("outbox", outbox)])) is None def test_no_folder_returns_none(): - assert _lookup(_layout([])) is None + assert _lookup(_cached_client([])) is None def test_a_name_that_does_not_parse_is_ignored(): - assert _lookup(_layout([("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 456a344d0ce..c4b5c60ceb0 100644 --- a/tests/unit/test_peer_request_owner.py +++ b/tests/unit/test_peer_request_owner.py @@ -6,7 +6,7 @@ import pytest -from syft.sync.connections.drive.gdrive_layout import ( +from syft.sync.connections.drive.drive_cached_client import ( GDRIVE_P2P_FOLDER_DATASITE_PREFIX, GdriveP2PFolder, ) @@ -19,7 +19,7 @@ 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 drive_layout_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of DO_EMAIL = "do@test.com" DS_EMAIL = "ds@test.com" @@ -47,7 +47,7 @@ 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 = drive_layout_of( + backing_store = drive_cached_client_of( victim._connection_router.transport ).client.drive_service._backing_store owners = [{"emailAddress": owner_email}] if owner_email else [] diff --git a/tests/unit/test_peers_json_version.py b/tests/unit/test_peers_json_version.py index 3ca54276798..b887422f206 100644 --- a/tests/unit/test_peers_json_version.py +++ b/tests/unit/test_peers_json_version.py @@ -15,7 +15,7 @@ from unittest.mock import Mock, patch from syft.sync.connections.drive.gdrive_client import GDriveClient -from syft.sync.connections.drive.gdrive_layout import DriveLayout +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.platforms.gdrive_files_platform import GdriveFilesPlatform @@ -25,8 +25,8 @@ def _peering(peers_data): - layout = DriveLayout(GDriveClient(Mock(), Mock), "alice@example.com") - peering = DrivePeering(layout) + cached_client = DriveCachedClient(GDriveClient(Mock(), Mock), "alice@example.com") + peering = DrivePeering(cached_client) peering._peers_json_cache = dict(peers_data) return peering @@ -44,7 +44,9 @@ def test_a_write_stamps_the_reserved_entry(): peering = _peering({PEER: {"state": "accepted"}}) with ( patch.object(DrivePeering, "_get_peers_file_id", return_value="file-id"), - patch.object(DriveLayout, "get_syftbox_folder_id", return_value="folder-id"), + patch.object( + DriveCachedClient, "get_syftbox_folder_id", return_value="folder-id" + ), patch.object(GDriveClient, "create_file_payload", return_value=(Mock(), None)), ): peering._write_peers_json({PEER: {"state": "accepted"}}) diff --git a/tests/unit/test_rolling_state.py b/tests/unit/test_rolling_state.py index 24dd60b3aad..629656e1e65 100644 --- a/tests/unit/test_rolling_state.py +++ b/tests/unit/test_rolling_state.py @@ -289,7 +289,7 @@ def test_no_rolling_state_when_owner_state_disabled(): def test_rolling_state_read_retries_drive_error(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.layout) + store = DriveOwnerStore(do_conn.cached_client) store.put_rolling_state("rolling_state_1.0.tar.gz", b"state") fail_next_listing(monkeypatch, ROLLING_STATE_FILENAME_PREFIX) @@ -298,7 +298,7 @@ def test_rolling_state_read_retries_drive_error(monkeypatch): def test_rolling_state_read_while_drive_is_down_raises(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.layout) + store = DriveOwnerStore(do_conn.cached_client) store.put_rolling_state("rolling_state_1.0.tar.gz", b"state") def drive_down(self, file_id: str) -> bytes: @@ -311,7 +311,7 @@ def drive_down(self, file_id: str) -> bytes: def test_rolling_state_listing_while_drive_is_down_raises(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.layout) + store = DriveOwnerStore(do_conn.cached_client) store.put_rolling_state("rolling_state_1.0.tar.gz", b"state") def drive_down(self) -> dict: @@ -324,7 +324,7 @@ def drive_down(self) -> dict: def test_rolling_state_download_error_raises(monkeypatch): do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.layout) + store = DriveOwnerStore(do_conn.cached_client) store.put_rolling_state("rolling_state_1.0.tar.gz", b"state") def forbidden(self, file_id: str) -> bytes: diff --git a/tests/unit/test_sender_binding.py b/tests/unit/test_sender_binding.py index 67eef15ee56..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 drive_layout_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]: - do_layout = do_manager.datasite_owner_syncer.event_cache.file_connection - return [str(path) for path, _ in do_layout.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) - do_layout = drive_layout_of( + do_cached_client = drive_cached_client_of( do_manager.datasite_owner_syncer.connection_router.transport ) folder_names = [ - f.name for f in do_layout.client.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 f39705e45a4..93a25371b50 100644 --- a/tests/unit/test_shared_router.py +++ b/tests/unit/test_shared_router.py @@ -8,9 +8,9 @@ from syft.sync.connections.base_connection import ConnectionConfig from syft.sync.connections.drive.drive_connections import ( DriveWatcherStore, - drive_layout_of, + drive_cached_client_of, ) -from syft.sync.connections.drive.gdrive_layout import DriveLayout +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.connections.drive.grdrive_config import GdriveConnectionConfig from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service @@ -50,11 +50,11 @@ def test_every_part_of_manager_shares_one_router(): assert manager._peer_store is router.peer_store -def _drive_layouts(transport) -> list[DriveLayout]: - """The layout of each component that uses Drive. The watcher store does not +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._layout + component._cached_client for component in transport.components() if not isinstance(component, DriveWatcherStore) ] @@ -75,25 +75,25 @@ def build_mock_connection(connection_config): manager = SyftboxManager.from_config(config) [connection] = built - for layout in _drive_layouts(manager._connection_router.transport): - assert layout is connection.layout + for cached_client in _drive_cached_clients(manager._connection_router.transport): + assert cached_client is connection.cached_client -def test_every_drive_component_shares_one_layout(): +def test_every_drive_component_shares_one_cached_client(): _, do_manager = SyftboxManager.pair_with_mock_drive_service_connection() transport = do_manager._connection_router.transport - layout = drive_layout_of(transport) + cached_client = drive_cached_client_of(transport) - for component_layout in _drive_layouts(transport): - assert component_layout is layout + for component_cached_client in _drive_cached_clients(transport): + assert component_cached_client is cached_client -def test_drive_layout_of_transport_not_on_drive_raises(): +def test_drive_cached_client_of_transport_not_on_drive_raises(): transport = SyftboxManager._drive_transport(_mock_connection("do@example.com")) transport = transport.model_copy(update={"owner_store": Mock(spec=OwnerStore)}) with pytest.raises(TypeError, match="not built on Google Drive"): - drive_layout_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 04d220f6080..f9279ca8f99 100644 --- a/tests/unit/test_sync_manager.py +++ b/tests/unit/test_sync_manager.py @@ -25,7 +25,7 @@ get_mock_proposed_events_messages, grant_job_inbox_access, ) -from syft.sync.connections.drive.drive_connections import drive_layout_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: @@ -488,9 +488,11 @@ 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( email, - drive_layout_of(ds_manager._connection_router.transport).client.drive_service, + drive_cached_client_of( + ds_manager._connection_router.transport + ).client.drive_service, ) - DrivePeering(conn.layout).open_channels(do_manager.email) + DrivePeering(conn.cached_client).open_channels(do_manager.email) _add_known_peer(do_manager, email) return ConnectionRouter( transport=SyftboxManager._drive_transport(conn), diff --git a/tests/unit/test_version_negotiation.py b/tests/unit/test_version_negotiation.py index 0d456d5edd7..d1d0d19f0b9 100644 --- a/tests/unit/test_version_negotiation.py +++ b/tests/unit/test_version_negotiation.py @@ -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 drive_layout_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of def build_client_version(client_version: str) -> VersionInfo: @@ -240,7 +240,7 @@ def test_load_peer_versions_parallel(self): SyftboxManager.pair_with_mock_drive_service_connection() ) - backing_store = drive_layout_of( + backing_store = drive_cached_client_of( ds_manager._connection_router.transport ).client.drive_service._backing_store diff --git a/tests/unit/test_versioned_folder_adopt.py b/tests/unit/test_versioned_folder_adopt.py index a177ad75c85..fbb2cad084e 100644 --- a/tests/unit/test_versioned_folder_adopt.py +++ b/tests/unit/test_versioned_folder_adopt.py @@ -20,23 +20,26 @@ from syft.sync.checkpoints.checkpoint import CHECKPOINT_VERSION, Checkpoint from syft.sync.connections.drive.gdrive_client import GDriveClient -from syft.sync.connections.drive.gdrive_layout import DriveLayout, _partition_by_version +from syft.sync.connections.drive.drive_cached_client import ( + DriveCachedClient, + _partition_by_version, +) 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 drive_layout_of +from syft.sync.connections.drive.drive_connections import drive_cached_client_of EMAIL = "alice@example.com" -def _layout(): - return DriveLayout(GDriveClient(Mock(), Mock), EMAIL) +def _cached_client(): + return DriveCachedClient(GDriveClient(Mock(), Mock), EMAIL) -def _renames(layout): - """Return the (fileId, new name) pairs the layout 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 layout.client.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", {}) ] @@ -77,65 +80,65 @@ def test_partition_returns_empty_for_a_bad_current_version(): def test_a_compatible_folder_wins_and_nothing_is_renamed(): - layout = _layout() + cached_client = _cached_client() folders = [("same", f"0.2.5#{EMAIL}"), ("old", f"0.1.9#{EMAIL}")] - got = layout.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(layout) == [] + assert _renames(cached_client) == [] def test_an_older_folder_is_adopted_by_rename(): - layout = _layout() + cached_client = _cached_client() folders = [("old", f"0.1.9#{EMAIL}")] - got = layout.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(layout) == [("old", f"0.2.7#{EMAIL}")] + assert _renames(cached_client) == [("old", f"0.2.7#{EMAIL}")] def test_the_highest_older_folder_is_adopted(): - layout = _layout() + cached_client = _cached_client() folders = [ ("v1", f"0.1.9#{EMAIL}"), ("v2", f"0.1.20#{EMAIL}"), ("v0", f"0.0.4#{EMAIL}"), ] - got = layout.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(layout) == [("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. - layout = _layout() + cached_client = _cached_client() folders = [("new", f"0.3.0#{EMAIL}")] with pytest.raises(RuntimeError, match="0.3.0"): - layout.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(layout) == [] + assert _renames(cached_client) == [] def test_no_folder_returns_none_so_the_caller_creates_one(): - layout = _layout() - got = layout.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(layout) == [] + assert _renames(cached_client) == [] def test_two_compatible_folders_still_raise(): - layout = _layout() + cached_client = _cached_client() folders = [("a", f"0.2.1#{EMAIL}"), ("b", f"0.2.2#{EMAIL}")] with pytest.raises(RuntimeError): - layout.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" ) @@ -147,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, layout, + pre-upgrade client would have left it. Returns (router, cached_client, older_name). """ _, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( @@ -155,7 +158,7 @@ def _checkpoints_owner_with_version(checkpoint_version: int): sync_automatically=False, ) router = do_manager._connection_router - layout = drive_layout_of(router.transport) + cached_client = drive_cached_client_of(router.transport) router.upload_checkpoint( Checkpoint(email=do_manager.email, version=checkpoint_version) ) @@ -163,19 +166,19 @@ def _checkpoints_owner_with_version(checkpoint_version: int): folder_id = router.transport.owner_store._get_checkpoints_folder_id() older_name = f"{do_manager.email}-0.0.1-checkpoints" execute_with_retries( - layout.client.drive_service.files().update( + cached_client.client.drive_service.files().update( fileId=folder_id, body={"name": older_name} ) ) - return router, layout, older_name + return router, cached_client, older_name -def _checkpoints_folder_names(layout) -> list[str]: +def _checkpoints_folder_names(cached_client) -> list[str]: return [ name - for _, name in layout.client.find_folders( - name_contains=[f"{layout.email}-", "-checkpoints"], - parent_id=layout.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(), ) ] @@ -183,16 +186,16 @@ def _checkpoints_folder_names(layout) -> 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, layout, 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(layout) == [ + 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(layout) + assert older_name not in _checkpoints_folder_names(cached_client) def test_an_adopted_folder_still_refuses_a_later_checkpoint(capsys): @@ -200,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, layout, _ = _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 ff9c7e5eb19..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_layout import ( +from syft.sync.connections.drive.drive_cached_client import ( _extract_version_from_name, _looks_like_version, _partition_by_version, From 2afdaf9e27288af1c854a19216e0610fa64b1291 Mon Sep 17 00:00:00 2001 From: Koen van der Veen Date: Fri, 30 Oct 2026 13:11:55 +0100 Subject: [PATCH 28/28] refactor: merge GDriveConnection into DriveCachedClient GDriveConnection only logged in and built the DriveCachedClient, so the login constructors (from_config, from_token_path, from_service) now live on DriveCachedClient and gdrive_transport.py is gone. build_drive_service and SCOPES move to gdrive_client.py. SyftboxPlatformConnection and GdriveConnectionConfig.connection_type had no other users and are removed. --- .../syft-enclave/src/syft_enclaves/utils.py | 4 +- packages/syft-rds/tests/test_sync_manager.py | 10 +- .../tests/test_version_mismatch_flow.py | 12 +- scripts/check_token.py | 2 +- scripts/reproduce_stale_socket.py | 2 +- syft/gdrive_utils.py | 20 ++- syft/sync/connections/base_connection.py | 12 +- .../connections/drive/drive_cached_client.py | 82 +++++++++- syft/sync/connections/drive/gdrive_client.py | 42 ++++- .../connections/drive/gdrive_transport.py | 152 ------------------ syft/sync/connections/drive/grdrive_config.py | 3 - .../connections/drive/mock_drive_service.py | 18 +-- syft/sync/login_utils.py | 8 +- syft/sync/syftbox_manager.py | 30 ++-- tests/unit/test_checkpoints.py | 12 +- tests/unit/test_drive_connections.py | 16 +- tests/unit/test_encryption_policy.py | 8 +- tests/unit/test_gdrive_client.py | 18 +-- tests/unit/test_gdrive_pagination.py | 4 +- tests/unit/test_mock_drive_service.py | 26 ++- tests/unit/test_rolling_state.py | 16 +- tests/unit/test_shared_router.py | 21 ++- tests/unit/test_sync_manager.py | 8 +- tests/unit/utils.py | 4 +- 24 files changed, 237 insertions(+), 293 deletions(-) delete mode 100644 syft/sync/connections/drive/gdrive_transport.py 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/tests/test_sync_manager.py b/packages/syft-rds/tests/test_sync_manager.py index f7e496368e5..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 @@ -1271,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 f8be407a889..7b4728ac12d 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -12,7 +12,7 @@ from syft.sync.connections.drive.drive_cached_client import ( GDRIVE_P2P_FOLDER_DATASITE_PREFIX, ) -from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.mock_drive_service import ( MockDriveService, ) @@ -74,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( @@ -99,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 DrivePeering(mock_conn.cached_client).read_own_peer_info().version_info + return DrivePeering(cached_client).read_own_peer_info().version_info with ( patch( 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 fce56384bdd..0139a2c63c9 100644 --- a/syft/gdrive_utils.py +++ b/syft/gdrive_utils.py @@ -12,13 +12,15 @@ from googleapiclient.http import MediaIoBaseDownload from syft.sync.config.config import settings -from syft.sync.connections.drive.gdrive_client import GOOGLE_FOLDER_MIME_TYPE +from syft.sync.connections.drive.gdrive_client import ( + GOOGLE_FOLDER_MIME_TYPE, + build_drive_service, +) +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.connections.drive.gdrive_transport import SCOPES as GDRIVE_SCOPES -from syft.sync.connections.drive.gdrive_transport import build_drive_service from syft.sync.environments.environment import Environment from syft.sync.utils.syftbox_utils import ( _resolve_email, @@ -299,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. @@ -311,13 +313,15 @@ def delete_remote_syftbox( this is auto-detected if not provided. verbose: If True (default), print deletion progress. """ + from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.drive_owner_store import DriveOwnerStore - from syft.sync.connections.drive.gdrive_transport import GDriveConnection email = _resolve_email(email) token_path = _resolve_token_path(token_path) - conn = GDriveConnection.from_token_path(email=email, token_path=token_path) - store = DriveOwnerStore(conn.cached_client) + 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(store.gather_all_file_and_folder_ids()) @@ -328,7 +332,7 @@ def delete_remote_syftbox( all_file_ids = list(folder_file_ids | orphaned_file_ids) start = time.time() - conn.cached_client.client.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: diff --git a/syft/sync/connections/base_connection.py b/syft/sync/connections/base_connection.py index da2ccc33e37..d1b297089fa 100644 --- a/syft/sync/connections/base_connection.py +++ b/syft/sync/connections/base_connection.py @@ -1,4 +1,3 @@ -from typing import ClassVar, Type from pydantic import BaseModel @@ -16,12 +15,5 @@ class FileCollection(BaseModel): class ConnectionConfig(BaseModel): - connection_type: ClassVar[Type["SyftboxPlatformConnection"]] - - -class SyftboxPlatformConnection(BaseModel): - config: ConnectionConfig | None = None - - @classmethod - def from_config(cls, config: ConnectionConfig): - return config.connection_type.from_config(config) + """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 index 785db426de9..4d6d31d0bcb 100644 --- a/syft/sync/connections/drive/drive_cached_client.py +++ b/syft/sync/connections/drive/drive_cached_client.py @@ -1,25 +1,36 @@ -"""Where syft keeps its folders in one account's Google Drive. +"""One Google Drive account, and where syft keeps its folders in it. -``DriveCachedClient`` 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. +``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" @@ -179,9 +190,9 @@ def valid_events_from_file_metadatas( class DriveCachedClient: """The Drive client of one account, and where syft's folders are in it. - It 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. + 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 @@ -199,6 +210,63 @@ class DriveCachedClient: # 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.""" diff --git a/syft/sync/connections/drive/gdrive_client.py b/syft/sync/connections/drive/gdrive_client.py index f71db155dd9..46db30bd4aa 100644 --- a/syft/sync/connections/drive/gdrive_client.py +++ b/syft/sync/connections/drive/gdrive_client.py @@ -1,23 +1,28 @@ """The Google Drive API client that the Drive transport jobs share. -One client per connection. It holds the Drive service, which is not safe to +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 +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. @@ -27,6 +32,39 @@ 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 diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py deleted file mode 100644 index b1c13a5d227..00000000000 --- a/syft/sync/connections/drive/gdrive_transport.py +++ /dev/null @@ -1,152 +0,0 @@ -"""A Google Drive account as a syft connection.""" - -import logging -from functools import partial -from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable - -from google.oauth2.credentials import Credentials as GoogleCredentials -from google_auth_httplib2 import AuthorizedHttp -from googleapiclient.discovery import build -from googleapiclient.http import build_http - -from syft.sync.connections.base_connection import ( - SyftboxPlatformConnection, -) -from syft.sync.connections.drive.gdrive_client import GDriveClient -from syft.sync.connections.drive.drive_cached_client import DriveCachedClient -from syft.sync.environments.environment import Environment -from syft.sync.utils.syftbox_utils import check_env - -if TYPE_CHECKING: - from syft.sync.connections.drive.grdrive_config import ( - GdriveConnectionConfig, - ) - from syft.sync.connections.drive.mock_drive_service import MockDriveService - -# 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) - - -class GDriveConnection(SyftboxPlatformConnection): - """A Google Drive account as a syft connection. - - It logs in and builds the ``DriveCachedClient`` that the transport components - share: ``SyftboxManager._drive_transport`` builds them over it. - """ - - class Config: - arbitrary_types_allowed = True - - cached_client: DriveCachedClient | None = None - verbose: bool = True - email: str - token_path: Path | None = None - _is_setup: bool = False - - @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": - """A connection with the OAuth token at ``token_path``, or with Colab's - own login when ``token_path`` is None.""" - 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( - partial(build_drive_service, credentials, environment=res.environment), - warm_syftbox_folder_id_cache=warm_syftbox_folder_id_cache, - ) - return res - - @classmethod - def from_service( - cls, - email: str, - mock_service: "MockDriveService", - warm_syftbox_folder_id_cache: bool = True, - ) -> "GDriveConnection": - """A connection for ``email`` over the backing store of ``mock_service``, - for tests. - - Args: - email: Email of the user - mock_service: MockDriveService whose backing store to use - warm_syftbox_folder_id_cache: See ``setup``. - """ - from syft.sync.connections.drive.mock_drive_service import ( - MockDriveService, - ) - - res = cls(email=email, token_path=None) - res.setup( - partial(MockDriveService, mock_service._backing_store, email), - warm_syftbox_folder_id_cache=warm_syftbox_folder_id_cache, - ) - return res - - def setup( - self, - new_service: Callable[[], Any], - warm_syftbox_folder_id_cache: bool = True, - ): - """Build the ``DriveCachedClient`` over a Drive service from - ``new_service``, which also builds the service of each copy for another - thread. - - Args: - new_service: Builds a Drive service, real or mock. - 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. - """ - self.cached_client = DriveCachedClient( - GDriveClient(new_service(), new_service), self.email - ) - if warm_syftbox_folder_id_cache: - self.cached_client.get_personal_syftbox_folder_id() - self._is_setup = True - - @property - def environment(self) -> Environment: - return check_env() 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 162fd763176..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,8 +17,8 @@ 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 @@ -29,7 +29,7 @@ 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" @@ -989,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 @@ -999,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 5025fefc9a1..e3f5d028e81 100644 --- a/syft/sync/login_utils.py +++ b/syft/sync/login_utils.py @@ -19,11 +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.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.drive_peering import DrivePeering - from syft.sync.connections.drive.gdrive_transport import GDriveConnection - conn = GDriveConnection.from_token_path(email=email, token_path=token_path) - peer_info = DrivePeering(conn.cached_client).read_own_peer_info() + 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 a4367fbeb6b..60f4847eeb4 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -13,6 +13,7 @@ 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 ( DriveWatcherStore, @@ -21,7 +22,6 @@ 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_transport import GDriveConnection 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.cached_client), - owner_store=DriveOwnerStore(drive.cached_client), + p2p=DriveP2PConnection(cached_client), + owner_store=DriveOwnerStore(cached_client), watcher_store=DriveWatcherStore(), - collections=DriveCollections(drive.cached_client), - peering=DrivePeering(drive.cached_client), + 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 @@ -1330,8 +1330,8 @@ def _copy(self): ).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, 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/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index c888eeec713..ede30eddc23 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -752,8 +752,8 @@ def test_owner_state_is_written_by_default(): def test_latest_checkpoint_read_retries_drive_error(monkeypatch): - do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.cached_client) + 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) @@ -761,8 +761,8 @@ def test_latest_checkpoint_read_retries_drive_error(monkeypatch): def test_incremental_count_retries_drive_error(monkeypatch): - do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.cached_client) + 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) @@ -770,8 +770,8 @@ def test_incremental_count_retries_drive_error(monkeypatch): def test_latest_checkpoint_download_error_raises(monkeypatch): - do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.cached_client) + 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: diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index ef38cdffacc..2ae768ca054 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -73,17 +73,17 @@ def _events_name(ts: float) -> str: @pytest.fixture def pair(): """DO and DS cached clients over one mock Drive, with the DS's channels open.""" - do_conn, ds_conn = pair_with_mock_service(DO, DS) - DrivePeering(ds_conn.cached_client).open_channels(DO) - return do_conn.cached_client, ds_conn.cached_client + 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) @@ -566,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_policy.py b/tests/unit/test_encryption_policy.py index 588d7833823..63f963b0e56 100644 --- a/tests/unit/test_encryption_policy.py +++ b/tests/unit/test_encryption_policy.py @@ -170,16 +170,14 @@ 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" ) - DrivePeering(a_conn.cached_client).grant_peer_info_read( + DrivePeering(a_cached_client).grant_peer_info_read( "b@example.com", _peer_info(REQUIRED) ) assert ( - DrivePeering(b_conn.cached_client) - .read_peer_info("a@example.com") - .encryption_policy + DrivePeering(b_cached_client).read_peer_info("a@example.com").encryption_policy == REQUIRED ) diff --git a/tests/unit/test_gdrive_client.py b/tests/unit/test_gdrive_client.py index 957261e2ef8..8dc8226cc48 100644 --- a/tests/unit/test_gdrive_client.py +++ b/tests/unit/test_gdrive_client.py @@ -12,9 +12,9 @@ def test_batch_delete_raises_error_whose_text_contains_404(monkeypatch): - do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - folder_id = do_conn.cached_client.client.create_folder( - "folder", do_conn.cached_client.get_syftbox_folder_id() + 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: @@ -22,13 +22,13 @@ def forbidden(self) -> None: monkeypatch.setattr(MockDeleteRequest, "execute", forbidden) with pytest.raises(HttpError): - do_conn.cached_client.client.delete_multiple_files_by_ids([folder_id]) + do_cached_client.client.delete_multiple_files_by_ids([folder_id]) def test_batch_share_raises_error_whose_text_says_already_shared(monkeypatch): - do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - folder_id = do_conn.cached_client.client.create_folder( - "folder", do_conn.cached_client.get_syftbox_folder_id() + 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: @@ -36,6 +36,4 @@ def forbidden(self) -> None: monkeypatch.setattr(MockPermissionCreateRequest, "execute", forbidden) with pytest.raises(HttpError): - do_conn.cached_client.client.batch_add_permissions( - folder_id, ["ds@example.com"] - ) + 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 726903c4611..3355b013004 100644 --- a/tests/unit/test_gdrive_pagination.py +++ b/tests/unit/test_gdrive_pagination.py @@ -6,7 +6,6 @@ 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.gdrive_transport import GDriveConnection from syft.sync.connections.drive.mock_drive_service import ( MockDriveBackingStore, MockDriveService, @@ -258,8 +257,7 @@ def test_newer_message_listed_after_older_one_in_same_second_not_skipped(): def _mock_cached_clients(*emails: str) -> list[DriveCachedClient]: store = MockDriveBackingStore() return [ - GDriveConnection.from_service(e, MockDriveService(store, e)).cached_client - for e in emails + DriveCachedClient.from_service(e, MockDriveService(store, e)) for e in emails ] diff --git a/tests/unit/test_mock_drive_service.py b/tests/unit/test_mock_drive_service.py index 8a3b97d21b2..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,27 +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.cached_client.create_syftbox_folder() - folder_id = connection.cached_client.client.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"} @@ -468,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.cached_client.client.download_file(file.id) + content = cached_client.client.download_file(file.id) assert content == b"Test content" diff --git a/tests/unit/test_rolling_state.py b/tests/unit/test_rolling_state.py index 629656e1e65..b3723dedf20 100644 --- a/tests/unit/test_rolling_state.py +++ b/tests/unit/test_rolling_state.py @@ -288,8 +288,8 @@ def test_no_rolling_state_when_owner_state_disabled(): def test_rolling_state_read_retries_drive_error(monkeypatch): - do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.cached_client) + 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) @@ -297,8 +297,8 @@ def test_rolling_state_read_retries_drive_error(monkeypatch): def test_rolling_state_read_while_drive_is_down_raises(monkeypatch): - do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.cached_client) + 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: @@ -310,8 +310,8 @@ def drive_down(self, file_id: str) -> bytes: def test_rolling_state_listing_while_drive_is_down_raises(monkeypatch): - do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.cached_client) + 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: @@ -323,8 +323,8 @@ def drive_down(self) -> dict: def test_rolling_state_download_error_raises(monkeypatch): - do_conn, _ = pair_with_mock_service("do@example.com", "ds@example.com") - store = DriveOwnerStore(do_conn.cached_client) + 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: diff --git a/tests/unit/test_shared_router.py b/tests/unit/test_shared_router.py index 93a25371b50..add45eb5bf8 100644 --- a/tests/unit/test_shared_router.py +++ b/tests/unit/test_shared_router.py @@ -11,7 +11,6 @@ drive_cached_client_of, ) from syft.sync.connections.drive.drive_cached_client import DriveCachedClient -from syft.sync.connections.drive.gdrive_transport import GDriveConnection 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 @@ -29,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 @@ -63,20 +62,20 @@ def _drive_cached_clients(transport) -> list[DriveCachedClient]: 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 + [built_client] = built for cached_client in _drive_cached_clients(manager._connection_router.transport): - assert cached_client is connection.cached_client + assert cached_client is built_client def test_every_drive_component_shares_one_cached_client(): @@ -89,7 +88,7 @@ def test_every_drive_component_shares_one_cached_client(): def test_drive_cached_client_of_transport_not_on_drive_raises(): - transport = SyftboxManager._drive_transport(_mock_connection("do@example.com")) + 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"): diff --git a/tests/unit/test_sync_manager.py b/tests/unit/test_sync_manager.py index f9279ca8f99..07afb3083c0 100644 --- a/tests/unit/test_sync_manager.py +++ b/tests/unit/test_sync_manager.py @@ -5,7 +5,7 @@ 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 @@ -486,16 +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, drive_cached_client_of( ds_manager._connection_router.transport ).client.drive_service, ) - DrivePeering(conn.cached_client).open_channels(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/utils.py b/tests/unit/utils.py index 0b62748dae1..798fcc31f0f 100644 --- a/tests/unit/utils.py +++ b/tests/unit/utils.py @@ -249,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), )