diff --git a/README.md b/README.md index 2e9f7f96..ca2090df 100644 --- a/README.md +++ b/README.md @@ -1021,6 +1021,7 @@ Options: --image TEXT The LEAN engine image to use (defaults to quantconnect/lean:latest) --update Pull the LEAN engine image before running the Downloader Data Provider --no-update Use the local LEAN engine image instead of pulling the latest version + --project INTEGER The cloud project ID to use for brokerage OAuth authentication --lean-config FILE The Lean configuration file that should be used (defaults to the nearest lean.json) --verbose Enable debug logging --help Show this message and exit. diff --git a/lean/commands/data/download.py b/lean/commands/data/download.py index 1657acc2..64c4b605 100644 --- a/lean/commands/data/download.py +++ b/lean/commands/data/download.py @@ -559,6 +559,9 @@ def _replace_data_type(ctx, param, value): is_flag=True, default=False, help="Use the local LEAN engine image instead of pulling the latest version") +@option("--project", + type=int, + help="The cloud project ID to use for brokerage OAuth authentication") @pass_context def download(ctx: Context, data_provider_historical: Optional[str], @@ -576,6 +579,7 @@ def download(ctx: Context, image: Optional[str], update: bool, no_update: bool, + project: Optional[int], **kwargs) -> None: """Purchase and download data directly from QuantConnect or download from supported data providers @@ -678,6 +682,13 @@ def download(ctx: Context, engine_image, container_module_version, project_config = container.manage_docker_image(image, update, no_update) + # OAuth downloaders need a real cloud project ID for the Auth0 URL, but `data download` + # runs outside any project. Take it from --project or prompt instead of the -1 the API rejects. + if data_downloader_provider.requires_auth(): + lean_config["project-id"] = data_downloader_provider.get_project_id( + project if project is not None else lean_config["project-id"], + require_project_id=True) + data_downloader_provider = config_build_for_name(lean_config, data_downloader_provider.get_name(), cli_data_downloaders, kwargs, logger, interactive=True) data_downloader_provider.ensure_module_installed(organization.id, container_module_version) diff --git a/lean/models/json_module.py b/lean/models/json_module.py index 19411dba..e1001406 100644 --- a/lean/models/json_module.py +++ b/lean/models/json_module.py @@ -203,6 +203,13 @@ def get_user_name(self, lean_config: Dict[str, Any], configuration, user_provide lean_config[user_name_key] = user_name return user_name + def requires_auth(self) -> bool: + """Returns whether this module uses OAuth (Auth0) authentication. + + :return: True if any of the module's configurations is an AuthConfiguration + """ + return any(isinstance(config, AuthConfiguration) for config in self._lean_configs) + def get_project_id(self, default_project_id: int, require_project_id: bool) -> int: """Retrieve the project ID, prompting the user if required and default is invalid. diff --git a/tests/commands/data/test_download.py b/tests/commands/data/test_download.py index 442cef9d..e16396b2 100644 --- a/tests/commands/data/test_download.py +++ b/tests/commands/data/test_download.py @@ -204,6 +204,65 @@ def test_download_data_non_interactive_wrong_data_type(wrong_data_type: str): assert wrong_data_type in error_msg +def _run_oauth_download(extra_run_command: List[str], cli_input: str = "\n"): + """Runs `lean data download` for an OAuth data provider, returning the project_id passed to Auth0.""" + for data_provider in cli_data_downloaders: + data_provider.__setattr__("_specifications_url", "") + + create_fake_lean_cli_directory() + # A data download runs outside any project, so the project id is -1 + lean_config_path = Path.cwd() / "lean.json" + lean_config_path.write_text(json.dumps({"data-folder": "data", "organization-id": "abc", "project-id": -1})) + + container = initialize_container() + + captured_project_ids = [] + + def _fake_get_authorization(auth0_client, brokerage_id, logger, project_id, *args, **kwargs): + captured_project_ids.append(project_id) + auth = MagicMock() + auth.get_authorization_config_without_account.return_value = {} + auth.get_account_ids.return_value = [] + return auth + + with mock.patch.object(container.lean_runner, "get_basic_docker_config_without_algo", + return_value={"commands": [], "mounts": []}), \ + mock.patch.object(container.api_client.data, "download_public_file_json", + return_value=_get_data_provider_config()), \ + mock.patch.object(container.api_client.organizations, "get", return_value=create_api_organization()), \ + mock.patch("lean.models.json_module.get_authorization", side_effect=_fake_get_authorization): + run_parameters = [ + "data", "download", + "--data-provider-historical", "Alpaca", + "--data-type", "Trade", + "--resolution", "Hour", + "--security-type", "Equity", + "--ticker", "AAPL", + "--start", "20240101", + "--end", "20240202", + "--market", "USA", + "--alpaca-environment", "paper", + ] + run_parameters += extra_run_command + result = CliRunner().invoke(lean, run_parameters, input=cli_input) + + return result, captured_project_ids + + +def test_download_oauth_provider_uses_provided_project_id(): + result, captured_project_ids = _run_oauth_download(["--project", "12345"]) + + assert result.exit_code == 0 + assert captured_project_ids == [12345] + + +def test_download_oauth_provider_prompts_for_project_id_when_missing(): + result, captured_project_ids = _run_oauth_download([], cli_input="54321\n") + + assert result.exit_code == 0 + assert captured_project_ids == [54321] + + def test_non_interactive_bulk_select(): # TODO pass diff --git a/tests/test_main.py b/tests/test_main.py index 9acef3f5..e877505d 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -20,14 +20,14 @@ def test_lean_shows_help_when_called_without_arguments() -> None: result = CliRunner().invoke(lean, []) assert result.exit_code == 0 - assert "Usage: lean [OPTIONS] COMMAND [ARGS]..." in result.output + assert "Usage: lean [OPTIONS]" in result.output def test_lean_shows_help_when_called_with_help_option() -> None: result = CliRunner().invoke(lean, ["--help"]) assert result.exit_code == 0 - assert "Usage: lean [OPTIONS] COMMAND [ARGS]..." in result.output + assert "Usage: lean [OPTIONS]" in result.output def test_lean_shows_error_when_running_unknown_command() -> None: