Add Google Drive and AWS S3 as a Remote Storage option

datashuttle/configs/aws_regions.py

@@ -0,0 +1,45 @@
+from typing import Dict, List
+
+# -----------------------------------------------------------------------------
+# AWS regions
+# -----------------------------------------------------------------------------
+
+# These function are used for type checking and providing intellisense to the developer
+
+
+def get_aws_regions() -> Dict[str, str]:
+    """Return a dict of available AWS S3 bucket regions."""
+    aws_regions = {
+        "US_EAST_1": "us-east-1",
+        "US_EAST_2": "us-east-2",
+        "US_WEST_1": "us-west-1",
+        "US_WEST_2": "us-west-2",
+        "CA_CENTRAL_1": "ca-central-1",
+        "EU_WEST_1": "eu-west-1",
+        "EU_WEST_2": "eu-west-2",
+        "EU_WEST_3": "eu-west-3",
+        "EU_NORTH_1": "eu-north-1",
+        "EU_SOUTH_1": "eu-south-1",
+        "EU_CENTRAL_1": "eu-central-1",
+        "AP_SOUTHEAST_1": "ap-southeast-1",
+        "AP_SOUTHEAST_2": "ap-southeast-2",
+        "AP_NORTHEAST_1": "ap-northeast-1",
+        "AP_NORTHEAST_2": "ap-northeast-2",
+        "AP_NORTHEAST_3": "ap-northeast-3",
+        "AP_SOUTH_1": "ap-south-1",
+        "AP_EAST_1": "ap-east-1",
+        "SA_EAST_1": "sa-east-1",
+        "IL_CENTRAL_1": "il-central-1",
+        "ME_SOUTH_1": "me-south-1",
+        "AF_SOUTH_1": "af-south-1",
+        "CN_NORTH_1": "cn-north-1",
+        "CN_NORTHWEST_1": "cn-northwest-1",
+        "US_GOV_EAST_1": "us-gov-east-1",
+        "US_GOV_WEST_1": "us-gov-west-1",
+    }
+    return aws_regions
+
+
+def get_aws_regions_list() -> List[str]:
+    """Return AWS S3 bucket regions as a list."""
+    return list(get_aws_regions().values())

datashuttle/configs/canonical_configs.py

@@ -24,18 +24,28 @@
 
 import typeguard
 
+from datashuttle.configs.aws_regions import get_aws_regions_list
 from datashuttle.utils import folders, utils
 from datashuttle.utils.custom_exceptions import ConfigError
 
 
+def get_connection_methods_list() -> List[str]:
+    """Return the canonical connection methods."""
+    return ["ssh", "local_filesystem", "gdrive", "aws"]
+
+
 def get_canonical_configs() -> dict:
     """Return the only permitted types for DataShuttle config values."""
     canonical_configs = {
         "local_path": Union[str, Path],
         "central_path": Optional[Union[str, Path]],
-        "connection_method": Optional[Literal["ssh", "local_filesystem"]],
+        "connection_method": Optional[Literal[*get_connection_methods_list()]],
         "central_host_id": Optional[str],
         "central_host_username": Optional[str],
+        "gdrive_client_id": Optional[str],
+        "gdrive_root_folder_id": Optional[str],
+        "aws_access_key_id": Optional[str],
+        "aws_region": Optional[Literal[*get_aws_regions_list()]],
     }
 
     return canonical_configs
@@ -121,6 +131,29 @@ def check_dict_values_raise_on_fail(config_dict: Configs) -> None:
             ConfigError,
         )
 
+    # Check gdrive settings
+    elif config_dict["connection_method"] == "gdrive":
+        if not config_dict["gdrive_root_folder_id"]:
+            utils.log_and_raise_error(
+                "'gdrive_root_folder_id' is required if 'connection_method' "
+                "is 'gdrive'.",
+                ConfigError,
+            )
+
+        if not config_dict["gdrive_client_id"]:
+            utils.log_and_message(
+                "`gdrive_client_id` not found in config. default rlcone client will be used (slower)."
+            )
+
+    # Check AWS settings
+    elif config_dict["connection_method"] == "aws" and (
+        not config_dict["aws_access_key_id"] or not config_dict["aws_region"]
+    ):
+        utils.log_and_raise_error(
+            "Both aws_access_key_id and aws_region must be present for AWS connection.",
+            ConfigError,
+        )
+
     # Initialise the local project folder
     utils.print_message_to_user(
         f"Making project folder at: {config_dict['local_path']}"

datashuttle/configs/config_class.py

@@ -120,16 +120,48 @@ def dump_to_file(self) -> None:
     def load_from_file(self) -> None:
         """Load a config dict saved at .yaml file.
 
-        Note this will not automatically check the configs are valid,
-        this requires calling self.check_dict_values_raise_on_fail().
+        This will do a minimal backwards compatibility check and
+        add config keys to ensure backwards compatibility with new connection
+        methods added to Datashuttle.
+
+        However, this will not automatically check the configs are valid, this
+        requires calling self.check_dict_values_raise_on_fail()
         """
         with open(self.file_path) as config_file:
             config_dict = yaml.full_load(config_file)
 
         load_configs.convert_str_and_pathlib_paths(config_dict, "str_to_path")
 
+        self.update_config_for_backward_compatability_if_required(config_dict)
+
         self.data = config_dict
 
+    def update_config_for_backward_compatability_if_required(
+        self, config_dict: Dict
+    ):
+        """Add keys introduced in later Datashuttle versions if they are missing."""
+        canonical_config_keys_to_add = [
+            "gdrive_client_id",
+            "gdrive_root_folder_id",
+            "aws_access_key_id",
+            "aws_region",
+        ]
+
+        # All keys shall be missing for a backwards compatibility update
+        if not (
+            all(
+                key in config_dict.keys()
+                for key in canonical_config_keys_to_add
+            )
+        ):
+            assert not any(
+                key in config_dict.keys()
+                for key in canonical_config_keys_to_add
+            )
+
+            for key in canonical_config_keys_to_add:
+                config_dict[key] = None
+
     # -------------------------------------------------------------------------
     # Utils
     # -------------------------------------------------------------------------

datashuttle/datashuttle_class.py

@@ -36,9 +36,11 @@
 from datashuttle.configs.config_class import Configs
 from datashuttle.datashuttle_functions import _format_top_level_folder
 from datashuttle.utils import (
+    aws,
     ds_logger,
     folders,
     formatting,
+    gdrive,
     getters,
     rclone,
     ssh,
@@ -53,6 +55,7 @@
 from datashuttle.utils.decorators import (  # noqa
     check_configs_set,
     check_is_not_local_project,
+    requires_aws_configs,
     requires_ssh_configs,
 )
 
@@ -826,6 +829,87 @@ def write_public_key(self, filepath: str) -> None:
             public.write(key.get_base64())
         public.close()
 
+    # -------------------------------------------------------------------------
+    # Google Drive
+    # -------------------------------------------------------------------------
+
+    @check_configs_set
+    def setup_google_drive_connection(self) -> None:
+        """Set up a connection to Google Drive using the provided credentials.
+
+        Assumes `gdrive_root_folder_id` is set in configs.
+
+        First, the user will be prompted to enter their Google Drive client
+        secret if `gdrive_client_id` is set in the configs.
+
+        Next, the user will be asked if their machine has access to a browser.
+        If not, they will be prompted to input their service account file path.
+
+        Next, with the provided credentials, the final setup will be done. This
+        opens up a browser if the user confirmed access to a browser.
+        """
+        self._start_log(
+            "setup-google-drive-connection-to-central-server",
+            local_vars=locals(),
+        )
+
+        if self.cfg["gdrive_client_id"]:
+            gdrive_client_secret = gdrive.get_client_secret()
+        else:
+            gdrive_client_secret = None
+
+        browser_available = gdrive.ask_user_for_browser(log=True)
+
+        if not browser_available:
+            service_account_filepath = (
+                gdrive.prompt_and_get_service_account_filepath(
+                    log=True,
+                )
+            )
+        else:
+            service_account_filepath = None
+
+        self._setup_rclone_gdrive_config(
+            gdrive_client_secret, service_account_filepath, log=True
+        )
+
+        rclone.check_successful_connection_and_raise_error_on_fail(self.cfg)
+
+        utils.log_and_message("Google Drive Connection Successful.")
+
+        ds_logger.close_log_filehandler()
+
+    # -------------------------------------------------------------------------
+    # AWS S3
+    # -------------------------------------------------------------------------
+
+    @requires_aws_configs
+    @check_configs_set
+    def setup_aws_connection(self) -> None:
+        """Set up a connection to AWS S3 buckets using the provided credentials.
+
+        Assumes `aws_access_key_id` and `aws_region` are set in configs.
+
+        First, the user will be prompted to input their AWS secret access key.
+
+        Next, with the provided credentials, the final connection setup will be done.
+        """
+        self._start_log(
+            "setup-aws-connection-to-central-server",
+            local_vars=locals(),
+        )
+
+        aws_secret_access_key = aws.get_aws_secret_access_key()
+
+        self._setup_rclone_aws_config(aws_secret_access_key, log=True)
+
+        rclone.check_successful_connection_and_raise_error_on_fail(self.cfg)
+        aws.raise_if_bucket_absent(self.cfg)
+
+        utils.log_and_message("AWS Connection Successful.")
+
+        ds_logger.close_log_filehandler()
+
     # -------------------------------------------------------------------------
     # Configs
     # -------------------------------------------------------------------------
@@ -837,6 +921,10 @@ def make_config_file(
         connection_method: str | None = None,
         central_host_id: Optional[str] = None,
         central_host_username: Optional[str] = None,
+        gdrive_client_id: Optional[str] = None,
+        gdrive_root_folder_id: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_region: Optional[str] = None,
     ) -> None:
         """Initialize the configurations for datashuttle on the local machine.
 
@@ -877,6 +965,26 @@ def make_config_file(
             username for which to log in to central host.
             e.g. ``"jziminski"``
 
+        gdrive_client_id
+            The client ID used to authenticate with the Google Drive API via OAuth 2.0.
+            This is obtained from the Google Cloud Console when setting up API credentials.
+            e.g. "1234567890-abc123def456.apps.googleusercontent.com"
+
+        gdrive_root_folder_id
+            The folder ID for the Google Drive folder to connect to. This can be copied
+            directly from your browser when on the folder in Google Drive.
+            e.g. 1eoAnopd2ZHOd87LgiPtgViFE7u3R9sSw
+
+        aws_access_key_id
+            The AWS access key ID used to authenticate requests to AWS services.
+            This is part of your AWS credentials and can be generated via the AWS IAM console.
+            e.g. "AKIAIOSFODNN7EXAMPLE"
+
+        aws_region
+            The AWS region in which your resources are located.
+            This determines the data center your requests are routed to.
+            e.g. "us-west-2"
+
         """
         self._start_log(
             "make-config-file",
@@ -900,6 +1008,10 @@ def make_config_file(
                 "connection_method": connection_method,
                 "central_host_id": central_host_id,
                 "central_host_username": central_host_username,
+                "gdrive_client_id": gdrive_client_id,
+                "gdrive_root_folder_id": gdrive_root_folder_id,
+                "aws_access_key_id": aws_access_key_id,
+                "aws_region": aws_region,
             },
         )
 
@@ -1405,6 +1517,30 @@ def _setup_rclone_central_local_filesystem_config(self) -> None:
             self.cfg.get_rclone_config_name("local_filesystem"),
         )
 
+    def _setup_rclone_gdrive_config(
+        self,
+        gdrive_client_secret: str | None,
+        service_account_filepath: str | None,
+        log: bool,
+    ) -> None:
+        rclone.setup_rclone_config_for_gdrive(
+            self.cfg,
+            self.cfg.get_rclone_config_name("gdrive"),
+            gdrive_client_secret,
+            service_account_filepath,
+            log=log,
+        )
+
+    def _setup_rclone_aws_config(
+        self, aws_secret_access_key: str, log: bool
+    ) -> None:
+        rclone.setup_rclone_config_for_aws(
+            self.cfg,
+            self.cfg.get_rclone_config_name("aws"),
+            aws_secret_access_key,
+            log=log,
+        )
+
     # Persistent settings
     # -------------------------------------------------------------------------