add: call rclone with popen; refactor: google drive connection to use popen

Author: shrey

datashuttle/configs/aws_regions.py

@@ -4,8 +4,6 @@
 # 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."""

datashuttle/datashuttle_class.py

@@ -18,6 +18,8 @@
 )
 
 if TYPE_CHECKING:
+    import subprocess
+
     from datashuttle.utils.custom_types import (
         DisplayMode,
         OverwriteExistingFiles,
@@ -853,26 +855,27 @@ def setup_google_drive_connection(self) -> None:
             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 = None
+        gdrive_client_secret = None
+
+        if browser_available and self.cfg["gdrive_client_id"]:
+            gdrive_client_secret = gdrive.get_client_secret()
+
+        elif 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
+        process = self._setup_rclone_gdrive_config(
+            gdrive_client_secret, service_account_filepath
         )
 
+        rclone.await_call_rclone_with_popen_raise_on_fail(process, log=True)
+
         rclone.check_successful_connection_and_raise_error_on_fail(self.cfg)
 
         utils.log_and_message("Google Drive Connection Successful.")
@@ -1525,14 +1528,12 @@ 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(
+    ) -> subprocess.Popen:
+        return 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(

datashuttle/tui/interface.py

@@ -4,14 +4,16 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 if TYPE_CHECKING:
+    import subprocess
+
     import paramiko
 
     from datashuttle.configs.config_class import Configs
     from datashuttle.utils.custom_types import InterfaceOutput, TopLevelFolder
 
 from datashuttle import DataShuttle
 from datashuttle.configs import load_configs
-from datashuttle.utils import aws, rclone, ssh
+from datashuttle.utils import aws, rclone, ssh, utils
 
 
 class Interface:
@@ -36,6 +38,9 @@ def __init__(self) -> None:
         self.name_templates: Dict = {}
         self.tui_settings: Dict = {}
 
+        self.google_drive_rclone_setup_process: subprocess.Popen | None = None
+        self.gdrive_setup_process_killed: bool = False
+
     def select_existing_project(self, project_name: str) -> InterfaceOutput:
         """Load an existing project into `self.project`.
 
@@ -510,16 +515,46 @@ def setup_google_drive_connection(
     ) -> InterfaceOutput:
         """Try to set up and validate connection to Google Drive."""
         try:
-            self.project._setup_rclone_gdrive_config(
-                gdrive_client_secret, service_account_filepath, log=False
+            process = self.project._setup_rclone_gdrive_config(
+                gdrive_client_secret, service_account_filepath
             )
-            rclone.check_successful_connection_and_raise_error_on_fail(
-                self.project.cfg
+            self.google_drive_rclone_setup_process = process
+            self.gdrive_setup_process_killed = False
+
+            self.await_successful_gdrive_connection_setup_raise_on_fail(
+                process
             )
+
             return True, None
         except BaseException as e:
             return False, str(e)
 
+    def terminate_google_drive_setup(self) -> None:
+        """Terminate rclone setup for google drive."""
+        assert self.google_drive_rclone_setup_process is not None
+
+        process = self.google_drive_rclone_setup_process
+
+        if process.poll() is not None:
+            self.gdrive_setup_process_killed = True
+            process.kill()
+
+    def await_successful_gdrive_connection_setup_raise_on_fail(
+        self, process: subprocess.Popen
+    ):
+        """Wait for rclone setup for google drive to finish and verify successful connection."""
+        stdout, stderr = process.communicate()
+
+        if not self.gdrive_setup_process_killed:
+            if process.returncode != 0:
+                utils.log_and_raise_error(
+                    stderr.decode("utf-8"), ConnectionError
+                )
+
+            rclone.check_successful_connection_and_raise_error_on_fail(
+                self.project.cfg
+            )
+
     # Setup AWS
     # ----------------------------------------------------------------------------------
 

datashuttle/tui/screens/setup_gdrive.py

@@ -89,6 +89,7 @@ def on_button_pressed(self, event: Button.Pressed) -> None:
         ):
             if self.setup_worker and self.setup_worker.is_running:
                 self.setup_worker.cancel()  # fix
+                self.interface.terminate_google_drive_setup()
             self.dismiss()
 
         elif event.button.id == "setup_gdrive_ok_button":
@@ -197,6 +198,9 @@ def setup_gdrive_connection_using_service_account(
         self, service_account_filepath: Optional[str] = None
     ) -> None:
         """Set up the Google Drive connection using service account and show success message."""
+        message = "Setting up connection."
+        self.update_message_box_message(message)
+
         asyncio.create_task(
             self.setup_gdrive_connection_and_update_ui(
                 service_account_filepath=service_account_filepath

datashuttle/utils/rclone.py

@@ -10,6 +10,7 @@
 
 import os
 import platform
+import shlex
 import subprocess
 import tempfile
 from subprocess import CompletedProcess
@@ -94,6 +95,35 @@ def call_rclone_through_script(command: str) -> CompletedProcess:
     return output
 
 
+def call_rclone_with_popen(command: str) -> subprocess.Popen:
+    """Call rclone using `subprocess.Popen` for control over process termination.
+
+    It is not possible to kill a process while running it using `subprocess.run`.
+    """
+    command = "rclone " + command
+    process = subprocess.Popen(
+        shlex.split(command), stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    )
+    return process
+
+
+def await_call_rclone_with_popen_raise_on_fail(
+    process: subprocess.Popen, log: bool = True
+):
+    """Await rclone the subprocess.Popen call.
+
+    Calling `process.communicate()` waits for the process to complete and returns
+    the stdout and stderr.
+    """
+    stdout, stderr = process.communicate()
+
+    if process.returncode != 0:
+        utils.log_and_raise_error(stderr.decode("utf-8"), ConnectionError)
+
+    if log:
+        log_rclone_config_output()
+
+
 # -----------------------------------------------------------------------------
 # Setup
 # -----------------------------------------------------------------------------
@@ -182,10 +212,15 @@ def setup_rclone_config_for_gdrive(
     rclone_config_name: str,
     gdrive_client_secret: str | None,
     service_account_filepath: Optional[str] = None,
-    log: bool = True,
-):
+) -> subprocess.Popen:
     """Set up rclone config for connections to Google Drive.
 
+    This function uses `call_rclone_with_popen` instead of `call_rclone`. This
+    is done to have more control over the setup process in case the user wishes to
+    cancel the setup. Since the rclone setup for google drive uses a local web server
+    for authentication to google drive, the running process must be killed before the
+    setup can be started again.
+
     Parameters
     ----------
     cfg
@@ -203,9 +238,6 @@ def setup_rclone_config_for_gdrive(
     service_account_filepath : path to service account file path for connection
         without browser
 
-    log
-        Whether to log, if `True` logger must already be initialised.
-
     """
     client_id_key_value = (
         f"client_id {cfg['gdrive_client_id']} "
@@ -224,7 +256,7 @@ def setup_rclone_config_for_gdrive(
         else f"service_account_file {service_account_filepath}"
     )
 
-    output = call_rclone(
+    process = call_rclone_with_popen(
         f"config create "
         f"{rclone_config_name} "
         f"drive "
@@ -233,16 +265,9 @@ def setup_rclone_config_for_gdrive(
         f"scope drive "
         f"root_folder_id {cfg['gdrive_root_folder_id']} "
         f"{service_account_filepath_arg}",
-        pipe_std=True,
     )
 
-    if output.returncode != 0:
-        utils.log_and_raise_error(
-            output.stderr.decode("utf-8"), ConnectionError
-        )
-
-    if log:
-        log_rclone_config_output()
+    return process
 
 
 def setup_rclone_config_for_aws(