Cleanup database params - just pass the dict

Author: erezsh

data_diff/databases/connect.py

@@ -116,9 +116,22 @@ def connect_to_uri(db_uri: str, thread_count: Optional[int] = 1) -> Database:
     kw = matcher.match_path(dsn)
 
     if scheme == "bigquery":
-        return cls(dsn.host, **kw)
+        kw["project"] = dsn.host
+        return cls(**kw)
+
+    if scheme == "snowflake":
+        kw["account"] = dsn.host
+        assert not dsn.port
+        kw["user"] = dsn.user
+        kw["password"] = dsn.password
+    else:
+        kw["host"] = dsn.host
+        kw["port"] = dsn.port
+        kw["user"] = dsn.user
+        kw["password"] = dsn.password
+    kw = {k: v for k, v in kw.items() if v is not None}
 
     if issubclass(cls, ThreadedDatabase):
-        return cls(dsn.host, dsn.port, dsn.user, dsn.password, thread_count=thread_count, **kw)
+        return cls(thread_count=thread_count, **kw)
 
-    return cls(dsn.host, dsn.port, dsn.user, dsn.password, **kw)
+    return cls(**kw)

data_diff/databases/database_types.py

@@ -18,6 +18,11 @@ class ColType:
     pass
 
 
+class IKey(ABC):
+    "Interface for ColType, for using a column as a key in data-diff"
+    python_type: type
+
+
 @dataclass
 class PrecisionType(ColType):
     precision: int
@@ -54,7 +59,7 @@ class Float(FractionalType):
     pass
 
 
-class Decimal(FractionalType):
+class Decimal(FractionalType, IKey):
     @property
     def python_type(self) -> type:
         if self.precision == 0:
@@ -66,11 +71,6 @@ class StringType(ColType):
     pass
 
 
-class IKey(ABC):
-    "Interface for ColType, for using a column as a key in data-diff"
-    python_type: type
-
-
 class ColType_UUID(StringType, IKey):
     python_type = ArithUUID
 

data_diff/databases/mysql.py

@@ -29,14 +29,13 @@ class MySQL(ThreadedDatabase):
     }
     ROUNDS_ON_PREC_LOSS = True
 
-    def __init__(self, host, port, user, password, *, database, thread_count, **kw):
-        args = dict(host=host, port=port, database=database, user=user, password=password, **kw)
-        self._args = {k: v for k, v in args.items() if v is not None}
+    def __init__(self, *, thread_count, **kw):
+        self._args = kw
 
         super().__init__(thread_count=thread_count)
 
         # In MySQL schema and database are synonymous
-        self.default_schema = database
+        self.default_schema = kw["database"]
 
     def create_connection(self):
         mysql = import_mysql()
@@ -48,7 +47,7 @@ def create_connection(self):
             elif e.errno == mysql.errorcode.ER_BAD_DB_ERROR:
                 raise ConnectError("Database does not exist") from e
             else:
-                raise ConnectError(*e.args) from e
+                raise ConnectError(*e._args) from e
 
     def quote(self, s: str):
         return f"`{s}`"

data_diff/databases/oracle.py

@@ -24,11 +24,10 @@ class Oracle(ThreadedDatabase):
     }
     ROUNDS_ON_PREC_LOSS = True
 
-    def __init__(self, host, port, user, password, *, database, thread_count, **kw):
-        assert not port
-        self.kwargs = dict(user=user, password=password, dsn="%s/%s" % (host, database), **kw)
+    def __init__(self, *, host, database, thread_count, **kw):
+        self.kwargs = dict(dsn="%s/%s" % (host, database), **kw)
 
-        self.default_schema = user
+        self.default_schema = kw.get('user')
 
         super().__init__(thread_count=thread_count)
 

data_diff/databases/postgresql.py

@@ -35,8 +35,8 @@ class PostgreSQL(ThreadedDatabase):
 
     default_schema = "public"
 
-    def __init__(self, host, port, user, password, *, database, thread_count, **kw):
-        self.args = dict(host=host, port=port, database=database, user=user, password=password, **kw)
+    def __init__(self, *, thread_count, **kw):
+        self._args = kw
 
         super().__init__(thread_count=thread_count)
 
@@ -47,11 +47,11 @@ def _convert_db_precision_to_digits(self, p: int) -> int:
     def create_connection(self):
         pg = import_postgresql()
         try:
-            c = pg.connect(**self.args)
+            c = pg.connect(**self._args)
             # c.cursor().execute("SET TIME ZONE 'UTC'")
             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}"'

data_diff/databases/presto.py

@@ -35,11 +35,10 @@ class Presto(Database):
     }
     ROUNDS_ON_PREC_LOSS = True
 
-    def __init__(self, host, port, user, password, *, catalog, schema=None, **kw):
+    def __init__(self, **kw):
         prestodb = import_presto()
-        self.args = dict(host=host, user=user, catalog=catalog, schema=schema, **kw)
 
-        self._conn = prestodb.dbapi.connect(**self.args)
+        self._conn = prestodb.dbapi.connect(**kw)
 
     def quote(self, s: str):
         return f'"{s}"'

data_diff/databases/snowflake.py

@@ -25,19 +25,7 @@ class Snowflake(Database):
     }
     ROUNDS_ON_PREC_LOSS = False
 
-    def __init__(
-        self,
-        account: str,
-        _port: int,
-        user: str,
-        password: str,
-        *,
-        warehouse: str,
-        schema: str,
-        database: str,
-        role: str = None,
-        **kw,
-    ):
+    def __init__(self, *, schema: str, **kw):
         snowflake = import_snowflake()
         logging.getLogger("snowflake.connector").setLevel(logging.WARNING)
 
@@ -48,12 +36,6 @@ def __init__(
 
         assert '"' not in schema, "Schema name should not contain quotes!"
         self._conn = snowflake.connector.connect(
-            user=user,
-            password=password,
-            account=account,
-            role=role,
-            database=database,
-            warehouse=warehouse,
             schema=f'"{schema}"',
             **kw,
         )