Update databricks URI and tests

Author: pik94

README.md

@@ -125,7 +125,7 @@ $ data-diff \
 | BigQuery      | `bigquery://<project>/<dataset>`                                                                                                    |  💛    |
 | Redshift      | `redshift://<username>:<password>@<hostname>:5439/<database>`                                                                       |  💛    |
 | Presto        | `presto://<username>:<password>@<hostname>:8080/<database>`                                                                         |  💛    |
-| Databricks    | `databricks://<access_token>:<http_path>@<server_hostname>/<catalog>/<schema>`                                                      |  💛    |
+| Databricks    | `databricks://<http_path>:<access_token>@<server_hostname>/<catalog>/<schema>`                                                      |  💛    |
 | ElasticSearch |                                                                                                                                     |  📝    |                                                        |  📝    |
 | Planetscale   |                                                                                                                                     |  📝    |
 | Clickhouse    |                                                                                                                                     |  📝    |

data_diff/databases/__init__.py

@@ -7,5 +7,6 @@
 from .bigquery import BigQuery
 from .redshift import Redshift
 from .presto import Presto
+from .databricks import Databricks
 
 from .connect import connect_to_uri

data_diff/databases/connect.py

@@ -79,7 +79,7 @@ def match_path(self, dsn):
     "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:http_path@account/catalog/schema"
+        Databricks, ["catalog", "schema"], help_str="databricks://http_path:access_token@server_name/catalog/schema",
     )
 }
 
@@ -129,6 +129,12 @@ def connect_to_uri(db_uri: str, thread_count: Optional[int] = 1) -> Database:
         assert not dsn.port
         kw["user"] = dsn.user
         kw["password"] = dsn.password
+    if scheme == "databricks":
+        # dsn user - access token
+        # sdn password - http path (starting with /)
+        kw["http_path"] = dsn.user
+        kw["access_token"] = dsn.password
+        kw["server_hostname"] = dsn.host
     else:
         kw["host"] = dsn.host
         kw["port"] = dsn.port
@@ -137,12 +143,6 @@ def connect_to_uri(db_uri: str, thread_count: Optional[int] = 1) -> Database:
             kw["password"] = dsn.password
     kw = {k: v for k, v in kw.items() if v is not None}
 
-    if scheme == "databricks":
-        # dsn user - access token
-        # sdn password - http path (starting with /)
-
-        return cls(http_path=dsn.password, access_token=dsn.user, server_hostname=dsn.hostname, **kw)
-
     if issubclass(cls, ThreadedDatabase):
         return cls(thread_count=thread_count, **kw)
 

tests/common.py

@@ -15,6 +15,7 @@
 TEST_BIGQUERY_CONN_STRING: str = None
 TEST_REDSHIFT_CONN_STRING: str = None
 TEST_ORACLE_CONN_STRING: str = None
+TEST_DATABRICKS_CONN_STRING: str = None
 
 DEFAULT_N_SAMPLES = 50
 N_SAMPLES = int(os.environ.get("N_SAMPLES", DEFAULT_N_SAMPLES))
@@ -53,6 +54,7 @@ def get_git_revision_short_hash() -> str:
     db.Redshift: TEST_REDSHIFT_CONN_STRING,
     db.Oracle: TEST_ORACLE_CONN_STRING,
     db.Presto: TEST_PRESTO_CONN_STRING,
+    db.Databricks: TEST_DATABRICKS_CONN_STRING,
 }
 
 for k, v in CONN_STRINGS.items():

tests/test_database_types.py

@@ -350,6 +350,23 @@ def __iter__(self):
             "char(100)",
         ],
     },
+    db.Databricks: {
+        "int": [
+            "INT",
+            "BIGINT",
+        ],
+        "datetime": [
+            "TIMESTAMP",
+        ],
+        "float": [
+            "FLOAT",
+            "DOUBLE",
+            "DECIMAL(6, 2)",
+        ],
+        "uuid": [
+            "STRING",
+        ]
+    }
 }