diff --git a/chia/data_layer/data_layer.py b/chia/data_layer/data_layer.py index 75c422af35d6..b29775eb65d1 100644 --- a/chia/data_layer/data_layer.py +++ b/chia/data_layer/data_layer.py @@ -183,12 +183,32 @@ class DataLayer: @contextlib.asynccontextmanager async def manage(self) -> AsyncIterator[None]: - await self._start() + sql_log_path: Optional[Path] = None + if self.config.get("log_sqlite_cmds", False): + sql_log_path = path_from_root(self.root_path, "log/data_sql.log") + self.log.info(f"logging SQL commands to {sql_log_path}") + + self._data_store = await DataStore.create(database=self.db_path, sql_log_path=sql_log_path) + self._wallet_rpc = await self.wallet_rpc_init + + self.periodically_manage_data_task = asyncio.create_task(self.periodically_manage_data()) try: yield finally: - self._close() - await self._await_closed() + # TODO: review for anything else we need to do here + self._shut_down = True + if self._wallet_rpc is not None: + self.wallet_rpc.close() + + if self.periodically_manage_data_task is not None: + try: + self.periodically_manage_data_task.cancel() + except asyncio.CancelledError: + pass + if self._data_store is not None: + await self.data_store.close() + if self._wallet_rpc is not None: + await self.wallet_rpc.await_closed() def _set_state_changed_callback(self, callback: StateChangedProtocol) -> None: self.state_changed_callback = callback @@ -202,34 +222,6 @@ class DataLayer: def set_server(self, server: ChiaServer) -> None: self._server = server - async def _start(self) -> None: - sql_log_path: Optional[Path] = None - if self.config.get("log_sqlite_cmds", False): - sql_log_path = path_from_root(self.root_path, "log/data_sql.log") - self.log.info(f"logging SQL commands to {sql_log_path}") - - self._data_store = await DataStore.create(database=self.db_path, sql_log_path=sql_log_path) - self._wallet_rpc = await self.wallet_rpc_init - - self.periodically_manage_data_task = asyncio.create_task(self.periodically_manage_data()) - - def _close(self) -> None: - # TODO: review for anything else we need to do here - self._shut_down = True - if self._wallet_rpc is not None: - self.wallet_rpc.close() - - async def _await_closed(self) -> None: - if self.periodically_manage_data_task is not None: - try: - self.periodically_manage_data_task.cancel() - except asyncio.CancelledError: - pass - if self._data_store is not None: - await self.data_store.close() - if self._wallet_rpc is not None: - await self.wallet_rpc.await_closed() - async def wallet_log_in(self, fingerprint: int) -> int: result = await self.wallet_rpc.log_in(fingerprint) if not result.get("success", False):