diff options
Diffstat (limited to 'gn3')
| -rw-r--r-- | gn3/db_utils.py | 68 |
1 files changed, 54 insertions, 14 deletions
diff --git a/gn3/db_utils.py b/gn3/db_utils.py index 0d9bd0a..a855137 100644 --- a/gn3/db_utils.py +++ b/gn3/db_utils.py @@ -1,7 +1,7 @@ """module contains all db related stuff""" import contextlib import logging -from typing import Any, Iterator, Protocol, Tuple +from typing import Any, Iterator, Protocol from urllib.parse import urlparse import MySQLdb as mdb import xapian @@ -10,14 +10,59 @@ import xapian LOGGER = logging.getLogger(__file__) -def parse_db_url(sql_uri: str) -> Tuple: - """function to parse SQL_URI env variable note:there\ - is a default value for SQL_URI so a tuple result is\ - always expected""" +def __check_true__(val: str) -> bool: + """Check whether the variable 'val' has the string value `true`.""" + return val.strip().lower() == "true" + + +def __parse_db_opts__(opts: str) -> dict: + """Parse database options into their appropriate values. + + This assumes use of python-mysqlclient library.""" + allowed_opts = ( + "unix_socket", "connect_timeout", "compress", "named_pipe", + "init_command", "read_default_file", "read_default_group", + "cursorclass", "use_unicode", "charset", "collation", "auth_plugin", + "sql_mode", "client_flag", "multi_statements", "ssl_mode", "ssl", + "local_infile", "autocommit", "binary_prefix") + conversion_fns = { + **{opt: str for opt in allowed_opts}, + "connect_timeout": int, + "compress": __check_true__, + "use_unicode": __check_true__, + # "cursorclass": __load_cursor_class__ + "client_flag": int, + "multi_statements": __check_true__, + # "ssl": __parse_ssl_options__, + "local_infile": __check_true__, + "autocommit": __check_true__, + "binary_prefix": __check_true__ + } + keyvals = tuple(filter(bool, opts.split("&"))) + if len(keyvals) > 0: + keyvals = tuple(tuple(item.strip() for item in keyval.split("=")) for keyval in keyvals) + def __check_opt__(opt): + assert opt in allowed_opts, ( + f"Invalid database connection option ({opt}) provided.") + return opt + return { + __check_opt__(key): conversion_fns[key](val) + for key, val in keyvals + } + return {} + + +def parse_db_url(sql_uri: str) -> dict: + """Parse the `sql_uri` variable into a dict of connection parameters.""" parsed_db = urlparse(sql_uri) - return ( - parsed_db.hostname, parsed_db.username, parsed_db.password, - parsed_db.path[1:], parsed_db.port) + return { + "host": parsed_db.hostname, + "port": parsed_db.port or 3306, + "user": parsed_db.username, + "password": parsed_db.password, + "database": parsed_db.path.strip("/").strip(), + **__parse_db_opts__(parsed_db.query) + } # pylint: disable=missing-class-docstring, missing-function-docstring, too-few-public-methods @@ -30,12 +75,7 @@ class Connection(Protocol): @contextlib.contextmanager def database_connection(sql_uri: str, logger: logging.Logger = LOGGER) -> Iterator[Connection]: """Connect to MySQL database.""" - host, user, passwd, db_name, port = parse_db_url(sql_uri) - connection = mdb.connect(db=db_name, - user=user, - passwd=passwd or '', - host=host, - port=port or 3306) + connection = mdb.connect(**parse_db_url(sql_uri)) try: yield connection except mdb.Error as _mbde: |
