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"])