[GitHub] [airflow] ashb commented on a change in pull request #16030: Update `airflow tasks *` commands to lookup TaskInstances from DagRun Table

[GitBox <gi...@apache.org>]

ashb commented on a change in pull request #16030:
URL: https://github.com/apache/airflow/pull/16030#discussion_r641842538



##########
File path: airflow/models/dag.py
##########
@@ -905,22 +905,24 @@ def get_num_active_runs(self, external_trigger=None, session=None):
         return query.scalar()
 
     @provide_session
-    def get_dagrun(self, execution_date, session=None):
+    def get_dagrun(self, execution_date: str = None, run_id: str = None, session=None):
         """
-        Returns the dag run for a given execution date if it exists, otherwise
+        Returns the dag run for a given execution date or run_id if it exists, otherwise
         none.
 
         :param execution_date: The execution date of the DagRun to find.
+        :param run_id: The run_id of the DagRun to find.
         :param session:
         :return: The DagRun if found, otherwise None.
         """
-        dagrun = (
-            session.query(DagRun)
-            .filter(DagRun.dag_id == self.dag_id, DagRun.execution_date == execution_date)
-            .first()
-        )
-
-        return dagrun
+        if not (execution_date or run_id):
+            raise AirflowException("You must provide either the execution_date or the run_id")

Review comment:
       There is a built in error that has this meaning, so we should raise that.
   
   ```suggestion
               raise TypeError("You must provide either the execution_date or the run_id")
   ```

##########
File path: airflow/cli/commands/task_command.py
##########
@@ -42,9 +45,26 @@
     get_dags,
     suppress_logs_and_warning,
 )
+from airflow.utils.dates import timezone
 from airflow.utils.log.logging_mixin import StreamLogWriter
 from airflow.utils.net import get_hostname
-from airflow.utils.session import create_session
+from airflow.utils.session import create_session, provide_session
+
+
+def _get_ti(dag, task, exec_date_or_run_id):

Review comment:
       ```suggestion
   def _get_ti(dag, task_id, exec_date_or_run_id):
   ```




-- 
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.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org