ashb commented on a change in pull request #18439:
URL: https://github.com/apache/airflow/pull/18439#discussion_r764887554
##########
File path: airflow/utils/cli.py
##########
@@ -53,51 +54,57 @@ def _check_cli_args(args):
)
-def action_logging(f: T) -> T:
- """
- Decorates function to execute function at the same time submitting
action_logging
- but in CLI context. It will call action logger callbacks twice,
- one for pre-execution and the other one for post-execution.
-
- Action logger will be called with below keyword parameters:
- sub_command : name of sub-command
- start_datetime : start datetime instance by utc
- end_datetime : end datetime instance by utc
- full_command : full command line arguments
- user : current user
- log : airflow.models.log.Log ORM instance
- dag_id : dag id (optional)
- task_id : task_id (optional)
- execution_date : execution date (optional)
- error : exception instance if there's an exception
-
- :param f: function instance
- :return: wrapped function
- """
-
- @functools.wraps(f)
- def wrapper(*args, **kwargs):
+def action_cli(check_db=True):
+ def action_logging(f: T) -> T:
+ """
+ Decorates function to execute function at the same time submitting
action_logging
+ but in CLI context. It will call action logger callbacks twice,
+ one for pre-execution and the other one for post-execution.
+
+ Action logger will be called with below keyword parameters:
+ sub_command : name of sub-command
+ start_datetime : start datetime instance by utc
+ end_datetime : end datetime instance by utc
+ full_command : full command line arguments
+ user : current user
+ log : airflow.models.log.Log ORM instance
+ dag_id : dag id (optional)
+ task_id : task_id (optional)
+ execution_date : execution date (optional)
+ error : exception instance if there's an exception
+
+ :param f: function instance
+ :return: wrapped function
"""
- An wrapper for cli functions. It assumes to have Namespace instance
- at 1st positional argument
- :param args: Positional argument. It assumes to have Namespace instance
+ @functools.wraps(f)
+ def wrapper(*args, **kwargs):
+ """
+ An wrapper for cli functions. It assumes to have Namespace instance
at 1st positional argument
- :param kwargs: A passthrough keyword argument
- """
- _check_cli_args(args)
- metrics = _build_metrics(f.__name__, args[0])
- cli_action_loggers.on_pre_execution(**metrics)
- try:
- return f(*args, **kwargs)
- except Exception as e:
- metrics['error'] = e
- raise
- finally:
- metrics['end_datetime'] = datetime.utcnow()
- cli_action_loggers.on_post_execution(**metrics)
-
- return cast(T, wrapper)
+
+ :param args: Positional argument. It assumes to have Namespace
instance
+ at 1st positional argument
+ :param kwargs: A passthrough keyword argument
+ """
+ _check_cli_args(args)
+ metrics = _build_metrics(f.__name__, args[0])
+ cli_action_loggers.on_pre_execution(**metrics)
+ try:
+ # Check and run migrations if necessary
+ if check_db:
+ check_and_run_migrations()
+ return f(*args, **kwargs)
+ except Exception as e:
+ metrics['error'] = e
+ raise
+ finally:
+ metrics['end_datetime'] = datetime.utcnow()
+ cli_action_loggers.on_post_execution(**metrics)
+
+ return cast(T, wrapper)
+
+ return action_logging
Review comment:
And then
```suggestion
if func:
return action_logging(func)
return action_logging
```
##########
File path: airflow/utils/cli.py
##########
@@ -53,51 +54,57 @@ def _check_cli_args(args):
)
-def action_logging(f: T) -> T:
- """
- Decorates function to execute function at the same time submitting
action_logging
- but in CLI context. It will call action logger callbacks twice,
- one for pre-execution and the other one for post-execution.
-
- Action logger will be called with below keyword parameters:
- sub_command : name of sub-command
- start_datetime : start datetime instance by utc
- end_datetime : end datetime instance by utc
- full_command : full command line arguments
- user : current user
- log : airflow.models.log.Log ORM instance
- dag_id : dag id (optional)
- task_id : task_id (optional)
- execution_date : execution date (optional)
- error : exception instance if there's an exception
-
- :param f: function instance
- :return: wrapped function
- """
-
- @functools.wraps(f)
- def wrapper(*args, **kwargs):
+def action_cli(check_db=True):
Review comment:
Ah, this is why we have `action_cli()` instead of `action_cli` at usage.
See https://stackoverflow.com/a/69210737/439189
```suggestion
def action_cli(func=None, check_db=True):
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
