Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion datashuttle/configs/config_class.py
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,7 @@ def get_base_folder(
def get_rclone_config_name(
self, connection_method: Optional[str] = None
) -> str:
"""Generate the rclone configuration name for the project.
"""Generate the rclone configuration name for the central project.

These configs are created by datashuttle but managed and stored by rclone.
"""
Expand Down
33 changes: 18 additions & 15 deletions datashuttle/datashuttle_class.py
Original file line number Diff line number Diff line change
Expand Up @@ -845,7 +845,8 @@ def setup_google_drive_connection(self) -> None:
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.
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.
Expand All @@ -855,23 +856,25 @@ def setup_google_drive_connection(self) -> None:
local_vars=locals(),
)

browser_available = gdrive.ask_user_for_browser(log=True)

service_account_filepath = None
gdrive_client_secret = None

if browser_available and self.cfg["gdrive_client_id"]:
if self.cfg["gdrive_client_id"]:
gdrive_client_secret = gdrive.get_client_secret()
else:
gdrive_client_secret = None

elif not browser_available:
service_account_filepath = (
gdrive.prompt_and_get_service_account_filepath(
log=True,
)
browser_available = gdrive.ask_user_for_browser(log=True)

if not browser_available:
config_token = gdrive.prompt_and_get_config_token(
self.cfg,
gdrive_client_secret,
self.cfg.get_rclone_config_name("gdrive"),
log=True,
)
else:
config_token = None

process = self._setup_rclone_gdrive_config(
gdrive_client_secret, service_account_filepath
gdrive_client_secret, config_token
)

rclone.await_call_rclone_with_popen_raise_on_fail(process, log=True)
Expand Down Expand Up @@ -1527,13 +1530,13 @@ def _setup_rclone_central_local_filesystem_config(self) -> None:
def _setup_rclone_gdrive_config(
self,
gdrive_client_secret: str | None,
service_account_filepath: str | None,
config_token: str | None,
) -> subprocess.Popen:
return rclone.setup_rclone_config_for_gdrive(
self.cfg,
self.cfg.get_rclone_config_name("gdrive"),
gdrive_client_secret,
service_account_filepath,
config_token,
)

def _setup_rclone_aws_config(
Expand Down
21 changes: 18 additions & 3 deletions datashuttle/tui/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

from datashuttle import DataShuttle
from datashuttle.configs import load_configs
from datashuttle.utils import aws, rclone, ssh, utils
from datashuttle.utils import aws, gdrive, rclone, ssh, utils


class Interface:
Expand Down Expand Up @@ -511,7 +511,7 @@ def setup_key_pair_and_rclone_config(
def setup_google_drive_connection(
self,
gdrive_client_secret: Optional[str] = None,
service_account_filepath: Optional[str] = None,
config_token: Optional[str] = None,
) -> InterfaceOutput:
"""Try to set up and validate connection to Google Drive.

Expand All @@ -526,7 +526,7 @@ def setup_google_drive_connection(
"""
try:
process = self.project._setup_rclone_gdrive_config(
gdrive_client_secret, service_account_filepath
gdrive_client_secret, config_token
)
self.google_drive_rclone_setup_process = process
self.gdrive_setup_process_killed = False
Expand All @@ -539,6 +539,21 @@ def setup_google_drive_connection(
except BaseException as e:
return False, str(e)

def get_rclone_message_for_gdrive_without_browser(
self, gdrive_client_secret: Optional[str] = None
) -> InterfaceOutput:
"""Get the rclone message for Google Drive setup without a browser."""
try:
output = gdrive.preliminary_for_setup_without_browser(
self.project.cfg,
gdrive_client_secret,
self.project.cfg.get_rclone_config_name("gdrive"),
log=False,
)
return True, output
except BaseException as e:
return False, str(e)

def terminate_google_drive_setup(self) -> None:
"""Terminate rclone setup for google drive by killing the rclone process."""
assert self.google_drive_rclone_setup_process is not None
Expand Down
86 changes: 55 additions & 31 deletions datashuttle/tui/screens/setup_gdrive.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@ class SetupGdriveScreen(ModalScreen):

If the config contains a "gdrive_client_id", the user is prompted
to enter a client secret. If the user has access to a browser, a Google Drive
authentication page will open. Otherwise, the user is asked to enter path to a
service account file.
authentication page will open. Otherwise, the user is asked to run an rclone command
and input a config token.
"""

def __init__(self, interface: Interface) -> None:
Expand All @@ -37,6 +37,7 @@ def __init__(self, interface: Interface) -> None:
self.stage: int = 0
self.setup_worker: Worker | None = None
self.is_browser_available: bool = True
self.gdrive_client_secret: Optional[str] = None

# For handling credential inputs
self.input_box: Input = Input(
Expand Down Expand Up @@ -75,9 +76,9 @@ def on_button_pressed(self, event: Button.Pressed) -> None:
and proceeds to a browser authentication.

3) "no" button : A "no" answer to the availability of browser question. On click,
prompts the user to enter path to their service account file.
prompts the user to enter a config token by running an rclone command.

4) "enter" button : To enter the client secret or service account file path.
4) "enter" button : To enter the client secret or config token.

5) "finish" button : To finish the setup.

Expand All @@ -94,39 +95,44 @@ def on_button_pressed(self, event: Button.Pressed) -> None:
self.dismiss()

elif event.button.id == "setup_gdrive_ok_button":
self.ask_user_for_browser()

elif event.button.id == "setup_gdrive_yes_button":
self.remove_yes_no_buttons()
self.query_one("#setup_gdrive_ok_button").remove()

if self.interface.project.cfg["gdrive_client_id"]:
self.ask_user_for_gdrive_client_secret()
else:
self.open_browser_and_setup_gdrive_connection()
self.ask_user_for_browser()

elif event.button.id == "setup_gdrive_yes_button":
self.remove_yes_no_buttons()
self.open_browser_and_setup_gdrive_connection(
self.gdrive_client_secret
)

elif event.button.id == "setup_gdrive_no_button":
self.is_browser_available = False
self.remove_yes_no_buttons()
self.prompt_user_for_service_account_filepath()
self.prompt_user_for_config_token()

elif event.button.id == "setup_gdrive_enter_button":
if self.is_browser_available:
gdrive_client_secret = (
if (
self.interface.project.cfg["gdrive_client_id"]
and self.stage == 0
):
self.gdrive_client_secret = (
self.input_box.value.strip()
if self.input_box.value.strip()
else None
)
self.open_browser_and_setup_gdrive_connection(
gdrive_client_secret
)
self.stage += 1
self.ask_user_for_browser()
else:
service_account_filepath = (
config_token = (
self.input_box.value.strip()
if self.input_box.value.strip()
else None
)
self.setup_gdrive_connection_using_service_account(
service_account_filepath
self.setup_gdrive_connection_using_config_token(
self.gdrive_client_secret, config_token
)

def ask_user_for_browser(self) -> None:
Expand All @@ -137,7 +143,11 @@ def ask_user_for_browser(self) -> None:
)
self.update_message_box_message(message)

self.query_one("#setup_gdrive_ok_button").remove()
if self.enter_button.is_mounted:
self.enter_button.remove()

if self.input_box.is_mounted:
self.input_box.visible = False

# Mount the Yes and No buttons
yes_button = Button("Yes", id="setup_gdrive_yes_button")
Expand Down Expand Up @@ -184,34 +194,46 @@ def open_browser_and_setup_gdrive_connection(
)
)

def prompt_user_for_service_account_filepath(self) -> None:
"""Set up widgets and prompt user for their service account file path for browser-less connection."""
message = "Please enter your service account file path."
def prompt_user_for_config_token(self) -> None:
"""Prompt the user for the rclone config token for Google Drive setup."""
success, message = (
self.interface.get_rclone_message_for_gdrive_without_browser(
self.gdrive_client_secret
)
)

self.update_message_box_message(message)
if not success:
self.display_failed(message)
return

self.update_message_box_message(
message + "\nPress shift+click to copy."
)

self.enter_button = Button("Enter", id="setup_gdrive_enter_button")
self.query_one("#setup_gdrive_buttons_horizontal").mount(
self.enter_button, before="#setup_gdrive_cancel_button"
)
self.mount_input_box_before_buttons()

def setup_gdrive_connection_using_service_account(
self, service_account_filepath: Optional[str] = None
def setup_gdrive_connection_using_config_token(
self, gdrive_client_secret: str | None, config_token: str | None
) -> None:
"""Set up the Google Drive connection using service account and show success message."""
"""Set up the Google Drive connection using rclone config token."""
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
gdrive_client_secret=gdrive_client_secret,
config_token=config_token,
)
)

async def setup_gdrive_connection_and_update_ui(
self,
gdrive_client_secret: Optional[str] = None,
service_account_filepath: Optional[str] = None,
config_token: Optional[str] = None,
) -> None:
"""Start the Google Drive connection setup in a separate thread.

Expand All @@ -232,7 +254,7 @@ async def setup_gdrive_connection_and_update_ui(
self.enter_button.disabled = True

worker = self.setup_gdrive_connection(
gdrive_client_secret, service_account_filepath
gdrive_client_secret, config_token
)
self.setup_worker = worker
if worker.is_running:
Expand All @@ -250,7 +272,7 @@ async def setup_gdrive_connection_and_update_ui(
def setup_gdrive_connection(
self,
gdrive_client_secret: Optional[str] = None,
service_account_filepath: Optional[str] = None,
config_token: Optional[str] = None,
) -> Worker[InterfaceOutput]:
"""Authenticate the Google Drive connection.

Expand All @@ -260,7 +282,7 @@ def setup_gdrive_connection(
with Google Drive.
"""
success, output = self.interface.setup_google_drive_connection(
gdrive_client_secret, service_account_filepath
gdrive_client_secret, config_token
)
return success, output

Expand Down Expand Up @@ -302,6 +324,8 @@ def mount_input_box_before_buttons(
self.query_one("#gdrive_setup_messagebox_message_container").mount(
self.input_box, after="#gdrive_setup_messagebox_message"
)
self.input_box.visible = True
self.input_box.value = ""

def remove_yes_no_buttons(self) -> None:
"""Remove yes and no buttons."""
Expand Down
2 changes: 1 addition & 1 deletion datashuttle/tui/shared/configs_content.py
Original file line number Diff line number Diff line change
Expand Up @@ -530,7 +530,7 @@ def widget_configs_match_saved_configs(self):
for key, value in cfg_kwargs.items():
saved_val = self.interface.get_configs()[key]
if key in ["central_path", "local_path"]:
if value.name != project_name:
if value is not None and value.name != project_name:
value = value / project_name
if saved_val != value:
return False
Expand Down
Loading
Loading