[Prototyping] Play around with google drive and AWS testing with github secrets

.github/workflows/code_test_and_deploy.yml

@@ -55,8 +55,25 @@ jobs:
           conda uninstall datashuttle --force
           python -m pip install --upgrade pip
           pip install .[dev]
-      - name: Test
-        run: pytest
+#      - name: Test
+ #       run: pytest
+  #    - name: Set up Google Drive secrets
+   #     run: |
+    #      printf '%s' '${{ secrets.GDRIVE_SERVICE_ACCOUNT_JSON }}' > "$HOME/gdrive.json"
+     #     echo "GDRIVE_SERVICE_ACCOUNT_FILE=$HOME/gdrive.json" >> $GITHUB_ENV
+      #    echo "GDRIVE_ROOT_FOLDER_ID=${{ secrets.GDRIVE_ROOT_FOLDER_ID }}" >> $GITHUB_ENV
+
+ #     - name: Run Google Drive tests
+  #      run: pytest -q -k test_gdrive_connection
+
+      - name: Set up AWS secrets
+        run: |
+          echo "AWS_ACCESS_KEY_ID=${{ secrets.AWS_ACCESS_KEY_ID }}" >> $GITHUB_ENV
+          echo "AWS_ACCESS_KEY_ID_SECRET=${{ secrets.AWS_ACCESS_KEY_ID_SECRET }}" >> $GITHUB_ENV
+          echo "AWS_REGION=${{ secrets.AWS_REGION }}" >> $GITHUB_ENV
+
+      - name: Run AWS tests
+        run: pytest -q -k test_aws_connection
 
   build_sdist_wheels:
     name: Build source distribution

datashuttle/configs/aws_regions.py

@@ -0,0 +1,43 @@
+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]:
+    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 list(get_aws_regions().values())

datashuttle/configs/canonical_configs.py

@@ -39,9 +39,15 @@ def get_canonical_configs() -> dict:
     canonical_configs = {
         "local_path": Union[str, Path],
         "central_path": Optional[Union[str, Path]],
-        "connection_method": Optional[Literal["ssh", "local_filesystem"]],
+        "connection_method": Optional[
+            Literal["ssh", "local_filesystem", "gdrive", "aws"]
+        ],
         "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[str],
     }
 
     return canonical_configs
@@ -128,6 +134,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

@@ -116,16 +116,46 @@ 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
+        do a minimal backwards compatibility check and add config
+        keys to ensure backwards compatibility with new connection
+        methods added to Datashuttle.
+        But this will not automatically check the configs are valid, this
         requires calling self.check_dict_values_raise_on_fail()
         """
         with open(self.file_path, "r") 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
+    ):
+        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,
 )
 
@@ -892,6 +895,90 @@ 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:
+        """
+        Setup 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 a config_token after running an
+        rclone command displayed to the user on a machine with access to a browser.
+
+        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 = utils.get_user_input(
+                "Please input the path to your credentials json"
+            )  # TODO: add more explanation
+        #         config_token = gdrive.prompt_and_get_config_token(
+        #            self.cfg,
+        #           gdrive_client_secret,
+        #          self.cfg.get_rclone_config_name("gdrive"),
+        #         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)
+
+        ds_logger.close_log_filehandler()
+
+    # -------------------------------------------------------------------------
+    # AWS S3
+    # -------------------------------------------------------------------------
+
+    @requires_aws_configs
+    @check_configs_set
+    def setup_aws_connection(self) -> None:
+        """
+        Setup 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
     # -------------------------------------------------------------------------
@@ -903,6 +990,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:
         """
         Initialise the configurations for datashuttle to use on the
@@ -967,6 +1058,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,
             },
         )
 
@@ -1470,6 +1565,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
     # -------------------------------------------------------------------------