diff --git a/packages/syft-enclave/src/syft_enclaves/client.py b/packages/syft-enclave/src/syft_enclaves/client.py index 99a1772f900..05e245f710b 100644 --- a/packages/syft-enclave/src/syft_enclaves/client.py +++ b/packages/syft-enclave/src/syft_enclaves/client.py @@ -15,6 +15,7 @@ from syft_rds import SyftRDSClient, SyftRDSClientConfig from syft_rds.waiting import JobCheck, JobStatusArg, wait_for_job +from syft.sync.transport.protocols import Transport from syft.sync.encryption_policy import EncryptionPolicy from syft.sync.peers.peer import Peer from syft.sync.peers.peer_list import PeerList @@ -45,7 +46,6 @@ create_clients, create_configs, setup_callbacks, - setup_connections, wire_peers, write_versions, ) @@ -783,6 +783,7 @@ def from_config( cls, config: SyftRDSClientConfig, data_owners: list[str] | None = None, + transport: Transport | None = None, ) -> "SyftEnclaveClient": """Build a SyftEnclaveClient from an RDS config with a wrapped job_client. @@ -790,8 +791,10 @@ def from_config( ``peer_manager_config.encryption_policy``; ``SyftboxManager.from_config`` (invoked inside ``SyftRDSClient.from_config``) loads/generates this datasite's keys from its per-datasite key file, so no extra step is needed. + ``transport`` replaces the one ``config.sync`` describes; tests pass + one over a mock Drive. """ - rds = SyftRDSClient.from_config(config) + rds = SyftRDSClient.from_config(config, transport) rds.job_client = EnclaveJobClient(rds.job_client) sync_engine = rds.sync_engine @@ -885,12 +888,10 @@ def quad_with_mock_drive_service_connection( enclave, do1, do2, ds = clients enclave.data_owners = [do1.email, do2.email] # The setup helpers manipulate the generic sync engine directly - # (connections, file-watcher callbacks, version files, encryption, - # peering). Each client's sync engine lives at ``rds.sync_engine``; the + # (file-watcher callbacks, version files, encryption, peering). Each client's sync engine lives at ``rds.sync_engine``; the # RDS-owned job/dataset managers react via the peer-lifecycle callbacks # registered on that same engine during SyftRDSClient.__init__. managers = tuple(c._rds.sync_engine for c in clients) - setup_connections(managers) setup_callbacks(managers) write_versions(managers) wire_peers(managers) diff --git a/packages/syft-enclave/src/syft_enclaves/utils.py b/packages/syft-enclave/src/syft_enclaves/utils.py index a63c1b97aca..553c74bf1eb 100644 --- a/packages/syft-enclave/src/syft_enclaves/utils.py +++ b/packages/syft-enclave/src/syft_enclaves/utils.py @@ -4,6 +4,8 @@ MockDriveService, ) from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.transport.protocols import Transport def create_configs( @@ -49,18 +51,23 @@ def create_configs( def create_clients(configs) -> tuple: + """One client per config, all on one mock Drive.""" from syft_enclaves.client import SyftEnclaveClient - return tuple(SyftEnclaveClient.from_config(c) for c in configs) - - -def setup_connections(managers: tuple): backing_store = MockDriveBackingStore() + return tuple( + SyftEnclaveClient.from_config( + c, transport=_mock_transport(backing_store, c.sync.email) + ) + for c in configs + ) - for manager in managers: - service = MockDriveService(backing_store, manager.email) - connection = GDriveConnection.from_service(manager.email, service) - manager._add_connection(connection) + +def _mock_transport(backing_store: MockDriveBackingStore, email: str) -> Transport: + service = MockDriveService(backing_store, email) + return SyftboxManager._drive_transport( + GDriveConnection.from_service(email, service) + ) def setup_callbacks(managers: tuple): diff --git a/packages/syft-rds/src/syft_rds/client.py b/packages/syft-rds/src/syft_rds/client.py index 0066807958d..a27519c6392 100644 --- a/packages/syft-rds/src/syft_rds/client.py +++ b/packages/syft-rds/src/syft_rds/client.py @@ -13,6 +13,7 @@ from pydantic import BaseModel, ConfigDict from syft_migration import ProtocolSchema +from syft.sync.transport.protocols import Transport from syft.sync.syftbox_manager import SyftboxManager from syft.sync.peers.peer import Peer from syft.sync.utils.waiting import DEFAULT_POLL_INTERVAL, DEFAULT_WAIT_TIMEOUT @@ -148,8 +149,14 @@ def model_post_init(self, __context: Any) -> None: ) @classmethod - def from_config(cls, config: "SyftRDSClientConfig") -> "SyftRDSClient": - sync_engine = SyftboxManager.from_config(config.sync) + def from_config( + cls, + config: "SyftRDSClientConfig", + transport: Transport | None = None, + ) -> "SyftRDSClient": + """``transport`` replaces the one ``config.sync`` describes; tests pass + one over a mock Drive (see ``SyftboxManager.from_config``).""" + sync_engine = SyftboxManager.from_config(config.sync, transport) # The job client selects a job protocol version per peer from the # schema the peer advertises, as the dataset manager does for layouts. job_client = JobClient.from_config( diff --git a/packages/syft-rds/tests/test_dataset_upload_private.py b/packages/syft-rds/tests/test_dataset_upload_private.py index ef90e08f0d1..4f3b6a057a6 100644 --- a/packages/syft-rds/tests/test_dataset_upload_private.py +++ b/packages/syft-rds/tests/test_dataset_upload_private.py @@ -8,6 +8,7 @@ PRIVATE_DATASET_COLLECTION_PREFIX, ) from dataset_test_utils import create_tmp_dataset_files +from syft.sync.connections.drive.drive_connections import gdrive_connection_of class TestDatasetUploadPrivate: @@ -145,8 +146,8 @@ def test_ds_cannot_find_private_folders_via_gdrive_query(self): ) # Get the DS's raw GDrive connection and search for private folders - ds_connection: GDriveConnection = ( - ds_manager.sync_engine._connection_router.connections[0] + ds_connection: GDriveConnection = gdrive_connection_of( + ds_manager.sync_engine._connection_router.transport ) results = ( ds_connection.drive_service.files() @@ -159,8 +160,8 @@ def test_ds_cannot_find_private_folders_via_gdrive_query(self): assert len(results.get("files", [])) == 0 # DO should find it with the same query - do_connection: GDriveConnection = ( - do_manager.sync_engine._connection_router.connections[0] + do_connection: GDriveConnection = gdrive_connection_of( + do_manager.sync_engine._connection_router.transport ) do_results = ( do_connection.drive_service.files() diff --git a/packages/syft-rds/tests/test_sync_manager.py b/packages/syft-rds/tests/test_sync_manager.py index d375856d18e..6d83bb46852 100644 --- a/packages/syft-rds/tests/test_sync_manager.py +++ b/packages/syft-rds/tests/test_sync_manager.py @@ -9,6 +9,7 @@ from syft.sync.connections.drive import mock_drive_service from syft.sync.connections.drive.gdrive_transport import GDriveConnection +from syft.sync.syftbox_manager import SyftboxManager from syft_datasets import Dataset from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX from syft_rds import SyftRDSClient @@ -21,6 +22,7 @@ create_tmp_dataset_files, create_tmp_dataset_files_with_parquet, ) +from syft.sync.connections.drive.drive_connections import gdrive_connection_of def path_for_job(do_email: str, ds_email: str, filename: str = "test.job") -> str: @@ -231,9 +233,9 @@ def test_dataset_only_mock_data_uploaded(): # Sync so DS receives the dataset ds_manager.sync() - files = ds_manager.sync_engine._connection_router.connections[ - 0 - ].drive_service._backing_store.files + files = gdrive_connection_of( + ds_manager.sync_engine._connection_router.transport + ).drive_service._backing_store.files list(files.keys()) file_objs = list(files.values()) file_objs_ex_dataset_yaml = [ @@ -282,7 +284,9 @@ def test_jobs(): # we do this by not always syncing on a file change, currently this logic is a bit of # a short cut, but we could do this based on timing eventually (if there are items in the # queue for longer than a certain time we start pushing) - connection_do = do_manager.sync_engine._connection_router.connections[0] + connection_do = gdrive_connection_of( + do_manager.sync_engine._connection_router.transport + ) inbox_folder_id = connection_do._get_own_datasite_inbox_id(ds_manager.email) inbox_file_metadatas = connection_do.get_file_metadatas_from_folder(inbox_folder_id) assert len(inbox_file_metadatas) == 1 @@ -1036,9 +1040,9 @@ def test_ds_dataset_cache_aware_sync(): assert remote_hash is not None # Get the mock backing store and directories for creating second pair - mock_backing_store = ds_manager.sync_engine._connection_router.connections[ - 0 - ].drive_service._backing_store + mock_backing_store = gdrive_connection_of( + ds_manager.sync_engine._connection_router.transport + ).drive_service._backing_store ds_folder = ds_manager.syftbox_folder do_folder = do_manager.syftbox_folder ds_email = ds_manager.email @@ -1055,12 +1059,12 @@ 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 - do_manager2.sync_engine._connection_router.connections[ - 0 - ].drive_service._backing_store = mock_backing_store + gdrive_connection_of( + ds_manager2.sync_engine._connection_router.transport + ).drive_service._backing_store = mock_backing_store + gdrive_connection_of( + do_manager2.sync_engine._connection_router.transport + ).drive_service._backing_store = mock_backing_store # Load peers (already approved in shared backing store) ds_manager2.load_peers() @@ -1136,9 +1140,9 @@ def test_do_dataset_cache_aware_sync(): assert len(do_manager.datasets.get_all()) == 1 # Get the mock backing store and directories for creating second pair - mock_backing_store = ds_manager.sync_engine._connection_router.connections[ - 0 - ].drive_service._backing_store + mock_backing_store = gdrive_connection_of( + ds_manager.sync_engine._connection_router.transport + ).drive_service._backing_store ds_folder = ds_manager.syftbox_folder do_folder = do_manager.syftbox_folder ds_email = ds_manager.email @@ -1155,12 +1159,12 @@ 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 - do_manager2.sync_engine._connection_router.connections[ - 0 - ].drive_service._backing_store = mock_backing_store + gdrive_connection_of( + ds_manager2.sync_engine._connection_router.transport + ).drive_service._backing_store = mock_backing_store + gdrive_connection_of( + do_manager2.sync_engine._connection_router.transport + ).drive_service._backing_store = mock_backing_store # Load peers (already approved in shared backing store) ds_manager2.load_peers() @@ -1203,15 +1207,13 @@ 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 + backing_store = gdrive_connection_of( + do_manager1.sync_engine._connection_router.transport + ).drive_service._backing_store ds_folder = ds_manager1.syftbox_folder do_folder = do_manager1.syftbox_folder ds_email = ds_manager1.email do_email = do_manager1.email - ds_manager1.sync_engine._connection_router.connections[0] - do_manager1.sync_engine._connection_router.connections[0] # Make some changes (submit to DS's job folder where DS has write access) ds_manager1.sync_engine._send_file_change( @@ -1265,9 +1267,6 @@ def test_in_memory_connection_load_state(): check_versions=False, ) - do_manager2 = SyftRDSClient.from_config(do_config2) - ds_manager2 = SyftRDSClient.from_config(ds_config2) - # 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) @@ -1275,8 +1274,12 @@ def test_in_memory_connection_load_state(): service_ds = mock_drive_service.MockDriveService(backing_store, ds_email) ds_connection2 = GDriveConnection.from_service(ds_email, service_ds) - do_manager2.sync_engine._add_connection(do_connection2) - ds_manager2.sync_engine._add_connection(ds_connection2) + do_manager2 = SyftRDSClient.from_config( + do_config2, SyftboxManager._drive_transport(do_connection2) + ) + ds_manager2 = SyftRDSClient.from_config( + ds_config2, SyftboxManager._drive_transport(ds_connection2) + ) # Load peers do_manager2.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 5ebb44ae328..f41865e1a36 100644 --- a/packages/syft-rds/tests/test_version_mismatch_flow.py +++ b/packages/syft-rds/tests/test_version_mismatch_flow.py @@ -20,6 +20,8 @@ from syft.version import SYFT_VERSION from dataset_test_utils import create_test_project_folder, create_tmp_dataset_files +from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.syftbox_manager import SyftboxManager NEW_VERSION = "99.0.0" @@ -46,7 +48,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] + conn = gdrive_connection_of(manager.peer_manager.connection_router.transport) return conn.drive_service._backing_store @@ -66,11 +68,9 @@ def _reinitialize_manager( use_in_memory_cache=False, syftbox_folder=syftbox_folder, ) - manager = SyftRDSClient.from_config(config) - mock_service = MockDriveService(backing_store, email) conn = GDriveConnection.from_service(email, mock_service) - manager.sync_engine._add_connection(conn) + manager = SyftRDSClient.from_config(config, SyftboxManager._drive_transport(conn)) if has_ds_role: manager.sync_engine.file_writer.add_callback( @@ -172,7 +172,7 @@ def test_version_mismatch_continues_and_repairs(): assert do_manager.jobs[0].status == "done" # -- Step 4: Record P2P folders and personal folder id before upgrade -- - do_conn = do_manager.peer_manager.connection_router.connections[0] + do_conn = gdrive_connection_of(do_manager.peer_manager.connection_router.transport) do_p2p_current = _find_versioned_p2p_folders( do_conn, ds_manager.email, SYFT_VERSION ) @@ -211,7 +211,9 @@ 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] + do_conn_new = gdrive_connection_of( + do_manager.peer_manager.connection_router.transport + ) new_personal_name = f"{NEW_VERSION}#{do_email}" new_personal_id = do_conn_new._find_folder_by_name( new_personal_name, @@ -247,7 +249,9 @@ def test_version_mismatch_continues_and_repairs(): syftbox_folder=ds_syftbox, ) - ds_conn_new = ds_manager.peer_manager.connection_router.connections[0] + ds_conn_new = gdrive_connection_of( + ds_manager.peer_manager.connection_router.transport + ) ds_personal_name = f"{NEW_VERSION}#{ds_email}" ds_personal_id = ds_conn_new._find_folder_by_name( ds_personal_name, @@ -327,7 +331,7 @@ def test_login_writes_the_remote_version_file_too(): use_in_memory_cache=False, check_versions=True, ) - conn = do_manager.peer_manager.connection_router.connections[0] + conn = gdrive_connection_of(do_manager.peer_manager.connection_router.transport) assert ( conn.read_own_peer_info_file().version_info.syft_client_version == SYFT_VERSION ) @@ -337,7 +341,9 @@ def test_login_writes_the_remote_version_file_too(): patch("syft.sync.version.version_info.SYFT_VERSION", NEW_VERSION), ): upgraded = _upgraded_manager(do_manager) - new_conn = upgraded.peer_manager.connection_router.connections[0] + new_conn = gdrive_connection_of( + upgraded.peer_manager.connection_router.transport + ) # Still the pre-upgrade version: nothing has refreshed it yet. assert ( new_conn.read_own_peer_info_file().version_info.syft_client_version @@ -388,7 +394,7 @@ def test_the_mismatch_prompt_does_not_return_after_a_login(): handle_potential_version_mismatches_on_login, ) - conn = do_manager.peer_manager.connection_router.connections[0] + conn = gdrive_connection_of(do_manager.peer_manager.connection_router.transport) with patch( "syft.sync.login_utils._read_remote_version", side_effect=lambda e, t: conn.read_own_peer_info_file().version_info, diff --git a/syft/sync/connections/base_connection.py b/syft/sync/connections/base_connection.py index e81c43fc18f..5ed17aae862 100644 --- a/syft/sync/connections/base_connection.py +++ b/syft/sync/connections/base_connection.py @@ -1,10 +1,7 @@ -from typing import TYPE_CHECKING, ClassVar, Type +from typing import ClassVar, Type from pydantic import BaseModel from syft.sync.messages.proposed_filechange import ProposedFileChangesMessage -if TYPE_CHECKING: - from syft.sync.transport.protocols import Transport - class FileCollection(BaseModel): """A collection of files that can be shared and synced.""" @@ -39,11 +36,6 @@ def watcher_send_proposed_file_changes_message( def from_config(cls, config: ConnectionConfig): return config.connection_type.from_config(config) - def as_transport(self) -> "Transport": - """This connection as the transport jobs: P2P messages, owner and - watcher storage, collections and peering.""" - raise NotImplementedError() - def clone_for_thread(self) -> "SyftboxPlatformConnection": """A connection that another thread may use. diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index 65bec3b9101..8a77f684c82 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 ( @@ -12,11 +12,7 @@ ) from syft.sync.checkpoints.checkpoint import Checkpoint, IncrementalCheckpoint from syft.sync.checkpoints.rolling_state import RollingState -from syft.sync.connections.base_connection import ( - ConnectionConfig, - FileCollection, - SyftboxPlatformConnection, -) +from syft.sync.connections.base_connection import FileCollection from syft.sync.encryption_policy import ( MessageFormatError, ) @@ -43,66 +39,18 @@ class ConnectionRouter(BaseModel): - connections: List[SyftboxPlatformConnection] + """The one router of a manager. ``SyftboxManager.from_config`` builds it + and hands it to every part, so they share one transport, and through it + one connection, and one peer store.""" - peer_store: PeerStore + # Built once with the router and never replaced. A test may swap in a + # component that does not use Google Drive, such as a fake p2p. + transport: Transport - # 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 - - @classmethod - def from_configs( - cls, connection_configs: List[ConnectionConfig], peer_store: PeerStore - ) -> "ConnectionRouter": - """The one router of a manager. ``SyftboxManager.from_config`` builds - it and hands it to every part, so they share one connection, one - transport and one peer store.""" - return cls( - connections=[ - SyftboxPlatformConnection.from_config(config) - for config in connection_configs - ], - peer_store=peer_store, - ) - - def add_connection(self, connection: SyftboxPlatformConnection): - self.connections.append(connection) + peer_store: PeerStore def get_authenticated_email(self) -> str: - return self.connections[0].get_authenticated_email() - - def copy_connection( - self, connection: SyftboxPlatformConnection - ) -> 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] - - def connection_for_outbox(self) -> SyftboxPlatformConnection: - return self.connections[0] - - def connection_for_own_syftbox(self) -> SyftboxPlatformConnection: - return self.connections[0] + return self.transport.peering.whoami() # ========================================================================= # MESSAGE SEND/RECEIVE (with encryption) @@ -246,9 +194,12 @@ def download_outbox_message(self, file_id: str) -> bytes: """Download one outbox message by its handle (thread-safe).""" return self.transport.p2p.clone_for_thread().fetch_message(file_id) - 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 owner_get_outbox_file_metadatas(self, recipient_email: str) -> List[dict]: + """The messages in my outbox to ``recipient_email``, without downloading.""" + return self.transport.p2p.own_outbox_message_metadatas(recipient_email) + + def owner_delete_outbox_messages(self, file_ids: List[str]) -> None: + self.transport.p2p.delete_messages(file_ids) # ========================================================================= # EVENT LOG (self-encrypted at rest when encryption enabled) @@ -266,8 +217,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.owner_store.clone_for_thread().fetch(events_message_id) raw = self.peer_store.decrypt_and_verify_for_self_if_needed(raw) return FileChangeEventsMessage.from_compressed_data(raw) @@ -418,19 +368,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) @@ -465,40 +402,22 @@ def watcher_download_collection_file(self, file_id: str, owner_email: str) -> by # MISC # ========================================================================= - def delete_multiple_files_by_ids( - self, - file_ids: List[str], - ignore_permissions_errors: bool = True, - ignore_file_not_found: bool = True, - ): - connection = self.connection_for_own_syftbox() - connection.delete_multiple_files_by_ids( - file_ids, - ignore_permissions_errors=ignore_permissions_errors, - ignore_file_not_found=ignore_file_not_found, - ) - def owner_get_all_accepted_event_file_ids( self, since_timestamp: float | None = None ) -> List[str]: return self.transport.owner_store.list_event_handles(since_timestamp) - def gather_all_file_and_folder_ids(self) -> List[str]: - connection = self.connection_for_own_syftbox() - return connection.gather_all_file_and_folder_ids() - - def find_orphaned_syft_files(self) -> List[str]: - """Find the orphaned syft files of this account in its My Drive root.""" - connection = self.connection_for_own_syftbox() - return connection.find_orphaned_syft_files() + def delete_all(self) -> int: + """Delete everything this account stores; return how many files and + folders went.""" + return self.transport.owner_store.delete_all() def reset_caches(self): - connection = self.connection_for_own_syftbox() - connection.reset_caches() + for component in self.transport.components(): + component.reset_caches() def delete_file_by_id(self, file_id: str): - connection = self.connection_for_own_syftbox() - connection.delete_file_by_id(file_id) + self.transport.collections.delete_by_handle(file_id) # ========================================================================= # CHECKPOINT METHODS diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index b85c64cd395..75118f968e5 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -13,10 +13,9 @@ import threading import warnings from concurrent.futures import ThreadPoolExecutor -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self from syft.sync.connections.drive.gdrive_retry import raise_unavailable -from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform from syft.sync.transport.exceptions import TransportUnavailableError from syft.sync.transport.protocols import ( Collections, @@ -65,7 +64,21 @@ def _require_inbox_kind(kind: MessageKind) -> None: ) -class DriveP2PConnection(P2PConnection): +class DriveTransportComponent: + """A transport component over one Google Drive connection.""" + + def __init__(self, connection: GDriveConnection): + self._conn = connection + + def clone_for_thread(self) -> Self: + """A copy safe to use from another thread: the Drive client is not.""" + return type(self)(self._conn.clone_for_thread()) + + def reset_caches(self) -> None: + self._conn.reset_caches() + + +class DriveP2PConnection(DriveTransportComponent, P2PConnection): """Messages through the shared inbox and outbox folders. A proposed change goes to the data owner's inbox and file events go to the @@ -79,9 +92,6 @@ class DriveP2PConnection(P2PConnection): data scientist keeps what it received in its own storage. """ - def __init__(self, connection: GDriveConnection): - self._conn = connection - def send( self, peer: str, msg_id: str, data: bytes, kind: MessageKind ) -> SendStatus: @@ -167,18 +177,24 @@ def outbox_message_metadatas( def fetch_message(self, handle: str) -> bytes: return self._conn.download_file(handle) - def clone_for_thread(self) -> DriveP2PConnection: - """A copy safe to use from another thread: the Drive client is not.""" - return DriveP2PConnection(self._conn.clone_for_thread()) + def own_outbox_message_metadatas(self, peer: str) -> list[dict]: + """The messages in my outbox to ``peer``, without downloading. An + outbox whose folder is not found holds none.""" + folder_id = self._conn._get_own_datasite_outbox_id(peer) + if folder_id is None: + return [] + return self._conn.get_file_metadatas_from_folder( + folder_id, since_timestamp=None + ) + + def delete_messages(self, handles: list[str]) -> None: + self._conn.delete_multiple_files_by_ids(handles) -class DriveOwnerStore(OwnerStore): +class DriveOwnerStore(DriveTransportComponent, 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) @@ -226,14 +242,19 @@ def get_rolling_state(self) -> bytes | None: def delete_rolling_state(self) -> None: self._conn.delete_rolling_state() + def delete_all(self) -> int: + """Everything in /SyftBox, and the syft files Drive moved to the root + of My Drive when they lost their folder.""" + file_ids = set(self._conn.gather_all_file_and_folder_ids()) + file_ids |= set(self._conn.find_orphaned_syft_files()) + self._conn.delete_multiple_files_by_ids(list(file_ids)) + return len(file_ids) + -class DriveWatcherStore(WatcherStore): +class DriveWatcherStore(DriveTransportComponent, WatcherStore): """The data scientist's storage on Drive. Not built yet: today the data scientist keeps what it received in a local cache.""" - def __init__(self, connection: GDriveConnection): - self._conn = connection - def put_checkpoint(self, peer: str, name: str, data: bytes) -> str: raise NotImplementedError("Watcher storage on Drive is not built yet") @@ -250,12 +271,9 @@ def reset(self, peer: str) -> None: raise NotImplementedError("Watcher storage on Drive is not built yet") -class DriveCollections(Collections): +class DriveCollections(DriveTransportComponent, 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 @@ -283,6 +301,9 @@ def list_mine(self, prefix: str) -> list[str]: def delete(self, prefix: str, tag: str) -> None: self._conn.owner_delete_collection(prefix, tag) + def delete_by_handle(self, handle: str) -> None: + self._conn.delete_file_by_id(handle) + def list_shared_with_me(self, prefix: str) -> list[dict]: return self._conn.watcher_list_collections(prefix) @@ -304,13 +325,10 @@ def download( ) -class DrivePeering(Peering): +class DrivePeering(DriveTransportComponent, 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() @@ -353,13 +371,11 @@ 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( - p2p=DriveP2PConnection(connection), - owner_store=DriveOwnerStore(connection), - watcher_store=DriveWatcherStore(connection), - collections=DriveCollections(connection), - peering=DrivePeering(connection), - platform=GdriveFilesPlatform(), - ) +def gdrive_connection_of(transport: Transport) -> GDriveConnection: + """The Google Drive connection under ``transport``, for tests and tools + that need Drive itself. The sync code reaches Drive only through the + components. Raises if the transport is not built on Drive.""" + store = transport.owner_store + if not isinstance(store, DriveTransportComponent): + raise TypeError("This transport is not built on Google Drive") + return store._conn diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 653099b1dd1..a589b9b23a7 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -31,7 +31,6 @@ FileCollection, SyftboxPlatformConnection, ) -from syft.sync.connections.drive.drive_connections import drive_transport from syft.sync.connections.drive.gdrive_retry import ( batch_execute_with_retries, execute_with_retries, @@ -454,10 +453,6 @@ def setup( self.get_personal_syftbox_folder_id() self._is_setup = True - def as_transport(self): - """This connection as the transport jobs. See ``drive_connections.py``.""" - return drive_transport(self) - def clone_for_thread(self) -> "GDriveConnection": """A connection safe to use from another thread: the Drive client is not.""" return self.copy() diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index 2623754ce06..a4c22fec7d9 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -11,12 +11,16 @@ from pydantic import BaseModel, ConfigDict, PrivateAttr from syft.sync.callback_mixin import BaseModelCallbackMixin -from syft.sync.connections.base_connection import ( - ConnectionConfig, - SyftboxPlatformConnection, -) from syft.sync.connections.connection_router import ConnectionRouter from syft.sync.connections.drive import mock_drive_service +from syft.sync.connections.drive.drive_connections import ( + DriveCollections, + DriveOwnerStore, + DriveP2PConnection, + DrivePeering, + DriveWatcherStore, + gdrive_connection_of, +) 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 ( @@ -56,6 +60,9 @@ random_syftbox_folder_for_testing, ) from syft.sync.peers.peer_store import PeerStore +from syft.sync.transport.config import TransportConfig +from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform +from syft.sync.transport.protocols import Transport from syft.sync.version.peer_manager import ( PeerManager, PeerManagerConfig, @@ -124,10 +131,11 @@ class SyftboxManagerConfig(BaseModel): # log, rolling state, checkpoints. Off for datasites that can never # restore (see syft-enclave's EnclaveSettings). persist_owner_state: bool = True - # The backends to connect to. Only the manager reads them: it builds one - # connection router from them and shares it with every part. Empty for - # tests that add a mock connection afterwards (see _add_connection). - connection_configs: List[ConnectionConfig] = [] + # The connection each transport component uses. Only the manager reads + # it: it builds one transport from it and shares that with every part. + # None when the caller passes a transport to SyftboxManager.from_config + # instead, as tests do with a mock Drive. + transport_config: TransportConfig | None = None datasite_owner_syncer_config: DatasiteOwnerSyncerConfig peer_manager_config: PeerManagerConfig @@ -162,7 +170,9 @@ def for_colab( syftbox_folder, email, collection_specs ) collection_variants = default_collection_variants(collection_specs) - connection_configs = [GdriveConnectionConfig(email=email, token_path=None)] + transport_config = TransportConfig.one_connection( + GdriveConnectionConfig(email=email, token_path=None) + ) datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, persist_owner_state=persist_owner_state, @@ -199,7 +209,7 @@ def for_colab( return cls( email=email, syftbox_folder=syftbox_folder, - connection_configs=connection_configs, + transport_config=transport_config, has_ds_role=has_ds_role, has_do_role=has_do_role, use_in_memory_cache=False, @@ -239,9 +249,9 @@ def for_jupyter( ) collection_variants = default_collection_variants(collection_specs) - connection_configs = [ + transport_config = TransportConfig.one_connection( GdriveConnectionConfig(email=email, token_path=token_path) - ] + ) datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, persist_owner_state=persist_owner_state, @@ -278,7 +288,7 @@ def for_jupyter( return cls( email=email, syftbox_folder=syftbox_folder, - connection_configs=connection_configs, + transport_config=transport_config, has_ds_role=has_ds_role, has_do_role=has_do_role, use_in_memory_cache=False, @@ -390,9 +400,9 @@ def for_google_drive_testing_connection( syftbox_folder, email, collection_specs ) collection_variants = default_collection_variants(collection_specs) - connection_configs = [ + transport_config = TransportConfig.one_connection( GdriveConnectionConfig(email=email, token_path=token_path) - ] + ) datasite_owner_syncer_config = DatasiteOwnerSyncerConfig( email=email, persist_owner_state=persist_owner_state, @@ -428,7 +438,7 @@ def for_google_drive_testing_connection( return cls( email=email, syftbox_folder=syftbox_folder, - connection_configs=connection_configs, + transport_config=transport_config, write_files=write_files, persist_owner_state=persist_owner_state, datasite_owner_syncer_config=datasite_owner_syncer_config, @@ -516,21 +526,28 @@ def peers(self) -> PeerList: return combined @classmethod - def from_config(cls, config: SyftboxManagerConfig): + def from_config( + cls, + config: SyftboxManagerConfig, + transport: Transport | None = None, + ): """Build a manager and its parts from ``config``. Configs hold plain values only (strings, paths, flags), never live objects or state. A live object that several parts share is built here, once, and passed to each part's ``from_config`` next to its config: the peer store and the connection router that holds it, and through - the router one connection and one transport for the whole manager. + the router one transport, and its connection, for the whole manager. + + ``transport`` replaces the one ``config`` describes. + Tests pass one over a mock Drive, some with a fake p2p. """ peer_store = PeerManager.peer_store_from_config( config.peer_manager_config, email=config.email ) - connection_router = ConnectionRouter.from_configs( - config.connection_configs, peer_store - ) + if transport is None: + transport = cls._transport_from_config(config) + connection_router = ConnectionRouter(transport=transport, peer_store=peer_store) file_writer = FileWriter( base_path=config.syftbox_folder, write_files=config.write_files ) @@ -569,6 +586,38 @@ def from_config(cls, config: SyftboxManagerConfig): ) return manager_res + @staticmethod + def _transport_from_config(config: SyftboxManagerConfig) -> Transport: + transport_config = config.transport_config + if transport_config is None: + raise ValueError( + "The config names no transport: set transport_config, or pass a " + "transport to SyftboxManager.from_config" + ) + drive_config = transport_config.p2p + if not ( + isinstance(drive_config, GdriveConnectionConfig) + and transport_config.uses_one_connection() + ): + raise NotImplementedError( + "Only one Google Drive connection for every transport component " + "is built so far" + ) + drive = GDriveConnection.from_config(drive_config) + return SyftboxManager._drive_transport(drive) + + @staticmethod + def _drive_transport(drive: GDriveConnection) -> Transport: + """Every transport component over one Google Drive connection.""" + return Transport( + p2p=DriveP2PConnection(drive), + owner_store=DriveOwnerStore(drive), + watcher_store=DriveWatcherStore(drive), + collections=DriveCollections(drive), + peering=DrivePeering(drive), + platform=GdriveFilesPlatform(), + ) + @classmethod def for_colab( cls, @@ -780,18 +829,15 @@ def pair_with_mock_drive_service_connection( encryption_policy=ds_encryption_policy or default_policy, ) - # Create managers from configs - do_manager = cls.from_config(do_config) - ds_manager = cls.from_config(ds_config) - - # Create GDriveConnection instances with mock services do_connection, ds_connection = mock_drive_service.pair_with_mock_service( - do_manager.email, ds_manager.email + do_config.email, ds_config.email + ) + do_manager = cls.from_config( + do_config, SyftboxManager._drive_transport(do_connection) + ) + ds_manager = cls.from_config( + ds_config, SyftboxManager._drive_transport(ds_connection) ) - - # Add connections to managers - do_manager._add_connection(do_connection) - ds_manager._add_connection(ds_connection) # Set up callbacks for DS -> DO communication ds_manager.file_writer.add_callback( @@ -1062,19 +1108,6 @@ def wait_until_peered( f"Timed out after {timeout:g}s waiting for peer {peer_email}: {error}" ) from error - def _add_connection(self, connection: SyftboxPlatformConnection): - if not ( - isinstance(connection, GDriveConnection) - and isinstance( - connection.drive_service, mock_drive_service.MockDriveService - ) - ): - raise ValueError( - "Only MockDriveService connections can be added to the manager" - ) - - self._connection_router.add_connection(connection) - def _send_file_change(self, path: str | Path, content: str): self.file_writer.write_file(path, content) @@ -1166,28 +1199,10 @@ def delete_syftbox( peer_emails = [p.email for p in self.peer_manager.approved_peers] file_hashes = dict(self.datasite_owner_syncer.event_cache.file_hashes) - # Get files by folder hierarchy - folder_file_ids = set(self._connection_router.gather_all_file_and_folder_ids()) - - # Also find the orphaned syft files in My Drive root - orphaned_file_ids = set(self._connection_router.find_orphaned_syft_files()) - - # Combine both sets and delete from Google Drive - all_file_ids = list(folder_file_ids | orphaned_file_ids) - start = time.time() - self._connection_router.delete_multiple_files_by_ids(all_file_ids) - end = time.time() + deleted = self._connection_router.delete_all() if verbose: - orphan_count = len(orphaned_file_ids - folder_file_ids) - print( - f"Deleted {len(all_file_ids)} files/folders in {end - start:.2f}s", - end="", - ) - if orphan_count > 0: - print(f" (including {orphan_count} orphaned)") - else: - print() + print(f"Deleted {deleted} files/folders in {time.time() - start:.2f}s") # Broadcast delete events after file deletion but before cache reset if broadcast_delete_events and self.has_do_role and peer_emails and file_hashes: @@ -1310,15 +1325,15 @@ def _copy(self): from copy import deepcopy new_config = deepcopy(self.config) - new_manager = SyftboxManager.from_config(new_config) - if not isinstance(self._connection_router.connections[0], GDriveConnection): + connection = gdrive_connection_of(self._connection_router.transport) + if not isinstance(connection, GDriveConnection): raise ValueError("Only GDriveConnections can be copied") - if isinstance( - self._connection_router.connections[0].drive_service, - mock_drive_service.MockDriveService, - ): - # Create new connection pointing to the same backing store - drive_service = self._connection_router.connections[0].drive_service - new_do_connection = GDriveConnection.from_service(self.email, drive_service) - new_manager._add_connection(new_do_connection) - return new_manager + if isinstance(connection.drive_service, mock_drive_service.MockDriveService): + # A new connection to the same backing store + new_connection = GDriveConnection.from_service( + self.email, connection.drive_service + ) + return SyftboxManager.from_config( + new_config, SyftboxManager._drive_transport(new_connection) + ) + return SyftboxManager.from_config(new_config) diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 3d558ae6e06..99771f32fd5 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -641,8 +641,8 @@ 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( + collections = self.connection_router.transport.collections.clone_for_thread() + return collections.list_files( spec.wire_prefix(collection.variant), tag=collection.tag, content_hash=collection.content_hash, @@ -651,8 +651,8 @@ def _get_file_metadatas_with_new_connection( 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) + collections = self.connection_router.transport.collections.clone_for_thread() + return 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. @@ -945,8 +945,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( @@ -972,16 +971,7 @@ 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: - 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) + metas = self.connection_router.owner_get_outbox_file_metadatas(recipient_email) # Keep only well-formed events-message filenames; track timestamps so # we apply events in order in the merged message. @@ -1022,7 +1012,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.owner_delete_outbox_messages(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 359fa769250..9f217f98867 100644 --- a/syft/sync/sync/datasite_watcher_syncer.py +++ b/syft/sync/sync/datasite_watcher_syncer.py @@ -187,8 +187,8 @@ 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) + collections = self.connection_router.transport.collections.clone_for_thread() + data = collections.fetch_file(file_id) data = self.connection_router.peer_store.decrypt_dataset_if_needed( owner_email, data ) diff --git a/syft/sync/transport/config.py b/syft/sync/transport/config.py new file mode 100644 index 00000000000..01c943062db --- /dev/null +++ b/syft/sync/transport/config.py @@ -0,0 +1,31 @@ +"""The connection each transport component uses, as plain config.""" + +from typing import Self + +from pydantic import BaseModel + +from syft.sync.connections.base_connection import ConnectionConfig +from syft.sync.transport.protocols import TRANSPORT_COMPONENTS + + +class TransportConfig(BaseModel): + """The connection each transport component uses. + ``SyftboxManager.from_config`` builds the transport from it; only one + Google Drive connection for every component is built so far.""" + + p2p: ConnectionConfig + owner_store: ConnectionConfig + watcher_store: ConnectionConfig + collections: ConnectionConfig + peering: ConnectionConfig + + @classmethod + def one_connection(cls, connection_config: ConnectionConfig) -> Self: + """Every component over the same connection.""" + return cls(**{name: connection_config for name in TRANSPORT_COMPONENTS}) + + def connection_configs(self) -> list[ConnectionConfig]: + return [getattr(self, name) for name in TRANSPORT_COMPONENTS] + + def uses_one_connection(self) -> bool: + return all(config == self.p2p for config in self.connection_configs()) diff --git a/syft/sync/transport/protocols.py b/syft/sync/transport/protocols.py index f9b4744e0cf..fa67ae4602d 100644 --- a/syft/sync/transport/protocols.py +++ b/syft/sync/transport/protocols.py @@ -14,7 +14,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self from pydantic import BaseModel, ConfigDict @@ -26,7 +26,19 @@ from syft.sync.version.peer_info import PeerInfo -class P2PConnection(ABC): +class TransportComponent: + """What every component of a transport does besides its own work.""" + + def clone_for_thread(self) -> Self: + """A copy that another thread may use. This one, unless its + connection cannot be shared across threads.""" + return self + + def reset_caches(self) -> None: + """Forget what was cached from the connection.""" + + +class P2PConnection(TransportComponent, ABC): """Messages between peers, kept until the recipient has processed them. Sending and receiving each have their own confirmation: @@ -79,7 +91,7 @@ def mark_all_processed(self, messages: list[RawMessage]) -> None: self.mark_processed(message) -class OwnerStore(ABC): +class OwnerStore(TransportComponent, ABC): """The data owner's own storage: event log, checkpoints, rolling state. Only the owner writes and reads it, and it is never shared. @@ -153,8 +165,14 @@ def get_rolling_state(self) -> bytes | None: def delete_rolling_state(self) -> None: """Remove the rolling state.""" + # -- the whole account ----------------------------------------------------- + @abstractmethod + def delete_all(self) -> int: + """Delete everything this account stores, for every role, and return + how many files and folders went.""" + -class WatcherStore(ABC): +class WatcherStore(TransportComponent, ABC): """The data scientist's memory of what each owner sent it. One area per owner, in the scientist's own storage, never shared, so it @@ -185,7 +203,7 @@ def reset(self, peer: str) -> None: everything again.""" -class Collections(ABC): +class Collections(TransportComponent, ABC): """Sets of files a data owner shares with named readers (datasets).""" @abstractmethod @@ -221,6 +239,10 @@ def list_mine(self, prefix: str) -> list[str]: def delete(self, prefix: str, tag: str) -> None: """Remove every layout of ``tag`` published under ``prefix``.""" + @abstractmethod + def delete_by_handle(self, handle: str) -> None: + """Remove the one collection folder ``handle`` names.""" + @abstractmethod def list_shared_with_me(self, prefix: str) -> list[dict]: """Collections peers shared with me, each with owner_email, tag, @@ -244,7 +266,7 @@ def download( return {f["file_name"]: self.fetch_file(f["file_id"]) for f in files} -class Peering(ABC): +class Peering(TransportComponent, ABC): """Finding and trusting a peer: identity, channels, peer list, published documents.""" @@ -319,6 +341,16 @@ def grant_peer_info_read(self, peer: str, peer_info: PeerInfo) -> None: """ +# The components of a transport, as named on ``Transport``. +TRANSPORT_COMPONENTS = ( + "p2p", + "owner_store", + "watcher_store", + "collections", + "peering", +) + + class Transport(BaseModel): """One implementation of each job. The parts may come from different backends, though a peer's ``p2p`` and ``peering`` go together.""" @@ -332,3 +364,7 @@ class Transport(BaseModel): peering: Peering # Recorded on each Peer so a peer list says which backend reaches it. platform: BasePlatform + + def components(self) -> list[TransportComponent]: + """The components, in ``TRANSPORT_COMPONENTS`` order.""" + return [getattr(self, name) for name in TRANSPORT_COMPONENTS] diff --git a/syft/sync/version/peer_manager.py b/syft/sync/version/peer_manager.py index 225c3d09f8b..4ff7c370007 100644 --- a/syft/sync/version/peer_manager.py +++ b/syft/sync/version/peer_manager.py @@ -279,12 +279,10 @@ def _load_single_peer_info(self, peer_email: str) -> tuple[str, Optional[PeerInf Raises: TransportUnavailableError: when the backend cannot be reached. """ - # Create new connection for thread safety - connection = self.connection_router.connection_for_peer_file_read( - create_new=True - ) + # A copy per thread: the connection may not be thread-safe + peering = self.connection_router.transport.peering.clone_for_thread() try: - return (peer_email, connection.read_peer_info_file(peer_email)) + return (peer_email, peering.read_peer_info(peer_email)) except TransportUnavailableError: raise except Exception as e: @@ -652,10 +650,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.peering.clone_for_thread() + 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..95a8002caa9 100644 --- a/tests/integration/with_unit_coverage/test_sync_manager.py +++ b/tests/integration/with_unit_coverage/test_sync_manager.py @@ -30,6 +30,7 @@ create_test_pair, is_mock_mode, ) +from syft.sync.connections.drive.drive_connections import gdrive_connection_of def path_for_job(do_email: str, ds_email: str, filename: str = "test.job") -> str: @@ -552,7 +553,7 @@ def test_version_upgrade_breaks_communication(): ) # Phase 2: Simulate DS "upgrading" to a new incompatible version - ds_connection = ds_manager._connection_router.connections[0] + ds_connection = gdrive_connection_of(ds_manager._connection_router.transport) current = VersionInfo.current() new_version = VersionInfo( diff --git a/tests/migrations/p2p/test_message_protocol_downgrade.py b/tests/migrations/p2p/test_message_protocol_downgrade.py index 97febf9f831..7f3766544ee 100644 --- a/tests/migrations/p2p/test_message_protocol_downgrade.py +++ b/tests/migrations/p2p/test_message_protocol_downgrade.py @@ -29,6 +29,7 @@ from syft_migration import MigrationError, ProtocolSchema from tests.unit.utils import get_mock_events_messages, mock_message +from syft.sync.connections.drive.drive_connections import gdrive_connection_of def _client_schema( @@ -128,9 +129,9 @@ def _raw_proposal_version(do_manager, ds_email: str) -> str: def _raw_outbox_versions(ds_manager, do_email: str) -> list[str]: """The version fields of the DO's outbox blobs for us, unparsed.""" - raw_list = ds_manager._connection_router.connections[ - 0 - ].watcher_download_raw_events_from_outbox(do_email, None) + raw_list = gdrive_connection_of( + ds_manager._connection_router.transport + ).watcher_download_raw_events_from_outbox(do_email, None) return [_wire_version(ds_manager, do_email, raw) for raw in raw_list] diff --git a/tests/unit/test_cache_aware_sync.py b/tests/unit/test_cache_aware_sync.py index 663cecfe993..7cf609cc58b 100644 --- a/tests/unit/test_cache_aware_sync.py +++ b/tests/unit/test_cache_aware_sync.py @@ -10,7 +10,7 @@ from syft.sync.sync.collection_spec import CollectionSyncSpec from tests.unit.utils import ( - router_without_connections, + router_with_mock_connection, TEST_COLLECTION_PREFIX, TEST_COLLECTION_SUBPATH, get_mock_events_messages, @@ -247,7 +247,7 @@ def test_ds_cache_handles_deletions_correctly(): ], ) new_cache = DataSiteWatcherCache.from_config( - config, router_without_connections(do_manager.email) + config, router_with_mock_connection(do_manager.email) ) # Verify deleted file is not in reloaded cache diff --git a/tests/unit/test_delete_syftbox.py b/tests/unit/test_delete_syftbox.py index 35aefb338ed..cb1fd97436b 100644 --- a/tests/unit/test_delete_syftbox.py +++ b/tests/unit/test_delete_syftbox.py @@ -6,6 +6,7 @@ from syft.sync.login_utils import handle_potential_version_mismatches_on_login from syft.sync.syftbox_manager import SyftboxManager from syft.sync.version.version_info import VersionInfo +from syft.sync.connections.drive.drive_connections import gdrive_connection_of EMAIL = "test@example.com" @@ -200,7 +201,7 @@ def test_syftbox_folder_not_found(self): def _connection(manager: SyftboxManager): - return manager._connection_router.connections[0] + return gdrive_connection_of(manager._connection_router.transport) def _archive_one_proposal(): diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 31614b924f5..53f095cf690 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -16,7 +16,13 @@ import syft.sync.connections.connection_router as connection_router_module from syft.sync.sync import datasite_owner_syncer from syft.sync.connections.connection_router import ConnectionRouter -from syft.sync.connections.drive.drive_connections import DriveP2PConnection +from syft.sync.connections.drive.drive_connections import ( + DriveCollections, + DriveOwnerStore, + DriveP2PConnection, + DrivePeering, + DriveWatcherStore, +) from syft.sync.connections.drive.gdrive_transport import ( CollectionFolder, GDriveConnection, @@ -26,7 +32,6 @@ from syft.sync.encryption_policy import EncryptionPolicy from syft.sync.events.file_change_event import FileChangeEventsMessage from syft.sync.messages.proposed_filechange import ProposedFileChangesMessage -from syft.sync.peers.peer_store import PeerStore from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform from syft.sync.syftbox_manager import SyftboxManager from syft.sync.transport import ( @@ -66,7 +71,7 @@ def _events_name(ts: float) -> str: def pair(): """DO and DS connections over one mock Drive, with the DS's channels open.""" do_conn, ds_conn = pair_with_mock_service(DO, DS) - ds_conn.as_transport().peering.open_channels(DO) + DrivePeering(ds_conn).open_channels(DO) return do_conn, ds_conn @@ -75,7 +80,7 @@ def pair(): def test_drive_transport_does_every_job(): do_conn, _ = pair_with_mock_service(DO, DS) - transport = do_conn.as_transport() + transport = SyftboxManager._drive_transport(do_conn) assert isinstance(transport, Transport) assert isinstance(transport.p2p, P2PConnection) assert isinstance(transport.owner_store, OwnerStore) @@ -90,11 +95,11 @@ def test_drive_transport_does_every_job(): def test_a_proposed_change_sent_through_p2p_reaches_the_owner_inbox(pair): do_conn, ds_conn = pair - status = ds_conn.as_transport().p2p.send( + status = DriveP2PConnection(ds_conn).send( DO, _proposal_name(1.0), b"job", MessageKind.PROPOSED_CHANGES ) assert status is SendStatus.DELIVERED - [message] = do_conn.as_transport().p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + [message] = DriveP2PConnection(do_conn).receive(DS, MessageKind.PROPOSED_CHANGES) assert (message.sender, message.kind, message.data) == ( DS, MessageKind.PROPOSED_CHANGES, @@ -106,10 +111,10 @@ def test_inbox_messages_are_received_oldest_first_up_to_max(pair): do_conn, ds_conn = pair names = {submitted: _proposal_name(submitted) for submitted in (2.0, 1.0, 3.0)} for submitted, name in names.items(): - ds_conn.as_transport().p2p.send( + DriveP2PConnection(ds_conn).send( DO, name, f"job {submitted}".encode(), MessageKind.PROPOSED_CHANGES ) - p2p = do_conn.as_transport().p2p + p2p = DriveP2PConnection(do_conn) received = p2p.receive(DS, MessageKind.PROPOSED_CHANGES, max_messages=2) assert [(m.msg_id, m.data) for m in received] == [ @@ -156,7 +161,7 @@ def counted_get_media(self, *args, **kwargs): def _send_proposals(ds_conn, count: int) -> list[str]: names = [_proposal_name(float(i)) for i in range(count)] for name in names: - ds_conn.as_transport().p2p.send( + DriveP2PConnection(ds_conn).send( DO, name, name.encode(), MessageKind.PROPOSED_CHANGES ) return names @@ -164,7 +169,7 @@ def _send_proposals(ds_conn, count: int) -> list[str]: def test_receiving_inbox_lists_once_and_downloads_each_message(pair, monkeypatch): do_conn, ds_conn = pair - p2p = do_conn.as_transport().p2p + p2p = DriveP2PConnection(do_conn) p2p.receive(DS, MessageKind.PROPOSED_CHANGES) # finds the inbox folder once names = _send_proposals(ds_conn, 3) @@ -192,14 +197,14 @@ def slow_download(self, file_id: str) -> bytes: return original_download(self, file_id) monkeypatch.setattr(GDriveConnection, "download_file", slow_download) - do_conn.as_transport().p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + DriveP2PConnection(do_conn).receive(DS, MessageKind.PROPOSED_CHANGES) assert most_running[0] > 1 def test_marking_inbox_message_processed_is_one_drive_call(pair, monkeypatch): do_conn, ds_conn = pair _send_proposals(ds_conn, 2) - p2p = do_conn.as_transport().p2p + p2p = DriveP2PConnection(do_conn) first, second = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) p2p.mark_processed(first) # finds or creates the archive folder once @@ -213,7 +218,7 @@ def test_marking_inbox_message_processed_is_one_drive_call(pair, monkeypatch): def test_marking_inbox_messages_processed_is_one_batch_request(pair, monkeypatch): do_conn, ds_conn = pair _send_proposals(ds_conn, 4) - p2p = do_conn.as_transport().p2p + p2p = DriveP2PConnection(do_conn) first, *rest = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) p2p.mark_processed(first) # finds or creates the archive folder once @@ -227,7 +232,7 @@ def test_marking_inbox_messages_processed_is_one_batch_request(pair, monkeypatch def test_move_that_fails_in_batch_warns_and_rest_are_archived(pair): do_conn, ds_conn = pair _send_proposals(ds_conn, 2) - p2p = do_conn.as_transport().p2p + p2p = DriveP2PConnection(do_conn) first, second = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) gone = first.model_copy(update={"handle": "no-such-file"}) @@ -258,7 +263,7 @@ def download_all_but_broken(self, file_id: str) -> bytes: monkeypatch.setattr(GDriveConnection, "download_file", download_all_but_broken) with pytest.warns(UserWarning, match=f"Skipped inbox message {names[1]}"): - received = do_conn.as_transport().p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + received = DriveP2PConnection(do_conn).receive(DS, MessageKind.PROPOSED_CHANGES) assert [m.msg_id for m in received] == [names[0], names[2]] @@ -268,15 +273,15 @@ def test_receiving_inbox_while_drive_is_down_raises(pair, monkeypatch): _send_proposals(ds_conn, 2) monkeypatch.setattr(GDriveConnection, "download_file", _drive_down) with pytest.raises(TransportUnavailableError): - do_conn.as_transport().p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + DriveP2PConnection(do_conn).receive(DS, MessageKind.PROPOSED_CHANGES) def test_processed_inbox_message_is_not_received_again(pair): do_conn, ds_conn = pair - ds_conn.as_transport().p2p.send( + DriveP2PConnection(ds_conn).send( DO, _proposal_name(1.0), b"job", MessageKind.PROPOSED_CHANGES ) - p2p = do_conn.as_transport().p2p + p2p = DriveP2PConnection(do_conn) [message] = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) p2p.mark_processed(message) @@ -286,11 +291,11 @@ def test_processed_inbox_message_is_not_received_again(pair): def test_file_events_sent_through_p2p_reach_the_scientist_outbox(pair): do_conn, ds_conn = pair name = _events_name(1.0) - status = do_conn.as_transport().p2p.send( + status = DriveP2PConnection(do_conn).send( DS, name, b"results", MessageKind.FILE_EVENTS ) assert status is SendStatus.DELIVERED - p2p = ds_conn.as_transport().p2p + p2p = DriveP2PConnection(ds_conn) 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" @@ -299,10 +304,10 @@ def test_file_events_sent_through_p2p_reach_the_scientist_outbox(pair): def test_a_p2p_copy_for_another_thread_reads_the_same_outbox(pair): do_conn, ds_conn = pair - do_conn.as_transport().p2p.send( + DriveP2PConnection(do_conn).send( DS, _events_name(1.0), b"results", MessageKind.FILE_EVENTS ) - p2p = ds_conn.as_transport().p2p + p2p = DriveP2PConnection(ds_conn) [listed] = p2p.outbox_message_metadatas(DO, since_timestamp=None) copy = p2p.clone_for_thread() @@ -313,17 +318,17 @@ def test_a_p2p_copy_for_another_thread_reads_the_same_outbox(pair): def test_an_unknown_message_kind_is_refused(pair): do_conn, _ = pair with pytest.raises(ValueError): - do_conn.as_transport().p2p.send(DS, _events_name(1.0), b"x", "unknown-kind") + DriveP2PConnection(do_conn).send(DS, _events_name(1.0), b"x", "unknown-kind") def test_drive_never_has_sends_waiting_for_confirmation(pair): do_conn, _ = pair - assert do_conn.as_transport().p2p.delivered() == [] + assert DriveP2PConnection(do_conn).delivered() == [] def test_receiving_file_events_through_p2p_is_not_built_yet(pair): _, ds_conn = pair - p2p = ds_conn.as_transport().p2p + p2p = DriveP2PConnection(ds_conn) message = RawMessage(sender=DO, msg_id="m", kind=MessageKind.FILE_EVENTS, data=b"") with pytest.raises(NotImplementedError): p2p.receive(DO, MessageKind.FILE_EVENTS) @@ -336,7 +341,7 @@ def test_receiving_file_events_through_p2p_is_not_built_yet(pair): def test_owner_event_log_is_read_oldest_first(pair): do_conn, _ = pair - store = do_conn.as_transport().owner_store + store = DriveOwnerStore(do_conn) # 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") @@ -352,7 +357,7 @@ def test_owner_event_log_is_read_oldest_first(pair): def test_owner_event_count_matches_the_listed_handles(pair): do_conn, _ = pair - store = do_conn.as_transport().owner_store + store = DriveOwnerStore(do_conn) store.append_event(_events_name(1.0), b"event") store.append_event("syfteventsmessagev3_not-a-valid-name.tar.gz", b"junk") @@ -361,7 +366,7 @@ def test_owner_event_count_matches_the_listed_handles(pair): def test_owner_store_keeps_the_newest_full_checkpoint(pair): do_conn, _ = pair - store = do_conn.as_transport().owner_store + store = DriveOwnerStore(do_conn) 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" @@ -369,7 +374,7 @@ def test_owner_store_keeps_the_newest_full_checkpoint(pair): def test_owner_store_counts_and_deletes_incremental_checkpoints(pair): do_conn, _ = pair - store = do_conn.as_transport().owner_store + store = DriveOwnerStore(do_conn) store.put_incremental("incremental_checkpoint_1_1.0.tar.gz", b"inc") assert store.incremental_count() == 1 assert store.next_incremental_sequence() == 2 @@ -380,7 +385,7 @@ def test_owner_store_counts_and_deletes_incremental_checkpoints(pair): def test_owner_store_overwrites_and_deletes_the_rolling_state(pair): do_conn, _ = pair - store = do_conn.as_transport().owner_store + store = DriveOwnerStore(do_conn) 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" @@ -390,7 +395,7 @@ def test_owner_store_overwrites_and_deletes_the_rolling_state(pair): def test_owner_store_files_carry_the_syft_tag(pair): do_conn, _ = pair - store = do_conn.as_transport().owner_store + store = DriveOwnerStore(do_conn) 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") @@ -411,7 +416,7 @@ def test_owner_store_files_carry_the_syft_tag(pair): def test_watcher_storage_on_drive_is_not_built_yet(pair): _, ds_conn = pair with pytest.raises(NotImplementedError): - ds_conn.as_transport().watcher_store.latest_checkpoint(DO) + DriveWatcherStore(ds_conn).latest_checkpoint(DO) # -- Collections ------------------------------------------------------------------ @@ -419,26 +424,26 @@ 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_conn, ds_conn = pair - owner, reader = do_conn.as_transport(), ds_conn.as_transport() + owner, reader = DriveCollections(do_conn), DriveCollections(ds_conn) files = {"mock.csv": b"a,b\n1,2\n"} content_hash = CollectionFolder.compute_hash(files) - owner.collections.create(DATASETS, "census", content_hash, DO) - owner.collections.upload(DATASETS, "census", content_hash, files) - owner.collections.grant_read(DATASETS, "census", content_hash, [DS]) + owner.create(DATASETS, "census", content_hash, DO) + owner.upload(DATASETS, "census", content_hash, files) + owner.grant_read(DATASETS, "census", content_hash, [DS]) - assert owner.collections.list_mine(DATASETS) == ["census"] - mine = owner.collections.list_mine_with_permissions(DATASETS) + assert owner.list_mine(DATASETS) == ["census"] + mine = owner.list_mine_with_permissions(DATASETS) assert [c.tag for c in mine] == ["census"] - shared = reader.collections.list_shared_with_me(DATASETS) + shared = reader.list_shared_with_me(DATASETS) assert [(c["owner_email"], c["tag"]) for c in shared] == [(DO, "census")] - assert reader.collections.download(DATASETS, "census", content_hash, DO) == files - [listed] = reader.collections.list_files(DATASETS, "census", content_hash, DO) - assert reader.collections.fetch_file(listed["file_id"]) == files["mock.csv"] + assert reader.download(DATASETS, "census", content_hash, DO) == files + [listed] = reader.list_files(DATASETS, "census", content_hash, DO) + assert reader.fetch_file(listed["file_id"]) == files["mock.csv"] - owner.collections.delete(DATASETS, "census") - assert owner.collections.list_mine(DATASETS) == [] + owner.delete(DATASETS, "census") + assert owner.list_mine(DATASETS) == [] # -- Peering ---------------------------------------------------------------------- @@ -446,21 +451,21 @@ def test_a_published_collection_is_listed_and_downloaded_by_its_reader(pair): def test_peering_finds_the_peer_that_opened_channels(pair): do_conn, _ = pair - peering = do_conn.as_transport().peering + peering = DrivePeering(do_conn) assert peering.whoami() == DO assert DS in peering.incoming_requests() def test_peering_keeps_the_peer_registry(pair): do_conn, _ = pair - peering = do_conn.as_transport().peering + peering = DrivePeering(do_conn) peering.update_peer_registry(DS, "accepted") assert peering.read_peer_registry(force_download=True)[DS]["state"] == "accepted" def test_peer_registry_update_while_drive_is_down_keeps_other_peers(pair, monkeypatch): do_conn, _ = pair - peering = do_conn.as_transport().peering + peering = DrivePeering(do_conn) peering.update_peer_registry(DS, "accepted") do_conn._peers_json_cache = None # a new process has no cache @@ -473,37 +478,37 @@ 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 + owner = DrivePeering(pair[0]) + assert owner.read_own_key_bundle(DS) is None assert pair[0]._find_encryption_bundles_folder_id() is None def test_peering_publishes_a_key_bundle_both_sides_can_read(pair): do_conn, ds_conn = pair - owner, scientist = do_conn.as_transport(), ds_conn.as_transport() + owner, scientist = DrivePeering(do_conn), DrivePeering(ds_conn) bundle = '{"bundle": "do-public-key"}' - owner.peering.publish_key_bundle(DS, bundle) - owner.peering.grant_key_bundle_read(DS) + owner.publish_key_bundle(DS, bundle) + owner.grant_key_bundle_read(DS) - assert scientist.peering.read_peer_key_bundle(DO) == bundle - assert owner.peering.read_own_key_bundle(DS) == bundle - assert owner.peering.read_own_key_bundle("someone-else@example.com") is None + assert scientist.read_peer_key_bundle(DO) == bundle + assert owner.read_own_key_bundle(DS) == bundle + assert owner.read_own_key_bundle("someone-else@example.com") is None def test_peering_publishes_peer_info_the_peer_can_read(pair): do_conn, ds_conn = pair - owner, scientist = do_conn.as_transport(), ds_conn.as_transport() + owner, scientist = DrivePeering(do_conn), DrivePeering(ds_conn) info = PeerInfo( version_info=VersionInfo.current(), encryption_policy=EncryptionPolicy.REQUIRED, ) - owner.peering.publish_peer_info(info) - owner.peering.grant_peer_info_read(DS, info) + owner.publish_peer_info(info) + owner.grant_peer_info_read(DS, info) - assert owner.peering.read_own_peer_info() == info - assert scientist.peering.read_peer_info(DO) == info + assert owner.read_own_peer_info() == info + assert scientist.read_peer_info(DO) == info # -- ConnectionRouter ------------------------------------------------------------- @@ -511,11 +516,9 @@ def test_peering_publishes_peer_info_the_peer_can_read(pair): def _router_with_mock_p2p(): do_conn, _ = pair_with_mock_service(DO, DS) - router = ConnectionRouter(connections=[do_conn], peer_store=_store_knowing(DO, DS)) - drive = do_conn.as_transport() p2p = Mock(spec=DriveP2PConnection) - router._transport = drive.model_copy(update={"p2p": p2p}) - router._transport_source = do_conn + transport = SyftboxManager._drive_transport(do_conn).model_copy(update={"p2p": p2p}) + router = ConnectionRouter(transport=transport, peer_store=_store_knowing(DO, DS)) return router, p2p @@ -594,16 +597,6 @@ def test_the_router_reads_the_outbox_through_p2p(): p2p.outbox_message_metadatas.assert_called_once_with(DO, 5.0) -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)) - first = router.transport - assert router.transport is first - - router.connections[0] = other_conn - assert router.transport is not first - - def test_the_router_holds_no_drive_specific_code(): source = inspect.getsource(connection_router_module).lower() assert "gdrive" not in source diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index 365450ac0df..33613382a58 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -12,6 +12,7 @@ from syft.sync.syftbox_manager import SyftboxManager from tests.unit.test_sync_manager import path_for_job from tests.unit.utils import grant_job_inbox_access +from syft.sync.connections.drive.drive_connections import gdrive_connection_of # ========================================================================= @@ -623,7 +624,7 @@ def test_tampered_inbox_message_is_archived_with_warning(): ds_manager._send_file_change(file_path, "will be tampered") # Tamper with raw bytes in the mock backing store - do_conn = do_manager._connection_router.connections[0] + do_conn = gdrive_connection_of(do_manager._connection_router.transport) inbox_id = do_conn._get_own_datasite_inbox_id(ds_manager.email) files = do_conn.get_file_metadatas_from_folder(inbox_id) assert len(files) > 0 diff --git a/tests/unit/test_login_email_check.py b/tests/unit/test_login_email_check.py index 2d828f1dc50..6eb060be798 100644 --- a/tests/unit/test_login_email_check.py +++ b/tests/unit/test_login_email_check.py @@ -2,6 +2,7 @@ from syft.sync.login import _verify_token_matches_email from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.connections.drive.drive_connections import gdrive_connection_of def test_get_authenticated_email_matches_constructor_email(): @@ -17,7 +18,7 @@ def test_verify_token_matches_email_passes_when_matching(): def _drive_service(manager): - return manager._connection_router.connections[0].drive_service + return gdrive_connection_of(manager._connection_router.transport).drive_service def test_verify_token_matches_email_raises_on_mismatch(): diff --git a/tests/unit/test_outbox_compaction.py b/tests/unit/test_outbox_compaction.py index cbc867d2663..ff43ddfbfe1 100644 --- a/tests/unit/test_outbox_compaction.py +++ b/tests/unit/test_outbox_compaction.py @@ -3,6 +3,7 @@ from syft.sync.events.file_change_event import FileChangeEventsMessage from syft.sync.syftbox_manager import SyftboxManager from tests.unit.utils import get_mock_event +from syft.sync.connections.drive.drive_connections import gdrive_connection_of def test_compact_outboxes_if_needed_collapses_to_single_message(): @@ -20,7 +21,7 @@ def test_compact_outboxes_if_needed_collapses_to_single_message(): ds_manager.email, msg ) - conn_do = do_manager._connection_router.connections[0] + conn_do = gdrive_connection_of(do_manager._connection_router.transport) outbox_id = conn_do._get_own_datasite_outbox_id(ds_manager.email) pre = conn_do.get_file_metadatas_from_folder(outbox_id, since_timestamp=None) assert len(pre) == n_messages diff --git a/tests/unit/test_peer_request_owner.py b/tests/unit/test_peer_request_owner.py index d93d74e5092..36df24d72e2 100644 --- a/tests/unit/test_peer_request_owner.py +++ b/tests/unit/test_peer_request_owner.py @@ -19,6 +19,7 @@ from syft.sync.syftbox_manager import SyftboxManager from syft.sync.version.peer_manager import PeerState from syft.version import SYFT_VERSION +from syft.sync.connections.drive.drive_connections import gdrive_connection_of DO_EMAIL = "do@test.com" DS_EMAIL = "ds@test.com" @@ -46,9 +47,9 @@ def share_request_folder( victim: SyftboxManager, owner_email: str | None, claimed_email: str ) -> None: """Put a request folder for `victim` in Drive and share it with `victim`.""" - backing_store = victim._connection_router.connections[ - 0 - ].drive_service._backing_store + backing_store = gdrive_connection_of( + victim._connection_router.transport + ).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 17178da8c18..179c41fa38c 100644 --- a/tests/unit/test_peers_json_version.py +++ b/tests/unit/test_peers_json_version.py @@ -20,6 +20,7 @@ GDriveConnection, ) from syft.sync.peers.peer import PeerState +from syft.sync.syftbox_manager import SyftboxManager PEER = "bob@example.com" @@ -33,7 +34,7 @@ def _conn(peers_data): def _router(conn): router = Mock() - router.transport = conn.as_transport() + router.transport = SyftboxManager._drive_transport(conn) from syft.sync.connections.connection_router import ConnectionRouter return ConnectionRouter.get_all_peers_from_json.__get__(router, ConnectionRouter) diff --git a/tests/unit/test_sender_binding.py b/tests/unit/test_sender_binding.py index c1a8a968b89..cd084e4dc75 100644 --- a/tests/unit/test_sender_binding.py +++ b/tests/unit/test_sender_binding.py @@ -12,6 +12,7 @@ ProposedFileChangesMessage, ) from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.connections.drive.drive_connections import gdrive_connection_of def send_spoofed_message( @@ -80,7 +81,9 @@ def test_spoofed_sender_does_not_name_an_archive_folder(): send_spoofed_message(ds_manager, do_email, claimed_email, "any.txt") sync_do(do_manager) - connection = do_manager.datasite_owner_syncer.connection_router.connections[0] + connection = gdrive_connection_of( + do_manager.datasite_owner_syncer.connection_router.transport + ) folder_names = [ f.name for f in connection.drive_service._backing_store.files.values() ] diff --git a/tests/unit/test_shared_router.py b/tests/unit/test_shared_router.py index bd454416773..288e32c1325 100644 --- a/tests/unit/test_shared_router.py +++ b/tests/unit/test_shared_router.py @@ -1,10 +1,18 @@ """A manager builds one connection router and every part shares it, so the -whole manager has one connection, one transport and one peer store.""" +whole manager has one transport, one connection and one peer store.""" -from syft.sync.connections.base_connection import SyftboxPlatformConnection +from unittest.mock import Mock + +import pytest + +from syft.sync.connections.base_connection import ConnectionConfig +from syft.sync.connections.drive.drive_connections import gdrive_connection_of +from syft.sync.connections.drive.gdrive_transport import GDriveConnection from syft.sync.connections.drive.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 +from syft.sync.transport.config import TransportConfig +from syft.sync.transport.protocols import OwnerStore def _config_with_both_roles() -> SyftboxManagerConfig: @@ -13,10 +21,22 @@ def _config_with_both_roles() -> SyftboxManagerConfig: ) +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 test_every_part_of_manager_shares_one_router(): - manager = SyftboxManager.from_config(_config_with_both_roles()) + config = _config_with_both_roles() + transport = SyftboxManager._drive_transport(_mock_connection(config.email)) + manager = SyftboxManager.from_config(config, transport) router = manager._connection_router + assert router.transport is transport assert manager.peer_manager.connection_router is router assert manager.datasite_owner_syncer.connection_router is router assert manager.datasite_watcher_syncer.connection_router is router @@ -26,29 +46,57 @@ def test_every_part_of_manager_shares_one_router(): assert manager._peer_store is router.peer_store -def test_manager_builds_one_connection_from_its_config(monkeypatch): - config = _config_with_both_roles() - config.connection_configs = [ - GdriveConnectionConfig(email=config.email, token_path=None) - ] +def test_manager_builds_one_drive_connection_from_its_config(monkeypatch): built = [] def build_mock_connection(connection_config): - connection, _ = pair_with_mock_service(config.email, "peer@example.com") - built.append(connection) - return connection + built.append(_mock_connection(connection_config.email)) + return built[-1] - monkeypatch.setattr( - SyftboxPlatformConnection, "from_config", staticmethod(build_mock_connection) + monkeypatch.setattr(GDriveConnection, "from_config", build_mock_connection) + config = _config_with_both_roles() + config.transport_config = TransportConfig.one_connection( + _drive_config(config.email) ) manager = SyftboxManager.from_config(config) - assert manager._connection_router.connections == built - assert len(built) == 1 + [connection] = built + for component in manager._connection_router.transport.components(): + assert component._conn is connection -def test_added_mock_connection_is_added_once_and_backs_the_transport(): +def test_every_transport_component_shares_one_mock_connection(): _, do_manager = SyftboxManager.pair_with_mock_drive_service_connection() - [connection] = do_manager._connection_router.connections + transport = do_manager._connection_router.transport + connection = gdrive_connection_of(transport) + + for component in transport.components(): + assert component._conn is connection + + +def test_gdrive_connection_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"): + gdrive_connection_of(transport) + + +def test_manager_without_transport_config_fails_at_init(): + with pytest.raises(ValueError, match="names no transport"): + SyftboxManager.from_config(_config_with_both_roles()) + + +@pytest.mark.parametrize( + "p2p_config", + [ConnectionConfig(), _drive_config("other@example.com")], + ids=["not_drive", "other_drive_login"], +) +def test_transport_other_than_one_drive_connection_is_not_built_yet(p2p_config): + config = _config_with_both_roles() + config.transport_config = TransportConfig.one_connection( + _drive_config(config.email) + ).model_copy(update={"p2p": p2p_config}) - assert do_manager._connection_router.transport.p2p._conn is connection + with pytest.raises(NotImplementedError, match="one Google Drive connection"): + SyftboxManager.from_config(config) diff --git a/tests/unit/test_sync_manager.py b/tests/unit/test_sync_manager.py index 4f84b621823..a9f763cdaed 100644 --- a/tests/unit/test_sync_manager.py +++ b/tests/unit/test_sync_manager.py @@ -24,6 +24,7 @@ get_mock_proposed_events_messages, grant_job_inbox_access, ) +from syft.sync.connections.drive.drive_connections import gdrive_connection_of def path_for_job(do_email: str, ds_email: str, filename: str = "test.job") -> str: @@ -148,7 +149,7 @@ def test_sync_existing_datasite_state_do(): ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( use_in_memory_cache=False ) - connection_ds = ds_manager._connection_router.connections[0] + connection_ds = gdrive_connection_of(ds_manager._connection_router.transport) events_messages = get_mock_events_messages(2) for message in events_messages: @@ -186,7 +187,7 @@ 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] + connection_ds = gdrive_connection_of(ds_manager._connection_router.transport) proposed_events_messages = get_mock_proposed_events_messages( 2, email=ds_manager.email @@ -487,12 +488,14 @@ 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, + gdrive_connection_of(ds_manager._connection_router.transport).drive_service, ) conn.add_peer(do_manager.email) _add_known_peer(do_manager, email) return ConnectionRouter( - connections=[conn], peer_store=_store_knowing(email, do_manager.email) + transport=SyftboxManager._drive_transport(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 59486163bda..7211011c447 100644 --- a/tests/unit/test_version_negotiation.py +++ b/tests/unit/test_version_negotiation.py @@ -15,6 +15,7 @@ from syft.sync.version.peer_info import PeerInfo from syft.sync.version.peer_manager import CompatAction from syft.sync.version.version_info import CompatibilityStatus, VersionInfo +from syft.sync.connections.drive.drive_connections import gdrive_connection_of def build_client_version(client_version: str) -> VersionInfo: @@ -239,9 +240,9 @@ def test_load_peer_versions_parallel(self): SyftboxManager.pair_with_mock_drive_service_connection() ) - backing_store = ds_manager._connection_router.connections[ - 0 - ].drive_service._backing_store + backing_store = gdrive_connection_of( + ds_manager._connection_router.transport + ).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..8b899d4e0f3 100644 --- a/tests/unit/test_versioned_folder_adopt.py +++ b/tests/unit/test_versioned_folder_adopt.py @@ -25,6 +25,7 @@ execute_with_retries, ) from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.connections.drive.drive_connections import gdrive_connection_of EMAIL = "alice@example.com" @@ -158,7 +159,7 @@ def _checkpoints_owner_with_version(checkpoint_version: int): sync_automatically=False, ) router = do_manager._connection_router - connection = router.connections[0] + connection = gdrive_connection_of(router.transport) router.upload_checkpoint( Checkpoint(email=do_manager.email, version=checkpoint_version) ) diff --git a/tests/unit/test_watcher_event_datasite_check.py b/tests/unit/test_watcher_event_datasite_check.py index 9a6ba6e596b..faa4b66389f 100644 --- a/tests/unit/test_watcher_event_datasite_check.py +++ b/tests/unit/test_watcher_event_datasite_check.py @@ -19,7 +19,7 @@ ) from syft.sync.syftbox_manager import SyftboxManager from syft.sync.utils.syftbox_utils import get_event_hash_from_content -from tests.unit.utils import router_without_connections +from tests.unit.utils import router_with_mock_connection VICTIM = "victim@example.com" PEER = "peer@example.com" @@ -43,7 +43,7 @@ def _event( def _cache() -> DataSiteWatcherCache: return DataSiteWatcherCache.from_config( DataSiteWatcherCacheConfig(email=VICTIM, use_in_memory_cache=True), - router_without_connections(VICTIM), + router_with_mock_connection(VICTIM), ) @@ -129,7 +129,7 @@ def test_dropped_event_is_not_restored_after_restart(tmp_path): ) forged = _event(VICTIM, "syft.pub.yaml") legitimate = _event(PEER, "data/file.txt") - router = router_without_connections(VICTIM) + router = router_with_mock_connection(VICTIM) DataSiteWatcherCache.from_config(config, router).apply_event_message( FileChangeEventsMessage(events=[forged, legitimate]), PEER ) diff --git a/tests/unit/utils.py b/tests/unit/utils.py index 01c01dedb27..1045405dfac 100644 --- a/tests/unit/utils.py +++ b/tests/unit/utils.py @@ -10,7 +10,9 @@ from syft.sync.messages.proposed_filechange import ProposedFileChange from syft.sync.utils.syftbox_utils import get_event_hash_from_content from syft.sync.connections.connection_router import ConnectionRouter +from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service from syft.sync.peers.peer_store import PeerStore +from syft.sync.syftbox_manager import SyftboxManager def get_mock_event(path: str = "email@email.com/test.job") -> FileChangeEvent: @@ -220,8 +222,12 @@ def grant_job_inbox_access(do_manager, ds_email: str) -> None: ).grant_write_access(ds_email) -def router_without_connections(email: str) -> ConnectionRouter: - """A connection router with no backend, for a part a test builds on its own. - In a manager every part shares the router ``SyftboxManager.from_config`` - builds.""" - return ConnectionRouter(connections=[], peer_store=PeerStore(email=email)) +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") + return ConnectionRouter( + transport=SyftboxManager._drive_transport(connection), + peer_store=PeerStore(email=email), + )