diff --git a/mycli/main.py b/mycli/main.py index fac68808..7258712c 100755 --- a/mycli/main.py +++ b/mycli/main.py @@ -820,7 +820,10 @@ def output_res(results: Generator[SQLResult], start: float) -> None: nonlocal mutating result_count = 0 for result in results: - title, cur, headers, status = result.get_output() + title = result.title + cur = result.results + headers = result.headers + status = result.status command = result.command logger.debug("title: %r", title) logger.debug("headers: %r", headers) @@ -837,7 +840,7 @@ def output_res(results: Generator[SQLResult], start: float) -> None: except ValueError as e: self.echo(f"Invalid watch sleep time provided ({e}).", err=True, fg="red") sys.exit(1) - if is_select(status) and cur and cur.rowcount > threshold: + if is_select(status) and isinstance(cur, Cursor) and cur.rowcount > threshold: self.echo( f"The result set has more than {threshold} rows.", fg="red", @@ -887,7 +890,10 @@ def output_res(results: Generator[SQLResult], start: float) -> None: if self.show_warnings and isinstance(cur, Cursor) and cur.warning_count > 0: warnings = sqlexecute.run("SHOW WARNINGS") for warning in warnings: - title, cur, headers, status = warning.get_output() + title = warning.title + cur = warning.results + headers = warning.headers + status = warning.status formatted = self.format_output( title, cur, @@ -1351,7 +1357,9 @@ def run_query(self, query: str, new_line: bool = True) -> None: assert self.sqlexecute is not None results = self.sqlexecute.run(query) for result in results: - title, cur, headers, _status = result.get_output() + title = result.title + cur = result.results + headers = result.headers self.main_formatter.query = query self.redirect_formatter.query = query output = self.format_output( @@ -1369,7 +1377,9 @@ def run_query(self, query: str, new_line: bool = True) -> None: if self.show_warnings and isinstance(cur, Cursor) and cur.warning_count > 0: warnings = self.sqlexecute.run("SHOW WARNINGS") for warning in warnings: - title, cur, headers, _status = warning.get_output() + title = warning.title + cur = warning.results + headers = warning.headers output = self.format_output( title, cur, @@ -1385,7 +1395,7 @@ def format_output( self, title: str | None, cur: Cursor | list[tuple] | None, - headers: list[str] | None, + headers: list[str] | str | None, expanded: bool = False, is_redirected: bool = False, null_string: str | None = None, diff --git a/mycli/packages/sqlresult.py b/mycli/packages/sqlresult.py index 008af447..9572ea44 100644 --- a/mycli/packages/sqlresult.py +++ b/mycli/packages/sqlresult.py @@ -11,9 +11,6 @@ class SQLResult: status: str | None = None command: dict[str, str | float] | None = None - def get_output(self): - return self.title, self.results, self.headers, self.status - def __iter__(self): return self diff --git a/mycli/sqlexecute.py b/mycli/sqlexecute.py index f1e1f664..800a5381 100644 --- a/mycli/sqlexecute.py +++ b/mycli/sqlexecute.py @@ -491,8 +491,12 @@ def reset_connection_id(self) -> None: try: results = self.run("select connection_id()") for result in results: - _title, cur, _headers, _status = result.get_output() - self.connection_id = cur.fetchone()[0] + cur = result.results + if isinstance(cur, Cursor): + v = cur.fetchone() + self.connection_id = v[0] if v is not None else -1 + else: + raise ValueError except Exception as e: # See #1054 self.connection_id = -1 diff --git a/test/utils.py b/test/utils.py index 3278f9ce..e9010952 100644 --- a/test/utils.py +++ b/test/utils.py @@ -51,12 +51,10 @@ def run(executor, sql, rows_as_list=True): results = [] for result in executor.run(sql): - ( - title, - rows, - headers, - status, - ) = result.get_output() + title = result.title + rows = result.results + headers = result.headers + status = result.status rows = list(rows) if (rows_as_list and rows) else rows results.append({"title": title, "rows": rows, "headers": headers, "status": status})