From 4270e2e87540bd7e15503a871636e355f7c477dc Mon Sep 17 00:00:00 2001
From: Maxime Armstrong <46797220+maximearmstrong@users.noreply.github.com>
Date: Wed, 30 Oct 2024 13:55:21 -0400
Subject: [PATCH] [cli] Set `DAGSTER_IS_DEFS_VALIDATION_CLI` in `dagster
 definitions validate` (#25624)

## Summary & Motivation

Our design partners would like to have more control on the code and its
behavior when in validation mode: different credentials, etc.

This PR updates `dagster definitions validate` to set the env var
`DAGSTER_IS_DEFS_VALIDATION_CLI` when the command is running , like we
do so `DAGSTER_IS_DEV_CLI` and `dagster dev`. This will allow users to
control the code using the env var like:

```
import os

if os.getenv("DAGSTER_IS_DEFS_VALIDATION_CLI"):
    ...
```

## How I Tested These Changes

Additional unit test
---
 python_modules/dagster/dagster/_cli/definitions.py   |  3 +++
 .../valid_project/valid_project/gated_definitions.py | 12 ++++++++++++
 .../test_definitions_validate_command.py             |  9 +++++++++
 3 files changed, 24 insertions(+)
 create mode 100644 python_modules/dagster/dagster_tests/cli_tests/command_tests/definitions_command_projects/valid_project/valid_project/gated_definitions.py

diff --git a/python_modules/dagster/dagster/_cli/definitions.py b/python_modules/dagster/dagster/_cli/definitions.py
index b5a3a939134e1..50c0306d507fa 100644
--- a/python_modules/dagster/dagster/_cli/definitions.py
+++ b/python_modules/dagster/dagster/_cli/definitions.py
@@ -1,4 +1,5 @@
 import logging
+import os
 import sys
 
 import click
@@ -60,6 +61,8 @@ def validate_command_options(f):
     """,
 )
 def definitions_validate_command(log_level: str, log_format: str, **kwargs: ClickArgValue):
+    os.environ["DAGSTER_IS_DEFS_VALIDATION_CLI"] = "1"
+
     configure_loggers(formatter=log_format, log_level=log_level.upper())
     logger = logging.getLogger("dagster")
 
diff --git a/python_modules/dagster/dagster_tests/cli_tests/command_tests/definitions_command_projects/valid_project/valid_project/gated_definitions.py b/python_modules/dagster/dagster_tests/cli_tests/command_tests/definitions_command_projects/valid_project/valid_project/gated_definitions.py
new file mode 100644
index 0000000000000..1f450daf16195
--- /dev/null
+++ b/python_modules/dagster/dagster_tests/cli_tests/command_tests/definitions_command_projects/valid_project/valid_project/gated_definitions.py
@@ -0,0 +1,12 @@
+import os
+
+from dagster import Definitions, asset, define_asset_job
+
+if os.getenv("DAGSTER_IS_DEFS_VALIDATION_CLI"):
+
+    @asset
+    def my_gated_asset() -> None: ...
+
+    my_gated_job = define_asset_job(name="my_gated_job", selection="my_gated_asset")
+
+    defs = Definitions(assets=[my_gated_asset], jobs=[my_gated_job])
diff --git a/python_modules/dagster/dagster_tests/cli_tests/command_tests/test_definitions_validate_command.py b/python_modules/dagster/dagster_tests/cli_tests/command_tests/test_definitions_validate_command.py
index 1e5c294b920bd..1383fb064990f 100644
--- a/python_modules/dagster/dagster_tests/cli_tests/command_tests/test_definitions_validate_command.py
+++ b/python_modules/dagster/dagster_tests/cli_tests/command_tests/test_definitions_validate_command.py
@@ -70,3 +70,12 @@ def test_invalid_project(options, monkeypatch):
         assert result.exit_code == 1
         assert "Validation failed" in result.output
         assert "Duplicate asset key: AssetKey(['my_asset'])" in result.output
+
+
+def test_env_var(monkeypatch):
+    with monkeypatch.context() as m:
+        m.chdir(VALID_PROJECT_PATH)
+        # Definitions in `gated_definitions.py` are gated by the "DAGSTER_IS_DEFS_VALIDATION_CLI" environment variable
+        result = invoke_validate(options=["-f", "valid_project/gated_definitions.py"])
+        assert result.exit_code == 0
+        assert "Validation successful for code location gated_definitions.py." in result.output