Refactor connect() into Connect class

Author: erezsh

data_diff/__init__.py

@@ -1,7 +1,7 @@
 from typing import Sequence, Tuple, Iterator, Optional, Union
 
 from .tracking import disable_tracking
-from .sqeleton.databases.connect import connect
+from .databases import connect
 from .sqeleton.databases.database_types import DbKey, DbTime, DbPath
 from .diff_tables import Algorithm
 from .hashdiff_tables import HashDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR

data_diff/__main__.py

@@ -14,8 +14,8 @@
 from .hashdiff_tables import HashDiffer, DEFAULT_BISECTION_THRESHOLD, DEFAULT_BISECTION_FACTOR
 from .joindiff_tables import TABLE_WRITE_LIMIT, JoinDiffer
 from .table_segment import TableSegment
-from .databases.database_types import create_schema
-from .databases.connect import connect
+from .sqeleton.databases.database_types import create_schema
+from .databases import connect
 from .parse_time import parse_time_before_now, UNITS_STR, ParseError
 from .config import apply_config_from_file
 from .tracking import disable_tracking

data_diff/databases/__init__.py

@@ -1,4 +1,4 @@
-# from data_diff.sqeleton.databases.base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, QueryError, ConnectError
+from data_diff.sqeleton.databases.base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, QueryError, ConnectError
 
 from .postgresql import PostgreSQL
 from .mysql import MySQL
@@ -13,4 +13,5 @@
 from .vertica import Vertica
 from .duckdb import DuckDB
 
-from .connect import connect_to_uri
+from ._connect import connect
+

data_diff/databases/_connect.py

@@ -0,0 +1,41 @@
+from data_diff.sqeleton.databases.connect import MatchUriPath, Connect
+
+from .postgresql import PostgreSQL
+from .mysql import MySQL
+from .oracle import Oracle
+from .snowflake import Snowflake
+from .bigquery import BigQuery
+from .redshift import Redshift
+from .presto import Presto
+from .databricks import Databricks
+from .trino import Trino
+from .clickhouse import Clickhouse
+from .vertica import Vertica
+
+
+
+MATCH_URI_PATH = {
+    "postgresql": MatchUriPath(PostgreSQL, ["database?"], help_str="postgresql://<user>:<pass>@<host>/<database>"),
+    "mysql": MatchUriPath(MySQL, ["database?"], help_str="mysql://<user>:<pass>@<host>/<database>"),
+    "oracle": MatchUriPath(Oracle, ["database?"], help_str="oracle://<user>:<pass>@<host>/<database>"),
+    # "mssql": MatchUriPath(MsSQL, ["database?"], help_str="mssql://<user>:<pass>@<host>/<database>"),
+    "redshift": MatchUriPath(Redshift, ["database?"], help_str="redshift://<user>:<pass>@<host>/<database>"),
+    "snowflake": MatchUriPath(
+        Snowflake,
+        ["database", "schema"],
+        ["warehouse"],
+        help_str="snowflake://<user>:<pass>@<account>/<database>/<SCHEMA>?warehouse=<WAREHOUSE>",
+    ),
+    "presto": MatchUriPath(Presto, ["catalog", "schema"], help_str="presto://<user>@<host>/<catalog>/<schema>"),
+    "bigquery": MatchUriPath(BigQuery, ["dataset"], help_str="bigquery://<project>/<dataset>"),
+    "databricks": MatchUriPath(
+        Databricks,
+        ["catalog", "schema"],
+        help_str="databricks://:access_token@server_name/http_path",
+    ),
+    "trino": MatchUriPath(Trino, ["catalog", "schema"], help_str="trino://<user>@<host>/<catalog>/<schema>"),
+    "clickhouse": MatchUriPath(Clickhouse, ["database?"], help_str="clickhouse://<user>:<pass>@<host>/<database>"),
+    "vertica": MatchUriPath(Vertica, ["database?"], help_str="vertica://<user>:<pass>@<host>/<database>"),
+}
+
+connect = Connect(MATCH_URI_PATH)

data_diff/databases/connect.py

@@ -1,8 +1,8 @@
-from data_diff.sqeleton.databases.postgresql import PostgresqlDialect, PostgreSQL
+from data_diff.sqeleton.databases import postgresql
 from .base import BaseDialect
 
-class PostgresqlDialect(BaseDialect, PostgresqlDialect):
+class PostgresqlDialect(BaseDialect, postgresql.PostgresqlDialect):
     pass
 
-class PostgreSQL(PostgreSQL):
+class PostgreSQL(postgresql.PostgreSQL):
     dialect = PostgresqlDialect()

data_diff/databases/postgresql.py

@@ -12,4 +12,3 @@
 from .clickhouse import Clickhouse
 from .vertica import Vertica
 
-from .connect import connect_to_uri

data_diff/sqeleton/databases/__init__.py

@@ -1,4 +1,4 @@
-from typing import Type, List, Optional, Union
+from typing import Type, List, Optional, Union, Dict
 from itertools import zip_longest
 import dsnparse
 
@@ -94,134 +94,138 @@ def match_path(self, dsn):
 }
 
 
-def connect_to_uri(db_uri: str, thread_count: Optional[int] = 1) -> Database:
-    """Connect to the given database uri
-
-    thread_count determines the max number of worker threads per database,
-    if relevant. None means no limit.
-
-    Parameters:
-        db_uri (str): The URI for the database to connect
-        thread_count (int, optional): Size of the threadpool. Ignored by cloud databases. (default: 1)
-
-    Note: For non-cloud databases, a low thread-pool size may be a performance bottleneck.
-
-    Supported schemes:
-    - postgresql
-    - mysql
-    - oracle
-    - snowflake
-    - bigquery
-    - redshift
-    - presto
-    - databricks
-    - trino
-    - clickhouse
-    - vertica
-    """
-
-    dsn = dsnparse.parse(db_uri)
-    if len(dsn.schemes) > 1:
-        raise NotImplementedError("No support for multiple schemes")
-    (scheme,) = dsn.schemes
-
-    try:
-        matcher = MATCH_URI_PATH[scheme]
-    except KeyError:
-        raise NotImplementedError(f"Scheme {scheme} currently not supported")
-
-    cls = matcher.database_cls
-
-    if scheme == "databricks":
-        assert not dsn.user
-        kw = {}
-        kw["access_token"] = dsn.password
-        kw["http_path"] = dsn.path
-        kw["server_hostname"] = dsn.host
-        kw.update(dsn.query)
-    elif scheme == 'duckdb':
-        kw = {}
-        kw['filepath'] = dsn.dbname
-        kw['dbname'] = dsn.user
-    else:
-        kw = matcher.match_path(dsn)
-
-        if scheme == "bigquery":
-            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
+@dataclass
+class Connect:
+    match_uri_path: Dict[str, MatchUriPath]
+
+    def connect_to_uri(self, db_uri: str, thread_count: Optional[int] = 1) -> Database:
+        """Connect to the given database uri
+
+        thread_count determines the max number of worker threads per database,
+        if relevant. None means no limit.
+
+        Parameters:
+            db_uri (str): The URI for the database to connect
+            thread_count (int, optional): Size of the threadpool. Ignored by cloud databases. (default: 1)
+
+        Note: For non-cloud databases, a low thread-pool size may be a performance bottleneck.
+
+        Supported schemes:
+        - postgresql
+        - mysql
+        - oracle
+        - snowflake
+        - bigquery
+        - redshift
+        - presto
+        - databricks
+        - trino
+        - clickhouse
+        - vertica
+        """
+
+        dsn = dsnparse.parse(db_uri)
+        if len(dsn.schemes) > 1:
+            raise NotImplementedError("No support for multiple schemes")
+        (scheme,) = dsn.schemes
+
+        try:
+            matcher = self.match_uri_path[scheme]
+        except KeyError:
+            raise NotImplementedError(f"Scheme {scheme} currently not supported")
+
+        cls = matcher.database_cls
+
+        if scheme == "databricks":
+            assert not dsn.user
+            kw = {}
+            kw["access_token"] = dsn.password
+            kw["http_path"] = dsn.path
+            kw["server_hostname"] = dsn.host
+            kw.update(dsn.query)
+        elif scheme == 'duckdb':
+            kw = {}
+            kw['filepath'] = dsn.dbname
+            kw['dbname'] = dsn.user
         else:
-            kw["host"] = dsn.host
-            kw["port"] = dsn.port
-            kw["user"] = dsn.user
-            if dsn.password:
-                kw["password"] = dsn.password
-
-    kw = {k: v for k, v in kw.items() if v is not None}
-
-    if issubclass(cls, ThreadedDatabase):
-        return cls(thread_count=thread_count, **kw)
-
-    return cls(**kw)
-
-
-def connect_with_dict(d, thread_count):
-    d = dict(d)
-    driver = d.pop("driver")
-    try:
-        matcher = MATCH_URI_PATH[driver]
-    except KeyError:
-        raise NotImplementedError(f"Driver {driver} currently not supported")
-
-    cls = matcher.database_cls
-    if issubclass(cls, ThreadedDatabase):
-        return cls(thread_count=thread_count, **d)
-
-    return cls(**d)
+            kw = matcher.match_path(dsn)
 
+            if scheme == "bigquery":
+                kw["project"] = dsn.host
+                return cls(**kw)
 
-def connect(db_conf: Union[str, dict], thread_count: Optional[int] = 1) -> Database:
-    """Connect to a database using the given database configuration.
-
-    Configuration can be given either as a URI string, or as a dict of {option: value}.
-
-    The dictionary configuration uses the same keys as the TOML 'database' definition given with --conf.
-
-    thread_count determines the max number of worker threads per database,
-    if relevant. None means no limit.
-
-    Parameters:
-        db_conf (str | dict): The configuration for the database to connect. URI or dict.
-        thread_count (int, optional): Size of the threadpool. Ignored by cloud databases. (default: 1)
-
-    Note: For non-cloud databases, a low thread-pool size may be a performance bottleneck.
-
-    Supported drivers:
-    - postgresql
-    - mysql
-    - oracle
-    - snowflake
-    - bigquery
-    - redshift
-    - presto
-    - databricks
-    - trino
-    - clickhouse
-    - vertica
-
-    Example:
-        >>> connect("mysql://localhost/db")
-        <data_diff.databases.mysql.MySQL object at 0x0000025DB45F4190>
-        >>> connect({"driver": "mysql", "host": "localhost", "database": "db"})
-        <data_diff.databases.mysql.MySQL object at 0x0000025DB3F94820>
-    """
-    if isinstance(db_conf, str):
-        return connect_to_uri(db_conf, thread_count)
-    elif isinstance(db_conf, dict):
-        return connect_with_dict(db_conf, thread_count)
-    raise TypeError(f"db configuration must be a URI string or a dictionary. Instead got '{db_conf}'.")
+            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
+                if dsn.password:
+                    kw["password"] = dsn.password
+
+        kw = {k: v for k, v in kw.items() if v is not None}
+
+        if issubclass(cls, ThreadedDatabase):
+            return cls(thread_count=thread_count, **kw)
+
+        return cls(**kw)
+
+
+    def connect_with_dict(self, d, thread_count):
+        d = dict(d)
+        driver = d.pop("driver")
+        try:
+            matcher = self.match_uri_path[driver]
+        except KeyError:
+            raise NotImplementedError(f"Driver {driver} currently not supported")
+
+        cls = matcher.database_cls
+        if issubclass(cls, ThreadedDatabase):
+            return cls(thread_count=thread_count, **d)
+
+        return cls(**d)
+
+
+    def __call__(self, db_conf: Union[str, dict], thread_count: Optional[int] = 1) -> Database:
+        """Connect to a database using the given database configuration.
+
+        Configuration can be given either as a URI string, or as a dict of {option: value}.
+
+        The dictionary configuration uses the same keys as the TOML 'database' definition given with --conf.
+
+        thread_count determines the max number of worker threads per database,
+        if relevant. None means no limit.
+
+        Parameters:
+            db_conf (str | dict): The configuration for the database to connect. URI or dict.
+            thread_count (int, optional): Size of the threadpool. Ignored by cloud databases. (default: 1)
+
+        Note: For non-cloud databases, a low thread-pool size may be a performance bottleneck.
+
+        Supported drivers:
+        - postgresql
+        - mysql
+        - oracle
+        - snowflake
+        - bigquery
+        - redshift
+        - presto
+        - databricks
+        - trino
+        - clickhouse
+        - vertica
+
+        Example:
+            >>> connect("mysql://localhost/db")
+            <data_diff.databases.mysql.MySQL object at 0x0000025DB45F4190>
+            >>> connect({"driver": "mysql", "host": "localhost", "database": "db"})
+            <data_diff.databases.mysql.MySQL object at 0x0000025DB3F94820>
+        """
+        if isinstance(db_conf, str):
+            return self.connect_to_uri(db_conf, thread_count)
+        elif isinstance(db_conf, dict):
+            return self.connect_with_dict(db_conf, thread_count)
+        raise TypeError(f"db configuration must be a URI string or a dictionary. Instead got '{db_conf}'.")

data_diff/sqeleton/databases/connect.py

@@ -1,12 +1,12 @@
 import unittest
 
 from .common import str_to_checksum, TEST_MYSQL_CONN_STRING
-from data_diff.databases import connect_to_uri
+from data_diff.databases import connect
 
 
 class TestDatabase(unittest.TestCase):
     def setUp(self):
-        self.mysql = connect_to_uri(TEST_MYSQL_CONN_STRING)
+        self.mysql = connect(TEST_MYSQL_CONN_STRING)
 
     def test_connect_to_db(self):
         self.assertEqual(1, self.mysql.query("SELECT 1", int))
@@ -21,9 +21,9 @@ def test_md5_as_int(self):
 
 class TestConnect(unittest.TestCase):
     def test_bad_uris(self):
-        self.assertRaises(ValueError, connect_to_uri, "p")
-        self.assertRaises(ValueError, connect_to_uri, "postgresql:///bla/foo")
-        self.assertRaises(ValueError, connect_to_uri, "snowflake://user:pass@bya42734/xdiffdev/TEST1")
+        self.assertRaises(ValueError, connect, "p")
+        self.assertRaises(ValueError, connect, "postgresql:///bla/foo")
+        self.assertRaises(ValueError, connect, "snowflake://user:pass@bya42734/xdiffdev/TEST1")
         self.assertRaises(
-            ValueError, connect_to_uri, "snowflake://user:pass@bya42734/xdiffdev/TEST1?warehouse=ha&schema=dup"
+            ValueError, connect, "snowflake://user:pass@bya42734/xdiffdev/TEST1?warehouse=ha&schema=dup"
         )