Cloud API: Improve environment handling for Jupyter Notebooks

.github/workflows/main.yml

@@ -65,7 +65,7 @@ jobs:
         pip install "setuptools>=64" --upgrade
 
         # Install package in editable mode.
-        pip install --use-pep517 --prefer-binary --editable=.[io,test,develop]
+        pip install --use-pep517 --prefer-binary --editable=.[all,examples,develop,test]
 
     - name: Run linter and software tests
       run: |

.gitignore

@@ -6,4 +6,5 @@ __pycache__
 *.pyc
 dist
 .coverage*
+.env
 coverage.xml

cratedb_toolkit/__init__.py

@@ -10,6 +10,10 @@
 except PackageNotFoundError:  # pragma: no cover
     __version__ = "unknown"
 
-from .api import ManagedCluster  # noqa: F401
-from .config import configure  # noqa: F401
-from .model import InputOutputResource, TableAddress  # noqa: F401
+from .config import preconfigure
+
+preconfigure()
+
+from .api import ManagedCluster  # noqa: E402, F401
+from .config import configure  # noqa: E402, F401
+from .model import InputOutputResource, TableAddress  # noqa: E402, F401

cratedb_toolkit/api/main.py

@@ -306,6 +306,13 @@ def get_client_bundle(self, username: str = None, password: str = None) -> Clien
             sqlalchemy=adapter.engine,
         )
 
+    def query(self, sql: str):
+        """
+        Shortcut method to submit a database query in SQL format, and retrieve the results.
+        """
+        client_bundle = self.get_client_bundle()
+        return client_bundle.adapter.run_sql(sql, records=True)
+
 
 @dataclasses.dataclass
 class StandaloneCluster(ClusterBase):

cratedb_toolkit/config.py

@@ -1,4 +1,7 @@
 import dataclasses
+import os
+import sys
+import types
 
 
 @dataclasses.dataclass
@@ -14,6 +17,14 @@ class EnvironmentConfiguration:
     settings_accept_env: bool = False
     settings_errors: str = "raise"
 
+    @property
+    def RUNNING_ON_JUPYTER(self):
+        return "JPY_SESSION_NAME" in os.environ
+
+    @property
+    def RUNNING_ON_PYTEST(self):
+        return "PYTEST_CURRENT_TEST" in os.environ
+
 
 # The global toolkit environment.
 # TODO: Provide as context manager.
@@ -37,3 +48,50 @@ def configure(
     CONFIG.settings_accept_cli = settings_accept_cli
     CONFIG.settings_accept_env = settings_accept_env
     CONFIG.settings_errors = settings_errors
+
+
+def preconfigure():
+    """
+    Configure a few environment details very early.
+    """
+    configure_croud()
+
+    # When running on Jupyter Notebooks, accept environment variables only.
+    if CONFIG.RUNNING_ON_JUPYTER or CONFIG.RUNNING_ON_PYTEST:
+        configure(
+            settings_accept_cli=False,
+            settings_accept_env=True,
+        )
+
+
+def configure_croud(no_spinner: bool = None, use_spinner: bool = None):
+    """
+    Turn off croud's Halo spinner when running in Jupyter Notebooks. It does not work well.
+
+    - https://github.com/ManrajGrover/halo/issues/32
+    - https://github.com/manrajgrover/halo/issues/179
+    """
+    if no_spinner or ((CONFIG.RUNNING_ON_JUPYTER or CONFIG.RUNNING_ON_PYTEST) and not use_spinner):
+        mod = types.ModuleType(
+            "croud.tools.spinner", "Mocking the croud.tools.spinner module, to turn off the Halo spinner"
+        )
+        setattr(mod, "HALO", NoopContextManager())  # noqa: B010
+        sys.modules["croud.tools.spinner"] = mod
+
+
+class NoopContextManager:
+    """
+    For making the Halo progressbar a no-op.
+    """
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        pass
+
+    def stop(self):
+        pass

cratedb_toolkit/util/runtime.py

@@ -11,9 +11,9 @@ def flexfun(domain: str = None):
     """
     Function decorator, which honors toolkit environment settings wrt. error handling.
 
-    It is sorting out whether to raise exceptions, or whether to just `exit(1)`.
+    It is sorting out whether to raise exceptions, or whether to just `exit({1,2})`.
 
-    This detail important to handle well depending on the runtime environment. It can
+    This detail is important to handle well depending on the runtime environment. It can
     either be a standalone program, used on behalf of a library, or within a Jupyter
     Notebook.
 
@@ -28,7 +28,7 @@ def flexfun(domain: str = None):
         Generic error code.
     2
 
-        Parse error—for instance, when parsing command-line options, the ‘.wgetrc’ or ‘.netrc’...
+        Parse error. For instance, when parsing command-line options or config files.
 
     -- https://www.pythontutorial.net/advanced-python/python-decorator-arguments/
     """

cratedb_toolkit/util/setting.py

@@ -4,6 +4,9 @@
 import typing as t
 
 import click
+from dotenv import find_dotenv, load_dotenv
+
+from cratedb_toolkit.config import CONFIG
 
 logger = logging.getLogger(__name__)
 
@@ -36,11 +39,22 @@ def obtain_settings(specs: t.List[Setting], prog_name: str = None) -> t.Dict[str
     - Environment variable. Example: `export FOOBAR=bazqux`.
     - Environment variable prefix. Example: `export APPNAME_FOOBAR=bazqux`.
     """
+    # Load environment variables from `.env` file.
+    if not CONFIG.RUNNING_ON_PYTEST:
+        dotenv_file = find_dotenv()
+        logger.info(f"Loading environment variables from .env file: {dotenv_file}")
+        load_dotenv(dotenv_file)
+
+    # Decode settings from command-line arguments or environment variables.
     prog_name = prog_name or sys.argv[0]
     click_specs = [spec.click for spec in specs]
     command = click.Command(prog_name, params=click_specs)
+    if CONFIG.RUNNING_ON_JUPYTER or CONFIG.RUNNING_ON_PYTEST:
+        args = []
+    else:
+        args = sys.argv[1:]
     try:
-        with command.make_context(prog_name, args=sys.argv[1:]) as ctx:
+        with command.make_context(prog_name, args=args) as ctx:
             return ctx.params
     except click.exceptions.Exit as ex:
         if ex.exit_code != 0:

examples/cloud_import.py

@@ -4,8 +4,9 @@
 About
 =====
 
-Example program demonstrating how to load data from files using
-the CrateDB Cloud Import API.
+This is an example program demonstrating how to load data from
+files using the Import API interface of CrateDB Cloud into
+a CrateDB Cloud Cluster.
 
 The supported file types are CSV, JSON, Parquet, optionally with
 gzip compression. They can be acquired from the local filesystem,
@@ -55,6 +56,7 @@
 
 """
 import logging
+from pprint import pprint  # noqa: F401
 
 import cratedb_toolkit
 from cratedb_toolkit import InputOutputResource, ManagedCluster, TableAddress
@@ -84,6 +86,12 @@ def import_csv():
     # table name will be derived from the input file name.
     cluster.load_table(source=source)
 
+    # Query data.
+    """
+    results = cluster.query('SELECT * FROM "nab-machine-failure" LIMIT 5;')
+    pprint(results)  # noqa: T203
+    """
+
 
 def import_parquet():
     """
@@ -106,6 +114,12 @@ def import_parquet():
     # Invoke import job. The destination table name is explicitly specified.
     cluster.load_table(source=source, target=target)
 
+    # Query data.
+    """
+    results = cluster.query('SELECT * FROM "yc-201907" LIMIT 5;')
+    pprint(results)  # noqa: T203
+    """
+
 
 def main():
     import_csv()

tests/conftest.py

@@ -1,10 +1,12 @@
 # Copyright (c) 2021-2023, Crate.io Inc.
 # Distributed under the terms of the AGPLv3 license, see LICENSE.
 import json
+import os
 
 import pytest
 import responses
 
+from cratedb_toolkit.api.main import ManagedClusterSettings
 from cratedb_toolkit.testing.testcontainers.util import PytestTestcontainerAdapter
 from cratedb_toolkit.util import DatabaseAdapter
 from cratedb_toolkit.util.common import setup_logging
@@ -191,4 +193,20 @@ def mock_cloud_import():
     )
 
 
+@pytest.fixture(scope="session", autouse=True)
+def reset_environment():
+    """
+    Reset all environment variables in use, so that they do not pollute the test suite.
+    """
+    envvars = []
+    specs = ManagedClusterSettings.settings_spec
+    for spec in specs:
+        envvars.append(spec.click.envvar)
+    for envvar in envvars:
+        try:
+            del os.environ[envvar]
+        except KeyError:
+            pass
+
+
 setup_logging()