Add support for loading from .env files

README.md

@@ -25,24 +25,38 @@ pip install 'great_expectations_cloud[sql]'
 
 ```console
 $ gx-agent --help
-usage: gx-agent [-h] [--log-level LOG_LEVEL] [--skip-log-file SKIP_LOG_FILE] [--log-cfg-file LOG_CFG_FILE] [--version]
+usage: gx-agent [-h] [--env-file ENV_FILE] [--log-level LOG_LEVEL] [--skip-log-file SKIP_LOG_FILE] [--json-log] [--log-cfg-file LOG_CFG_FILE]
+                [--custom-log-tags CUSTOM_LOG_TAGS] [--version]
 
 optional arguments:
   -h, --help            show this help message and exit
+  --env-file ENV_FILE   Path to an environment file to load environment variables from. Defaults to None.
   --log-level LOG_LEVEL
                         Level of logging to use. Defaults to WARNING.
   --skip-log-file SKIP_LOG_FILE
                         Skip writing debug logs to a file. Defaults to False. Does not affect logging to stdout/stderr.
+  --json-log            Output logs in JSON format. Defaults to False.
   --log-cfg-file LOG_CFG_FILE
-                        Path to a logging configuration json file. Supersedes --log-level and --skip-log-file.
-  --version             Show the GX Agent version.
+                        Path to a logging configuration file in JSON format. Supersedes --log-level and --skip-log-file.
+  --custom-log-tags CUSTOM_LOG_TAGS
+                        Additional tags for logs in form of key-value pairs
+  --version             Show the gx agent version.
 ```
 
 #### Set env variables
 
+Set the following environment variables, either directly or with an [environment file](/example.env).
+
+Note: The access token should have the `Editor` role.
+
 `GX_CLOUD_ACCESS_TOKEN`
 `GX_CLOUD_ORGANIZATION_ID`
 
+These can be found in the [GX Cloud](https://app.greatexpectations.io/) Settings -> Users section page.
+
+![Sidebar](/examples/agent/imgs/sidebar.png)
+
+
 ### Start the Agent
 
 If you intend to run the Agent against local services (Cloud backend or datasources) run the Agent outside of the container.

example.env

@@ -0,0 +1,3 @@
+# These can be found in the GX Cloud: Settings -> Users page
+GX_CLOUD_ACCESS_TOKEN=<YOUR_ACCESS_TOKEN>
+GX_CLOUD_ORGANIZATION_ID=<YOUR_ORGANIZATION_ID>

great_expectations_cloud/agent/cli.py

@@ -15,6 +15,7 @@
 
 @dc.dataclass(frozen=True)
 class Arguments:
+    env_file: pathlib.Path | None
     log_level: LogLevel
     skip_log_file: bool
     json_log: bool
@@ -29,6 +30,12 @@ def _parse_args() -> Arguments:
     `Arguments` dataclass.
     """
     parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--env-file",
+        help="Path to an environment file to load environment variables from. Defaults to None.",
+        default=None,
+        type=pathlib.Path,
+    )
     parser.add_argument(
         "--log-level",
         help="Level of logging to use. Defaults to WARNING.",
@@ -61,6 +68,7 @@ def _parse_args() -> Arguments:
     parser.add_argument("--version", help="Show the gx agent version.", action="store_true")
     args = parser.parse_args()
     return Arguments(
+        env_file=args.env_file,
         log_level=args.log_level,
         skip_log_file=args.skip_log_file,
         log_cfg_file=args.log_cfg_file,
@@ -70,9 +78,38 @@ def _parse_args() -> Arguments:
     )
 
 
+def load_dotenv(env_file: pathlib.Path) -> set[str]:
+    """
+    Load environment variables from a file.
+
+    Returns a set of the environment variables that were loaded.
+    """
+    import os
+
+    # throw error if file does not exist
+    env_file.resolve(strict=True)
+    # lazy import to keep the cli fast
+    from dotenv import dotenv_values
+
+    # os.environ does not allow None values, so we filter them out
+    loaded_values: dict[str, str] = {
+        k: v for (k, v) in dotenv_values(env_file).items() if v is not None
+    }
+
+    os.environ.update(loaded_values)  # noqa: TID251 # needed for OSS to pickup the env vars
+
+    return set(loaded_values.keys())
+
+
 def main() -> None:
     # lazy imports ensure our cli is fast and responsive
     args: Arguments = _parse_args()
+
+    if args.env_file:
+        print(f"Loading environment variables from {args.env_file}")
+        loaded_env_vars = load_dotenv(args.env_file)
+        LOGGER.info(f"Loaded {len(loaded_env_vars)} environment variables.")
+
     custom_tags: dict[str, Any] = {}
     try:
         custom_tags = json.loads(args.custom_log_tags)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "great_expectations_cloud"
-version = "20240814.0.dev0"
+version = "20240816.0.dev0"
 description = "Great Expectations Cloud"
 authors = ["The Great Expectations Team <team@greatexpectations.io>"]
 repository = "https://github.com/great-expectations/cloud"
@@ -30,6 +30,7 @@ pika = "^1.3.1"
 orjson = "^3.9.7, !=3.9.10" # TODO: remove inequality once dep resolution issue is resolved
 # relying on packaging in agent code so declaring it explicitly here
 packaging = ">=21.3,<25.0"
+python-dotenv = "^1.0.1"
 tenacity = ">=8.2.3,<10.0.0"
 sqlalchemy = { version = ">=2.0", optional = false }
 # optional dependencies
@@ -189,12 +190,12 @@ required-imports = ["from __future__ import annotations"]
 max-complexity = 10
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
-"os.environ".msg = """Please do not use os.environ, instead use a pydantic.BaseSettings model"""
+"os.environ".msg = "Please do not use os.environ, instead use a pydantic.BaseSettings model"
+"os.getenv".msg = "Please do not use os.getenv, instead use a pydantic.BaseSettings model"
 "great_expectations.compatibility.pydantic".msg = "Import pydantic directly."
 "great_expectations.compatibility.sqlalchemy".msg = "Import sqlalchemy directly."
 "great_expectations.compatibility".msg = "Favor specific version imports over compatibility module."
 
-
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = [
     "F401", # unused import

tests/agent/conftest.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import logging
+import os
 import time
 import uuid
 from collections import deque
-from typing import TYPE_CHECKING, Any, Iterable, NamedTuple, TypedDict
+from typing import TYPE_CHECKING, Any, Generator, Iterable, NamedTuple, TypedDict
 
 import pytest
 from great_expectations import (  # type: ignore[attr-defined] # TODO: fix this
@@ -27,6 +28,32 @@
 LOGGER = logging.getLogger(__name__)
 
 
+@pytest.fixture
+def clean_gx_env(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]:
+    """
+    Cleanup the GX_CLOUD environment variables before and after test run.
+
+    If the environment variables are already set, save their values and restore them after the test run.
+
+    GX_CLOUD_ACCESS_TOKEN
+    GX_CLOUD_ORGANIZATION_ID
+    GX_BASE_URL
+    """
+    env_vars = ["GX_CLOUD_ACCESS_TOKEN", "GX_CLOUD_ORGANIZATION_ID", "GX_BASE_URL"]
+    prior_values = {var: os.environ[var] for var in env_vars if var in os.environ}
+
+    for var in env_vars:
+        monkeypatch.delenv(var, raising=False)
+    yield None
+    for var in env_vars:
+        monkeypatch.delenv(var, raising=False)
+
+    # TODO: remove the action of restoring the environment variables once our tests are no longer
+    # relying on prior state of the environment variables.
+    if prior_values:
+        os.environ.update(prior_values)
+
+
 @pytest.fixture
 def mock_gx_version_check(
     monkeypatch: pytest.MonkeyPatch,

tests/agent/test_cli.py

@@ -1,15 +1,55 @@
 from __future__ import annotations
 
+import os
+import pathlib
 import subprocess
+from typing import TYPE_CHECKING
 
 import pytest
 
+from great_expectations_cloud.agent.cli import load_dotenv, main
+
+if TYPE_CHECKING:
+    from pytest_mock import MockerFixture, MockType
+
+
+@pytest.fixture
+def mock_agent_run(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch) -> MockType:
+    run_patch = mocker.patch("great_expectations_cloud.agent.run_agent")
+    return run_patch
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        ("--help",),
+        ("--version",),
+        ("--env-file", "example.env"),
+        ("--log-level", "DEBUG"),
+        ("--json-log",),
+    ],
+    ids=lambda x: " ".join(x),
+)
+def test_main(monkeypatch: pytest.MonkeyPatch, mock_agent_run: MockType, args: tuple[str, ...]):
+    """Ensure that the main function runs without error."""
+    monkeypatch.delenv("GX_CLOUD_ACCESS_TOKEN", raising=False)
+    monkeypatch.delenv("GX_CLOUD_ORGANIZATION_ID", raising=False)
+
+    monkeypatch.setattr("sys.argv", ["gx-agent", *args])
+
+    try:
+        main()
+    except SystemExit as exit:
+        # as long as the exit code is 0, we are good
+        assert exit.code == 0
+
 
 @pytest.mark.parametrize(
     "cmd",
     [
         "--help",
         "-h",
+        "--version",
     ],
 )
 def test_command_retuns_zero_exit_code(cmd: str):
@@ -28,5 +68,14 @@ def test_custom_log_tags_failure():
     assert cmplt_process.returncode != 0
 
 
+def test_load_dotenv(clean_gx_env: None):
+    env_file = pathlib.Path("example.env")
+    loaded_env_vars = load_dotenv(env_file)
+    assert loaded_env_vars == {"GX_CLOUD_ACCESS_TOKEN", "GX_CLOUD_ORGANIZATION_ID"}
+
+    assert os.environ["GX_CLOUD_ACCESS_TOKEN"] == "<YOUR_ACCESS_TOKEN>"
+    assert os.environ["GX_CLOUD_ORGANIZATION_ID"] == "<YOUR_ORGANIZATION_ID>"
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-vv"])