Skip to content

Commit

Permalink
mypy fixes for recent changes to mysql.connector package
Browse files Browse the repository at this point in the history
  • Loading branch information
fritzm committed Jan 30, 2023
1 parent cacebe3 commit 13ab579
Show file tree
Hide file tree
Showing 6 changed files with 44 additions and 27 deletions.
2 changes: 1 addition & 1 deletion src/admin/python/lsst/qserv/admin/cli/_integration_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ def get_databases() -> List[str]:
) as connection:
with closing(connection.cursor()) as cursor:
cursor.execute("show databases;")
return [row[0] for row in cursor.fetchall()]
return [str(row[0]) for row in cursor.fetchall()]

databases = get_databases()
checkdb = "qservMeta"
Expand Down
21 changes: 12 additions & 9 deletions src/admin/python/lsst/qserv/admin/itest_load.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,12 @@
import json
import logging
import mysql.connector
from mysql.connector.cursor import MySQLCursor
from mysql.connector.abstracts import MySQLConnectionAbstract, MySQLCursorAbstract
import os
import shutil
import subprocess
from tempfile import TemporaryDirectory
from typing import Any, Dict, Generator, List, NamedTuple, Optional, Tuple
from typing import Any, Dict, Generator, List, NamedTuple, Optional, Tuple, Union, cast
import yaml

from .constants import tmp_data_dir
Expand All @@ -57,7 +57,7 @@ def unzip(source: str, destination: str) -> None:
shutil.copyfileobj(_source, _target)


def execute(cursor: MySQLCursor, stmt: str, multi: bool = False) -> None:
def execute(cursor: MySQLCursorAbstract, stmt: str, multi: bool = False) -> None:
"""Execute a statement in a cursor, and clean up the cursor so it can be
closed without causing warnings on the server."""
warnings = []
Expand All @@ -73,6 +73,11 @@ def execute(cursor: MySQLCursor, stmt: str, multi: bool = False) -> None:
cursors = [cursor]

for cursor in cursors:
# Cast here because MySQLCursorAbtract does not have with_rows for some reason, even though both of
# its subclasses do...
cursor = cast(
Union[mysql.connector.cursor.MySQLCursor, mysql.connector.cursor_cext.CMySQLCursor], cursor
)
if cursor.with_rows:
for row in cursor.fetchall():
results.append(row)
Expand Down Expand Up @@ -121,12 +126,12 @@ def _create_ref_db(ref_db_admin: str, name: str) -> None:
execute(cursor, stmt)


def _create_ref_table(cnx: mysql.connector.connection, db: str, schema_file: str) -> None:
def _create_ref_table(cnx: MySQLConnectionAbstract, db: str, schema_file: str) -> None:
"""Create a table in the mysql database used for integration test reference.
Parameters
----------
cnx : `mysql.connector.connection`
cnx : `mysql.connector.abstracts.MySQLConnectionAbstract`
Connection to the reference database that has permission to create a
table.
db : `str`
Expand All @@ -145,14 +150,12 @@ def _create_ref_table(cnx: mysql.connector.connection, db: str, schema_file: str
execute(cursor, f.read(), multi=True)


def _load_ref_data(
cnx: mysql.connector.connection, data_file: str, db: str, table: LoadTable
) -> None:
def _load_ref_data(cnx: MySQLConnectionAbstract, data_file: str, db: str, table: LoadTable) -> None:
"""Load database data into the reference database.
Parameters
----------
cnx : `mysql.connector.connection`
cnx : `mysql.connector.abstracts.MySQLConnectionAbstract`
Connection to the reference database that has permission to load data
into the table.
data_file : `str`
Expand Down
19 changes: 12 additions & 7 deletions src/admin/python/lsst/qserv/admin/mysql_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@

import logging
import mysql.connector
from typing import cast
from urllib.parse import urlparse

_log = logging.getLogger(__name__)
Expand All @@ -34,7 +35,7 @@ def mysql_connection(
uri: str,
get_warnings: bool = True,
local_infile: bool = False,
) -> mysql.connector.connection:
) -> mysql.connector.abstracts.MySQLConnectionAbstract:
"""Create a mysql.connection that is connected to a database.
Parameters
Expand All @@ -58,12 +59,16 @@ def mysql_connection(
user = parsed.username
pw = parsed.password
_log.debug("mysql_connection hostname:%s, port:%s, user:%s", hostname, port, user)
cnx = mysql.connector.connect(
user=user,
password=pw,
host=hostname,
port=port,
allow_local_infile=local_infile,
# Cast justified because no pool args passed here to connect(), so cnx cannot be PooledMySQLConnection
cnx = cast(
mysql.connector.abstracts.MySQLConnectionAbstract,
mysql.connector.connect(
user=user,
password=pw,
host=hostname,
port=port,
allow_local_infile=local_infile,
),
)
cnx.get_warnings = get_warnings
return cnx
4 changes: 2 additions & 2 deletions src/admin/python/lsst/qserv/admin/watcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
import logging
import requests
from time import sleep
from typing import List, Sequence, Set, Type, Union
from typing import List, Sequence, Set, Type, Union, Any

_log = logging.getLogger(__name__)

Expand Down Expand Up @@ -168,7 +168,7 @@ def notify(self, msg: str) -> None:
_log.error(f"Failed to send notification {msg}, exception: {e}")
_log.debug(f"Sent notification: {msg}")

def query(self, uri: str, stmt: str) -> List[Sequence[str]]:
def query(self, uri: str, stmt: str) -> Sequence[Any]:
"""Execute a mysql query at the provided URI
Parameters
Expand Down
15 changes: 10 additions & 5 deletions src/replica/python/replicaConfig.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from contextlib import closing
import mysql.connector
import logging
from typing import cast
from urllib.parse import urlparse

from lsst.qserv.admin.backoff import qserv_backoff, on_backoff
Expand All @@ -44,11 +45,15 @@ def applyConfiguration(connection: str, sql: str) -> None:
"""Apply configuration sql to the replication controller database."""
c = urlparse(connection)
with closing(
mysql.connector.connect(
user=c.username,
password=c.password,
host=c.hostname,
port=c.port,
# Cast justified because no pool args passed here to connect(), so cannot be PooledMySQLConnection
cast(
mysql.connector.abstracts.MySQLConnectionAbstract,
mysql.connector.connect(
user=c.username,
password=c.password,
host=c.hostname,
port=c.port,
),
)
) as cnx:
cnx.database = replicaDb
Expand Down
10 changes: 7 additions & 3 deletions src/schema/python/schemaMigMgr.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
import jinja2
import logging
import mysql.connector
from typing import Callable, Dict, Union
from typing import Callable, Dict, Union, cast

# MySQLInterfaceError can get thrown, we need to catch it.
# It's not exposed as a public python object but *is* used in mysql.connector unit tests.
Expand Down Expand Up @@ -271,7 +271,7 @@ def max_migration_version(self) -> Optional[Version]:
on_backoff=on_backoff(log=_log),
max_time=max_backoff_sec,
)
def _connect(self, connection: str) -> mysql.connector.connection:
def _connect(self, connection: str) -> mysql.connector.abstracts.MySQLConnectionAbstract:
url = make_url(connection)
# The database is not always guaranteed to exist yet (sometimes we connect and then create it)
# so cache it, and it can be set in the connection before use when it is known to exist.
Expand All @@ -284,7 +284,11 @@ def _connect(self, connection: str) -> mysql.connector.connection:
kwargs["unix_socket"] = url.query["socket"]
else:
kwargs.update(host=url.host, port=url.port)
return mysql.connector.connect(**kwargs)
# Cast justified because no pool args passed here to connect() so cannot be PooledMySQLConnection
return cast(
mysql.connector.abstracts.MySQLConnectionAbstract,
mysql.connector.connect(**kwargs),
)

@backoff.on_exception(
exception=mysql.connector.errors.DatabaseError,
Expand Down

0 comments on commit 13ab579

Please sign in to comment.