Fixes for postgresql

Author: erezsh

data_diff/databases/postgresql.py

@@ -2,6 +2,7 @@
 from .base import ThreadedDatabase, import_helper, ConnectError
 from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, _CHECKSUM_BITSIZE, TIMESTAMP_PRECISION_POS
 
+SESSION_TIME_ZONE = None    # Changed by the tests
 
 @import_helper("postgresql")
 def import_postgresql():
@@ -47,13 +48,17 @@ def _convert_db_precision_to_digits(self, p: int) -> int:
         return super()._convert_db_precision_to_digits(p) - 2
 
     def create_connection(self):
+        if not self._args:
+            self._args['host'] = None   # psycopg2 requires 1+ arguments
+
         pg = import_postgresql()
         try:
             c = pg.connect(**self._args)
-            # c.cursor().execute("SET TIME ZONE 'UTC'")
+            if SESSION_TIME_ZONE:
+                c.cursor().execute(f"SET TIME ZONE '{SESSION_TIME_ZONE}'")
             return c
         except pg.OperationalError as e:
-            raise ConnectError(*e._args) from e
+            raise ConnectError(*e.args) from e
 
     def quote(self, s: str):
         return f'"{s}"'

tests/test_database_types.py

@@ -12,6 +12,7 @@
 from parameterized import parameterized
 
 from data_diff import databases as db
+import data_diff.databases.postgresql as postgresql
 from data_diff.utils import number_to_human
 from data_diff.diff_tables import TableDiffer, TableSegment, DEFAULT_BISECTION_THRESHOLD
 from .common import CONN_STRINGS, N_SAMPLES, N_THREADS, BENCHMARK, GIT_REVISION, random_table_suffix
@@ -20,6 +21,7 @@
 CONNS = {k: db.connect_to_uri(v, N_THREADS) for k, v in CONN_STRINGS.items()}
 
 CONNS[db.MySQL].query("SET @@session.time_zone='+00:00'", None)
+postgresql.SESSION_TIME_ZONE = 'UTC'
 
 
 class PaginatedTable: