Move table name parsing to dialects, where they semantically belong

Author: Sergey Vasilyev

data_diff/__init__.py

@@ -1,6 +1,7 @@
 from typing import Sequence, Tuple, Iterator, Optional, Union
 
 from data_diff.abcs.database_types import DbTime, DbPath
+from data_diff.databases import Database
 from data_diff.tracking import disable_tracking
 from data_diff.databases._connect import connect
 from data_diff.diff_tables import Algorithm
@@ -31,10 +32,10 @@ def connect_to_table(
     if isinstance(key_columns, str):
         key_columns = (key_columns,)
 
-    db = connect(db_info, thread_count=thread_count)
+    db: Database = connect(db_info, thread_count=thread_count)
 
     if isinstance(table_name, str):
-        table_name = db.parse_table_name(table_name)
+        table_name = db.dialect.parse_table_name(table_name)
 
     return TableSegment(db, table_name, key_columns, **kwargs)
 
@@ -161,7 +162,8 @@ def diff_tables(
         )
     elif algorithm == Algorithm.JOINDIFF:
         if isinstance(materialize_to_table, str):
-            materialize_to_table = table1.database.parse_table_name(eval_name_template(materialize_to_table))
+            table_name = eval_name_template(materialize_to_table)
+            materialize_to_table = table1.database.dialect.parse_table_name(table_name)
         differ = JoinDiffer(
             threaded=threaded,
             max_threadpool_size=max_threadpool_size,

data_diff/__main__.py

@@ -6,12 +6,13 @@
 import json
 import logging
 from itertools import islice
-from typing import Dict, Optional
+from typing import Dict, Optional, Tuple
 
 import rich
 from rich.logging import RichHandler
 import click
 
+from data_diff import Database
 from data_diff.schema import create_schema
 from data_diff.queries.api import current_timestamp
 
@@ -425,7 +426,7 @@ def _data_diff(
             logging.error(f"Error while parsing age expression: {e}")
             return
 
-    dbs = db1, db2
+    dbs: Tuple[Database, Database] = db1, db2
 
     if interactive:
         for db in dbs:
@@ -444,7 +445,7 @@ def _data_diff(
             materialize_all_rows=materialize_all_rows,
             table_write_limit=table_write_limit,
             materialize_to_table=materialize_to_table
-            and db1.parse_table_name(eval_name_template(materialize_to_table)),
+            and db1.dialect.parse_table_name(eval_name_template(materialize_to_table)),
         )
     else:
         assert algorithm == Algorithm.HASHDIFF
@@ -456,7 +457,7 @@ def _data_diff(
         )
 
     table_names = table1, table2
-    table_paths = [db.parse_table_name(t) for db, t in safezip(dbs, table_names)]
+    table_paths = [db.dialect.parse_table_name(t) for db, t in safezip(dbs, table_names)]
 
     schemas = list(differ._thread_map(_get_schema, safezip(dbs, table_paths)))
     schema1, schema2 = schemas = [

data_diff/abcs/database_types.py

@@ -176,6 +176,10 @@ class UnknownColType(ColType):
 class AbstractDialect(ABC):
     """Dialect-dependent query expressions"""
 
+    @abstractmethod
+    def parse_table_name(self, name: str) -> DbPath:
+        "Parse the given table name into a DbPath"
+
     @property
     @abstractmethod
     def name(self) -> str:
@@ -319,10 +323,6 @@ def _process_table_schema(
 
         """
 
-    @abstractmethod
-    def parse_table_name(self, name: str) -> DbPath:
-        "Parse the given table name into a DbPath"
-
     @abstractmethod
     def close(self):
         "Close connection(s) to the database instance. Querying will stop functioning."

data_diff/databases/base.py

@@ -156,6 +156,9 @@ class BaseDialect(AbstractDialect):
 
     PLACEHOLDER_TABLE = None  # Used for Oracle
 
+    def parse_table_name(self, name: str) -> DbPath:
+        return parse_table_name(name)
+
     def offset_limit(
         self, offset: Optional[int] = None, limit: Optional[int] = None, has_order_by: Optional[bool] = None
     ) -> str:
@@ -518,9 +521,6 @@ def _normalize_table_path(self, path: DbPath) -> DbPath:
 
         raise ValueError(f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected form: schema.table")
 
-    def parse_table_name(self, name: str) -> DbPath:
-        return parse_table_name(name)
-
     def _query_cursor(self, c, sql_code: str) -> QueryResult:
         assert isinstance(sql_code, str), sql_code
         try:

data_diff/databases/bigquery.py

@@ -212,6 +212,10 @@ def to_comparable(self, value: str, coltype: ColType) -> str:
     def set_timezone_to_utc(self) -> str:
         raise NotImplementedError()
 
+    def parse_table_name(self, name: str) -> DbPath:
+        path = parse_table_name(name)
+        return tuple(i for i in path if i is not None)
+
 
 class BigQuery(Database):
     CONNECT_URI_HELP = "bigquery://<project>/<dataset>"
@@ -288,10 +292,6 @@ def _normalize_table_path(self, path: DbPath) -> DbPath:
                 f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected form: [project.]schema.table"
             )
 
-    def parse_table_name(self, name: str) -> DbPath:
-        path = parse_table_name(name)
-        return tuple(i for i in self._normalize_table_path(path) if i is not None)
-
     @property
     def is_autocommit(self) -> bool:
         return True

data_diff/databases/databricks.py

@@ -94,6 +94,10 @@ def _convert_db_precision_to_digits(self, p: int) -> int:
     def set_timezone_to_utc(self) -> str:
         return "SET TIME ZONE 'UTC'"
 
+    def parse_table_name(self, name: str) -> DbPath:
+        path = parse_table_name(name)
+        return tuple(i for i in path if i is not None)
+
 
 class Databricks(ThreadedDatabase):
     dialect = Dialect()
@@ -178,10 +182,6 @@ def _process_table_schema(
         self._refine_coltypes(path, col_dict, where)
         return col_dict
 
-    def parse_table_name(self, name: str) -> DbPath:
-        path = parse_table_name(name)
-        return tuple(i for i in self._normalize_table_path(path) if i is not None)
-
     @property
     def is_autocommit(self) -> bool:
         return True

data_diff/queries/compiler.py

@@ -79,7 +79,8 @@ def new_unique_name(self, prefix="tmp"):
 
     def new_unique_table_name(self, prefix="tmp") -> DbPath:
         self._counter[0] += 1
-        return self.database.parse_table_name(f"{prefix}{self._counter[0]}_{'%x'%random.randrange(2**32)}")
+        table_name = f"{prefix}{self._counter[0]}_{'%x'%random.randrange(2**32)}"
+        return self.database.dialect.parse_table_name(table_name)
 
     def add_table_context(self, *tables: Sequence, **kw) -> Self:
         return self.replace(_table_context=self._table_context + list(tables), **kw)