diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 9c46197c8..cc51317b2 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -25,6 +25,7 @@ Commands: dag Render the DAG as an html file. diff Show the diff between the local state and the... dlt_refresh Attaches to a DLT pipeline with the option to... + environments Prints the list of SQLMesh environments with... evaluate Evaluate a model and return a dataframe with a... fetchdf Run a SQL query and display the results. format Format all SQL models and audits. @@ -146,6 +147,17 @@ Options: --help Show this message and exit. ``` +## environments +``` +Usage: sqlmesh environments [OPTIONS] + + Prints the list of SQLMesh environments with its expiry datetime. + +Options: + -e, --show-expiry Prints the expiry datetime of the environments. + --help Show this message and exit. +``` + ## evaluate ``` diff --git a/docs/reference/notebook.md b/docs/reference/notebook.md index 0d2c26ebb..555aeaebd 100644 --- a/docs/reference/notebook.md +++ b/docs/reference/notebook.md @@ -243,6 +243,14 @@ options: --force, -f If set it will overwrite existing models with the new generated models from the DLT tables. ``` +#### environments +``` +%environments [--show-expiry] +Prints the list of SQLMesh environments with its expiry datetime. +options: + --show-expiry, -e If set will print the expiry datetime of the environments +``` + #### fetchdf ``` %%fetchdf [df_var] diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py index 39189206b..03ebfa49c 100644 --- a/sqlmesh/cli/main.py +++ b/sqlmesh/cli/main.py @@ -966,3 +966,19 @@ def dlt_refresh( ctx.obj.console.log_success(f"Updated SQLMesh project with models:\n{model_names}") else: ctx.obj.console.log_success("All SQLMesh models are up to date.") + + +@cli.command("environments") +@click.option( + "-e", + "--show-expiry", + is_flag=True, + help="Prints the expiry datetime of the environments.", + default=False, +) +@click.pass_obj +@error_handler +@cli_analytics +def environments(obj: Context, show_expiry: bool) -> None: + """Prints the list of SQLMesh environments with its expiry datetime.""" + obj.print_environment_names(show_expiry=show_expiry) diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py index 9869096b4..f755a0b6c 100644 --- a/sqlmesh/core/context.py +++ b/sqlmesh/core/context.py @@ -113,7 +113,7 @@ from sqlmesh.core.user import User from sqlmesh.utils import UniqueKeyDict from sqlmesh.utils.dag import DAG -from sqlmesh.utils.date import TimeLike, now_ds, to_timestamp, format_tz_datetime +from sqlmesh.utils.date import TimeLike, now_ds, to_timestamp, format_tz_datetime, time_like_to_str from sqlmesh.utils.errors import ( CircuitBreakerError, ConfigError, @@ -1972,6 +1972,24 @@ def print_info(self, skip_connection: bool = False, verbose: bool = False) -> No if state_connection: self._try_connection("state backend", state_connection.connection_validator()) + @python_api_analytics + def print_environment_names(self, show_expiry: bool) -> None: + """Prints all environment names along with expiry datetime if show_expiry is True.""" + environment_names = self._new_state_sync().get_environment_names(get_expiry_ts=show_expiry) + if not environment_names: + error_msg = "Environments were not found." + raise SQLMeshError(error_msg) + output = ( + [ + f"{name} - {time_like_to_str(ts)}" if ts else f"{name} - No Expiry" + for name, ts in environment_names + ] + if show_expiry + else [name[0] for name in environment_names] + ) + output_str = "\n".join([str(len(output)), *output]) + self.console.log_status_update(f"Number of SQLMesh environments are: {output_str}") + def close(self) -> None: """Releases all resources allocated by this context.""" if self._snapshot_evaluator: diff --git a/sqlmesh/core/state_sync/base.py b/sqlmesh/core/state_sync/base.py index 9aeef86ec..2a61f5aab 100644 --- a/sqlmesh/core/state_sync/base.py +++ b/sqlmesh/core/state_sync/base.py @@ -135,6 +135,16 @@ def get_environments(self) -> t.List[Environment]: A list of all environments. """ + @abc.abstractmethod + def get_environment_names( + self, get_expiry_ts: bool = True + ) -> t.Optional[t.List[t.Tuple[str, ...]]]: + """Fetches all environment names along with expiry datetime if get_expiry_ts is True. + + Returns: + A list of all environment names along with expiry datetime if get_expiry_ts is True. + """ + @abc.abstractmethod def max_interval_end_per_model( self, diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py index ed24529c6..0ad555242 100644 --- a/sqlmesh/core/state_sync/engine_adapter.py +++ b/sqlmesh/core/state_sync/engine_adapter.py @@ -733,6 +733,21 @@ def get_environments(self) -> t.List[Environment]: self._environment_from_row(row) for row in self._fetchall(self._environments_query()) ] + def get_environment_names( + self, get_expiry_ts: bool = True + ) -> t.Optional[t.List[t.Tuple[str, ...]]]: + """Fetches all environment names along with expiry datetime if get_expiry_ts is True. + + Returns: + A list of all environment names along with expiry datetime if get_expiry_ts is True. + """ + name_field = ["name"] + return self._fetchall( + self._environments_query( + required_fields=name_field if not get_expiry_ts else name_field + ["expiration_ts"] + ), + ) + def _environment_from_row(self, row: t.Tuple[str, ...]) -> Environment: return Environment(**{field: row[i] for i, field in enumerate(Environment.all_fields())}) @@ -740,9 +755,11 @@ def _environments_query( self, where: t.Optional[str | exp.Expression] = None, lock_for_update: bool = False, + required_fields: t.Optional[t.List[str]] = None, ) -> exp.Select: + query_fields = required_fields if required_fields else Environment.all_fields() query = ( - exp.select(*(exp.to_identifier(field) for field in Environment.all_fields())) + exp.select(*(exp.to_identifier(field) for field in query_fields)) .from_(self.environments_table) .where(where) ) diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py index 9b3294abd..18fe79464 100644 --- a/sqlmesh/magics.py +++ b/sqlmesh/magics.py @@ -1009,6 +1009,20 @@ def clean(self, context: Context, line: str) -> None: context.clear_caches() context.console.log_success("SQLMesh cache and build artifacts cleared") + @magic_arguments() + @argument( + "--show-expiry", + "-e", + action="store_true", + help="Prints the expiration datetime of the environments.", + ) + @line_magic + @pass_sqlmesh_context + def environments(self, context: Context, line: str) -> None: + """Prints the list of SQLMesh environments with its expiry datetime.""" + args = parse_argstring(self.environments, line) + context.print_environment_names(show_expiry=args.show_expiry) + def register_magics() -> None: try: diff --git a/sqlmesh/schedulers/airflow/state_sync.py b/sqlmesh/schedulers/airflow/state_sync.py index 2b58395ad..5b9af9986 100644 --- a/sqlmesh/schedulers/airflow/state_sync.py +++ b/sqlmesh/schedulers/airflow/state_sync.py @@ -68,6 +68,14 @@ def get_environments(self) -> t.List[Environment]: """ return self._client.get_environments() + def get_environment_names( + self, get_expiry_ts: bool = True + ) -> t.Optional[t.List[t.Tuple[str, ...]]]: + """Fetches all environment names along with expiry datetime if get_expiry_ts is True.""" + raise NotImplementedError( + "get_environment_names method is not implemented for the Airflow state sync." + ) + def max_interval_end_per_model( self, environment: str, diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 8606138e9..4e5945d2a 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -10,7 +10,7 @@ from sqlmesh.cli.main import cli from sqlmesh.core.context import Context from sqlmesh.integrations.dlt import generate_dlt_models -from sqlmesh.utils.date import yesterday_ds +from sqlmesh.utils.date import now_ds, time_like_to_str, timedelta, to_datetime, yesterday_ds FREEZE_TIME = "2023-01-01 00:00:00 UTC" @@ -970,3 +970,115 @@ def test_init_project_dialects(tmp_path): assert config == f"{config_start}{expected_config}{config_end}" remove(tmp_path / "config.yaml") + + +def test_environments(runner, tmp_path): + create_example_project(tmp_path) + + # create dev environment and backfill + runner.invoke( + cli, + [ + "--log-file-dir", + tmp_path, + "--paths", + tmp_path, + "plan", + "dev", + "--no-prompts", + "--auto-apply", + ], + ) + + result = runner.invoke( + cli, + [ + "--log-file-dir", + tmp_path, + "--paths", + tmp_path, + "environments", + ], + ) + assert result.exit_code == 0 + assert result.output == "Number of SQLMesh environments are: 1\ndev\n" + + # # create dev2 environment from dev environment + # # Input: `y` to apply and virtual update + runner.invoke( + cli, + [ + "--log-file-dir", + tmp_path, + "--paths", + tmp_path, + "plan", + "dev2", + "--create-from", + "dev", + "--include-unmodified", + ], + input="y\n", + ) + + result = runner.invoke( + cli, + [ + "--log-file-dir", + tmp_path, + "--paths", + tmp_path, + "environments", + ], + ) + assert result.exit_code == 0 + assert result.output == "Number of SQLMesh environments are: 2\ndev\ndev2\n" + + result = runner.invoke( + cli, + [ + "--log-file-dir", + tmp_path, + "--paths", + tmp_path, + "environments", + "--show-expiry", + ], + ) + assert result.exit_code == 0 + ttl = time_like_to_str(to_datetime(now_ds()) + timedelta(days=7)) + assert result.output == f"Number of SQLMesh environments are: 2\ndev - {ttl}\ndev2 - {ttl}\n" + + # Example project models have start dates, so there are no date prompts + # for the `prod` environment. + # Input: `y` to apply and backfill + runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "plan"], input="y\n") + result = runner.invoke( + cli, + [ + "--log-file-dir", + tmp_path, + "--paths", + tmp_path, + "environments", + ], + ) + assert result.exit_code == 0 + assert result.output == "Number of SQLMesh environments are: 3\ndev\ndev2\nprod\n" + + result = runner.invoke( + cli, + [ + "--log-file-dir", + tmp_path, + "--paths", + tmp_path, + "environments", + "--show-expiry", + ], + ) + assert result.exit_code == 0 + assert ( + result.output + == f"Number of SQLMesh environments are: 3\ndev - {ttl}\ndev2 - {ttl}\nprod - No Expiry\n" + )