From 778c455776196600a5bc9599e8721005a1ceab5b Mon Sep 17 00:00:00 2001 From: Marcin Rudolf Date: Tue, 3 Dec 2024 17:26:43 +0100 Subject: [PATCH] checks notebook presence before finding userdata --- dlt/common/configuration/providers/toml.py | 6 ++++++ .../configuration/test_toml_provider.py | 21 ++++++++++++++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/dlt/common/configuration/providers/toml.py b/dlt/common/configuration/providers/toml.py index 3636565fae..e586fef225 100644 --- a/dlt/common/configuration/providers/toml.py +++ b/dlt/common/configuration/providers/toml.py @@ -124,6 +124,12 @@ def _read_google_colab_secrets(self, name: str, file_name: str) -> tomlkit.TOMLD """Try to load the toml from google colab userdata object""" try: from google.colab import userdata + from dlt.common.runtime.exec_info import is_notebook + + # make sure we work in interactive mode (get_ipython() is available) + # when dlt cli is run, userdata is available but without a kernel + if not is_notebook(): + return None try: return tomlkit.loads(userdata.get(file_name)) diff --git a/tests/common/configuration/test_toml_provider.py b/tests/common/configuration/test_toml_provider.py index 481c21b7bb..9538849976 100644 --- a/tests/common/configuration/test_toml_provider.py +++ b/tests/common/configuration/test_toml_provider.py @@ -4,6 +4,7 @@ import yaml from typing import Any, Dict, Type import datetime # noqa: I251 +from unittest.mock import Mock import dlt from dlt.common import pendulum, json @@ -538,11 +539,28 @@ def loader() -> Dict[str, Any]: def test_colab_toml() -> None: + import builtins + # use a path without any settings files try: sys.path.append("tests/common/cases/modules") - # secrets are in user data + + # ipython not present provider: SettingsTomlProvider = SecretsTomlProvider("tests/common/null", global_dir=None) + assert provider.is_empty + + get_ipython_m = Mock() + get_ipython_m.return_value = "google.colab.Shell" + # make it available to all modules + builtins.get_ipython = get_ipython_m # type: ignore[attr-defined] + # test mock + assert get_ipython() == "google.colab.Shell" # type: ignore[name-defined] # noqa + from dlt.common.runtime.exec_info import is_notebook + + assert is_notebook() + + # secrets are in user data + provider = SecretsTomlProvider("tests/common/null", global_dir=None) assert provider.to_toml() == 'api_key="api"' # config is not in userdata provider = ConfigTomlProvider("tests/common/null", "unknown") @@ -551,4 +569,5 @@ def test_colab_toml() -> None: provider = SecretsTomlProvider("tests/common/cases/configuration/.dlt", global_dir=None) assert provider.get_value("secret_value", str, None) == ("2137", "secret_value") finally: + delattr(builtins, "get_ipython") sys.path.pop()