Skip to content

Commit

Permalink
checks notebook presence before finding userdata
Browse files Browse the repository at this point in the history
  • Loading branch information
rudolfix committed Dec 3, 2024
1 parent 9c44290 commit 778c455
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 1 deletion.
6 changes: 6 additions & 0 deletions dlt/common/configuration/providers/toml.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
21 changes: 20 additions & 1 deletion tests/common/configuration/test_toml_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand All @@ -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()

0 comments on commit 778c455

Please sign in to comment.