[to#811]Fix special characters in PG url and Mysql connection reconnect

data_diff/databases/mysql.py

@@ -1,4 +1,4 @@
-from typing import Any, ClassVar, Dict, Type
+from typing import Any, ClassVar, Dict, Type, Union
 
 import attrs
 
@@ -19,7 +19,8 @@
  ThreadedDatabase,
  import_helper,
  ConnectError,
- BaseDialect,
+ BaseDialect, 
+ ThreadLocalInterpreter,
 )
 from data_diff.databases.base import (
  MD5_HEXDIGITS,
@@ -148,3 +149,12 @@ def create_connection(self):
  elif e.errno == mysql.errorcode.ER_BAD_DB_ERROR:
  raise ConnectError("Database does not exist") from e
  raise ConnectError(*e.args) from e
+
+ def _query_in_worker(self, sql_code: Union[str, ThreadLocalInterpreter]):
+ "This method runs in a worker thread"
+ if self._init_error:
+ raise self._init_error
+ if not self.thread_local.conn.is_connected():
+ self.thread_local.conn.ping(reconnect=True, attempts=3, delay=5)
+ return self._query_conn(self.thread_local.conn, sql_code)
+

data_diff/databases/postgresql.py

@@ -1,5 +1,5 @@
 from typing import Any, ClassVar, Dict, List, Type
-
+from urllib.parse import unquote
 import attrs
 
 from data_diff.abcs.database_types import (
@@ -168,6 +168,7 @@ def create_connection(self):
 
  pg = import_postgresql()
  try:
+ self._args["password"] = unquote(self._args["password"])
  self._conn = pg.connect(
  **self._args, keepalives=1, keepalives_idle=5, keepalives_interval=2, keepalives_count=2
  )

tests/test_postgresql.py

@@ -1,9 +1,11 @@
 import unittest
 
+from urllib.parse import quote
 from data_diff.queries.api import table, commit
 from data_diff import TableSegment, HashDiffer
 from data_diff import databases as db
-from tests.common import get_conn, random_table_suffix
+from tests.common import get_conn, random_table_suffix, connect
+from data_diff import connect_to_table
 
 
 class TestUUID(unittest.TestCase):
@@ -113,3 +115,42 @@ def test_100_fields(self):
  id_ = diff[0][1][0]
  result = (id_,) + tuple("1" for x in range(100))
  self.assertEqual(diff, [("-", result)])
+
+
+class TestSpecialCharacterPassword(unittest.TestCase):
+ def setUp(self) -> None:
+ self.connection = get_conn(db.PostgreSQL)
+
+ table_suffix = random_table_suffix()
+
+ self.table_name = f"table{table_suffix}"
+ self.table = table(self.table_name)
+
+ def test_special_char_password(self):
+ password = "passw!!!@rd"
+ # Setup user with special character '@' in password
+ self.connection.query("DROP USER IF EXISTS test;", None)
+ self.connection.query(f"CREATE USER test WITH PASSWORD '{password}';", None)
+
+ password_quoted = quote(password)
+ db_config = {
+ "driver": "postgresql",
+ "host": "localhost",
+ "port": 5432,
+ "dbname": "postgres",
+ "user": "test",
+ "password": password_quoted,
+ }
+
+ # verify pythonic connection method
+ connect_to_table(
+ db_config,
+ self.table_name,
+ )
+
+ # verify connection method with URL string unquoted after it's verified
+ db_url = f"postgresql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['dbname']}"
+
+ connection_verified = connect(db_url)
+ assert connection_verified._args.get('password') == password
+