|
37 | 37 | from prompt_toolkit.layout.processors import ConditionalProcessor, HighlightMatchingBracketProcessor |
38 | 38 | from prompt_toolkit.lexers import PygmentsLexer |
39 | 39 | from prompt_toolkit.shortcuts import CompleteStyle, PromptSession |
40 | | -from pymysql import OperationalError, err |
| 40 | +import pymysql |
41 | 41 | from pymysql.cursors import Cursor |
42 | 42 | import sqlglot |
43 | 43 | import sqlparse |
@@ -532,7 +532,7 @@ def _connect() -> None: |
532 | 532 | ssh_key_filename, |
533 | 533 | init_command, |
534 | 534 | ) |
535 | | - except OperationalError as e: |
| 535 | + except pymysql.OperationalError as e: |
536 | 536 | if e.args[0] == ERROR_CODE_ACCESS_DENIED: |
537 | 537 | if password_from_file is not None: |
538 | 538 | new_passwd = password_from_file |
@@ -566,7 +566,7 @@ def _connect() -> None: |
566 | 566 | self.echo(f"Connecting to socket {socket}, owned by user {socket_owner}", err=True) |
567 | 567 | try: |
568 | 568 | _connect() |
569 | | - except OperationalError as e: |
| 569 | + except pymysql.OperationalError as e: |
570 | 570 | # These are "Can't open socket" and 2x "Can't connect" |
571 | 571 | if [code for code in (2001, 2002, 2003) if code == e.args[0]]: |
572 | 572 | self.logger.debug("Database connection failed: %r.", e) |
@@ -919,7 +919,7 @@ def one_iteration(text: str | None = None) -> None: |
919 | 919 | output_res(res, start) |
920 | 920 | special.unset_once_if_written(self.post_redirect_command) |
921 | 921 | special.flush_pipe_once_if_written(self.post_redirect_command) |
922 | | - except err.InterfaceError: |
| 922 | + except pymysql.err.InterfaceError: |
923 | 923 | # attempt to reconnect |
924 | 924 | if not self.reconnect(): |
925 | 925 | return |
@@ -955,7 +955,7 @@ def one_iteration(text: str | None = None) -> None: |
955 | 955 | self.echo("Did not get a connection id, skip cancelling query", err=True, fg="red") |
956 | 956 | except NotImplementedError: |
957 | 957 | self.echo("Not Yet Implemented.", fg="yellow") |
958 | | - except OperationalError as e1: |
| 958 | + except pymysql.OperationalError as e1: |
959 | 959 | logger.debug("Exception: %r", e1) |
960 | 960 | if e1.args[0] in (2003, 2006, 2013): |
961 | 961 | # attempt to reconnect |
@@ -1044,7 +1044,7 @@ def reconnect(self, database: str = "") -> bool: |
1044 | 1044 | self.echo("Reconnecting...", fg="yellow") |
1045 | 1045 | try: |
1046 | 1046 | self.sqlexecute.connect() |
1047 | | - except OperationalError as e: |
| 1047 | + except pymysql.OperationalError as e: |
1048 | 1048 | self.logger.debug("Reconnect failed. e: %r", e) |
1049 | 1049 | self.echo(str(e), err=True, fg="red") |
1050 | 1050 | return False |
|
0 commit comments