diff --git a/cylc/flow/dbstatecheck.py b/cylc/flow/dbstatecheck.py index 4a7d960f3fc..5537f4eacc6 100644 --- a/cylc/flow/dbstatecheck.py +++ b/cylc/flow/dbstatecheck.py @@ -282,3 +282,11 @@ def task_state_met( return bool( self.workflow_state_query(task, cycle, status, output, flow_num) ) + + +class CylcWorkflowDBCheckerContext(CylcWorkflowDBChecker): + def __enter__(self): + return self + + def __exit__(self, *args): + self.conn.close() diff --git a/cylc/flow/scripts/workflow_state.py b/cylc/flow/scripts/workflow_state.py index c0b9df9867f..72c8bcd73e6 100755 --- a/cylc/flow/scripts/workflow_state.py +++ b/cylc/flow/scripts/workflow_state.py @@ -153,6 +153,9 @@ async def check(self): flow_num=self.args['flow_num'] ) + def disconnect(self): + self.checker.conn.close() + def get_option_parser() -> COP: parser = COP( @@ -285,3 +288,5 @@ def main(parser: COP, options: 'Values', workflow_id: str) -> None: output=options.output, flow_num=options.flow_num )) + + spoller.disconnect() diff --git a/cylc/flow/xtriggers/workflow_state.py b/cylc/flow/xtriggers/workflow_state.py index f1cb480eb39..c661fb1089d 100644 --- a/cylc/flow/xtriggers/workflow_state.py +++ b/cylc/flow/xtriggers/workflow_state.py @@ -19,7 +19,7 @@ from metomi.isodatetime.parsers import TimePointParser from cylc.flow.cycling.util import add_offset -from cylc.flow.dbstatecheck import CylcWorkflowDBChecker +from cylc.flow.dbstatecheck import CylcWorkflowDBCheckerContext from cylc.flow.pathutil import get_cylc_run_dir from cylc.flow.workflow_files import infer_latest_run_from_id from cylc.flow.exceptions import WorkflowConfigError @@ -88,27 +88,26 @@ def workflow_state( # It could mean the target workflow has not started yet, # but it could also mean a typo in the workflow ID, so # so don't hide the error. - checker = CylcWorkflowDBChecker(cylc_run_dir, workflow) - - # Point validity can only be checked at run time. - # Bad function arg templating can cause a syntax error. - if checker.point_fmt is None: - # Integer cycling: raises ValueError if bad. - int(point) - else: - # Datetime cycling: raises ISO8601SyntaxError if bad - point = str( - TimePointParser().parse( - point, dump_format=checker.point_fmt + with CylcWorkflowDBCheckerContext(cylc_run_dir, workflow) as checker: + # Point validity can only be checked at run time. + # Bad function arg templating can cause a syntax error. + if checker.point_fmt is None: + # Integer cycling: raises ValueError if bad. + int(point) + else: + # Datetime cycling: raises ISO8601SyntaxError if bad + point = str( + TimePointParser().parse( + point, dump_format=checker.point_fmt + ) ) - ) - if not output and not status: - status = "succeeded" + if not output and not status: + status = "succeeded" - satisfied: bool = checker.task_state_met( - task, point, output=output, status=status - ) + satisfied: bool = checker.task_state_met( + task, point, output=output, status=status + ) results = { 'workflow': workflow, diff --git a/tests/integration/test_dbstatecheck.py b/tests/integration/test_dbstatecheck.py index ff834a34a56..ab6895b0f00 100644 --- a/tests/integration/test_dbstatecheck.py +++ b/tests/integration/test_dbstatecheck.py @@ -21,7 +21,7 @@ from textwrap import dedent from typing import TYPE_CHECKING -from cylc.flow.dbstatecheck import CylcWorkflowDBChecker as Checker +from cylc.flow.dbstatecheck import CylcWorkflowDBCheckerContext as Checker if TYPE_CHECKING: