Merge pull request #313 from datafold/nov24

erezsh · web-flow · commit 779892d4b1f0 · 2022-11-24T13:16:28.000-03:00
Nov24 - Small fixes to tests
diff --git a/data_diff/sqeleton/databases/base.py b/data_diff/sqeleton/databases/base.py
@@ -121,6 +121,7 @@ def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable:
 
 class BaseDialect(AbstractDialect):
     SUPPORTS_PRIMARY_KEY = False
+    SUPPORTS_INDEXES = False
     TYPE_CLASSES: Dict[str, type] = {}
 
     def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None):
diff --git a/data_diff/sqeleton/databases/databricks.py b/data_diff/sqeleton/databases/databricks.py
@@ -13,7 +13,7 @@
     DbPath,
     ColType,
     UnknownColType,
-    Boolean
+    Boolean,
 )
 from ..abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue
 from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, BaseDialect, ThreadedDatabase, import_helper, parse_table_name
diff --git a/data_diff/sqeleton/databases/duckdb.py b/data_diff/sqeleton/databases/duckdb.py
@@ -58,6 +58,7 @@ class Dialect(BaseDialect, Mixin_Schema):
     name = "DuckDB"
     ROUNDS_ON_PREC_LOSS = False
     SUPPORTS_PRIMARY_KEY = True
+    SUPPORTS_INDEXES = True
 
     TYPE_CLASSES = {
         # Timestamps
diff --git a/data_diff/sqeleton/databases/mysql.py b/data_diff/sqeleton/databases/mysql.py
@@ -51,6 +51,7 @@ class Dialect(BaseDialect, Mixin_Schema):
     name = "MySQL"
     ROUNDS_ON_PREC_LOSS = True
     SUPPORTS_PRIMARY_KEY = True
+    SUPPORTS_INDEXES = True
     TYPE_CLASSES = {
         # Dates
         "datetime": Datetime,
diff --git a/data_diff/sqeleton/databases/oracle.py b/data_diff/sqeleton/databases/oracle.py
@@ -59,21 +59,23 @@ def normalize_number(self, value: str, coltype: FractionalType) -> str:
             format_str += "0." + "9" * (coltype.precision - 1) + "0"
         return f"to_char({value}, '{format_str}')"
 
+
 class Mixin_Schema(AbstractMixin_Schema):
     def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable:
         return (
-            table('ALL_TABLES')
+            table("ALL_TABLES")
             .where(
                 this.OWNER == table_schema,
                 this.TABLE_NAME.like(like) if like is not None else SKIP,
             )
-            .select(table_name = this.TABLE_NAME)
+            .select(table_name=this.TABLE_NAME)
         )
 
 
 class Dialect(BaseDialect, Mixin_Schema):
     name = "Oracle"
     SUPPORTS_PRIMARY_KEY = True
+    SUPPORTS_INDEXES = True
     TYPE_CLASSES: Dict[str, type] = {
         "NUMBER": Decimal,
         "FLOAT": Float,
diff --git a/data_diff/sqeleton/databases/postgresql.py b/data_diff/sqeleton/databases/postgresql.py
@@ -52,6 +52,7 @@ class PostgresqlDialect(BaseDialect, Mixin_Schema):
     name = "PostgreSQL"
     ROUNDS_ON_PREC_LOSS = True
     SUPPORTS_PRIMARY_KEY = True
+    SUPPORTS_INDEXES = True
 
     TYPE_CLASSES = {
         # Timestamps
diff --git a/data_diff/table_segment.py b/data_diff/table_segment.py
@@ -100,7 +100,9 @@ def source_table(self):
         return table(*self.table_path, schema=self._schema)
 
     def make_select(self):
-        return self.source_table.where(*self._make_key_range(), *self._make_update_range(), Code(self.where) if self.where else SKIP)
+        return self.source_table.where(
+            *self._make_key_range(), *self._make_update_range(), Code(self.where) if self.where else SKIP
+        )
 
     def get_values(self) -> list:
         "Download all the relevant values of the segment from the database"
diff --git a/tests/common.py b/tests/common.py
@@ -129,8 +129,11 @@ def str_to_checksum(str: str):
     return int(md5[half_pos:], 16)
 
 
-class TestPerDatabase(unittest.TestCase):
+class DiffTestCase(unittest.TestCase):
+    "Sets up two tables for diffing (doesn't create them)"
     db_cls = None
+    src_schema = None
+    dst_schema = None
 
     def setUp(self):
         assert self.db_cls, self.db_cls
@@ -144,12 +147,16 @@ def setUp(self):
         self.table_src_path = self.connection.parse_table_name(self.table_src_name)
         self.table_dst_path = self.connection.parse_table_name(self.table_dst_name)
 
-        self.table_src = ".".join(map(self.connection.dialect.quote, self.table_src_path))
-        self.table_dst = ".".join(map(self.connection.dialect.quote, self.table_dst_path))
-
         drop_table(self.connection, self.table_src_path)
         drop_table(self.connection, self.table_dst_path)
 
+        if self.src_schema:
+            self.src_table = table(self.table_src_path, schema=self.src_schema)
+            self.connection.query(self.src_table.create())
+        if self.dst_schema:
+            self.dst_table = table(self.table_dst_path, schema=self.dst_schema)
+            self.connection.query(self.dst_table.create())
+
         return super().setUp()
 
     def tearDown(self):
diff --git a/tests/sqeleton/test_database.py b/tests/sqeleton/test_database.py
@@ -2,7 +2,8 @@
 import unittest
 
 from ..common import str_to_checksum, TEST_MYSQL_CONN_STRING
-from ..common import str_to_checksum, test_each_database_in_list, TestPerDatabase, get_conn, random_table_suffix
+from ..common import str_to_checksum, test_each_database_in_list, DiffTestCase, get_conn, random_table_suffix
+
 # from data_diff.sqeleton import databases as db
 # from data_diff.sqeleton import connect
 
@@ -52,16 +53,16 @@ def test_bad_uris(self):
 
 
 @test_each_database
-class TestSchema(TestPerDatabase):
+class TestSchema(unittest.TestCase):
     def test_table_list(self):
-        name = self.table_src_name
-        db = self.connection
-        tbl = table(db.parse_table_name(name), schema={'id': int})
+        name = "tbl_" + random_table_suffix()
+        db = get_conn(self.db_cls)
+        tbl = table(db.parse_table_name(name), schema={"id": int})
         q = db.dialect.list_tables(db.default_schema, name)
         assert not db.query(q)
 
         db.query(tbl.create())
-        self.assertEqual( db.query(q, List[str] ), [name])
+        self.assertEqual(db.query(q, List[str]), [name])
 
-        db.query( tbl.drop() )
+        db.query(tbl.drop())
         assert not db.query(q)
diff --git a/tests/sqeleton/test_query.py b/tests/sqeleton/test_query.py
@@ -74,8 +74,7 @@ def test_basic(self):
         assert c.compile(t) == "SELECT * FROM point WHERE (x = 1) AND (y = 2)"
 
         t = table("person").where(this.name == "Albert")
-        self.assertEqual( c.compile(t), "SELECT * FROM person WHERE (name = 'Albert')" )
-
+        self.assertEqual(c.compile(t), "SELECT * FROM person WHERE (name = 'Albert')")
 
     def test_outerjoin(self):
         c = Compiler(MockDatabase())
diff --git a/tests/sqeleton/test_sql.py b/tests/sqeleton/test_sql.py
@@ -73,7 +73,9 @@ def test_in(self):
         expected_sql = "SELECT name FROM `marine_mammals`.`walrus` WHERE (id IN (1, 2, 3))"
         self.assertEqual(
             expected_sql,
-            self.compiler.compile(Select(table("marine_mammals", "walrus"), [Code("name")], [In(Code("id"), [1, 2, 3])])),
+            self.compiler.compile(
+                Select(table("marine_mammals", "walrus"), [Code("name")], [In(Code("id"), [1, 2, 3])])
+            ),
         )
 
     def test_count(self):
@@ -87,7 +89,9 @@ def test_count_with_column(self):
         expected_sql = "SELECT count(id) FROM `marine_mammals`.`walrus` WHERE (id IN (1, 2, 3))"
         self.assertEqual(
             expected_sql,
-            self.compiler.compile(Select(table("marine_mammals", "walrus"), [Count(Code("id"))], [In(Code("id"), [1, 2, 3])])),
+            self.compiler.compile(
+                Select(table("marine_mammals", "walrus"), [Count(Code("id"))], [In(Code("id"), [1, 2, 3])])
+            ),
         )
 
     def test_explain(self):
diff --git a/tests/test_database_types.py b/tests/test_database_types.py
@@ -592,38 +592,9 @@ def _insert_to_table(conn, table_path, values, type):
     conn.query(commit)
 
 
-def _create_indexes(conn, table):
-    # It is unfortunate that Presto doesn't support creating indexes...
-    # Technically we could create it in the backing Postgres behind the scenes.
-    if isinstance(
-        conn, (db.Snowflake, db.Redshift, db.Presto, db.BigQuery, db.Databricks, db.Trino, db.Clickhouse, db.Vertica)
-    ):
-        return
-
-    try:
-        if_not_exists = "IF NOT EXISTS" if not isinstance(conn, (db.MySQL, db.Oracle)) else ""
-        quote = conn.dialect.quote
-        conn.query(
-            f"CREATE INDEX {if_not_exists} xa_{table[1:-1]} ON {table} ({quote('id')}, {quote('col')})",
-            None,
-        )
-        conn.query(
-            f"CREATE INDEX {if_not_exists} xb_{table[1:-1]} ON {table} ({quote('id')})",
-            None,
-        )
-    except Exception as err:
-        if "Duplicate key name" in str(err):  #  mysql
-            pass
-        elif "such column list already indexed" in str(err):  #  oracle
-            pass
-        elif "name is already used" in str(err):  #  oracle
-            pass
-        else:
-            raise (err)
-
-
 def _create_table_with_indexes(conn, table_path, type_):
-    table_name = ".".join(map(conn.dialect.quote, table_path))
+    quote = conn.dialect.quote
+    table_name = ".".join(map(quote, table_path))
 
     tbl = table(
         table_path,
@@ -633,16 +604,16 @@ def _create_table_with_indexes(conn, table_path, type_):
         },
     )
 
-    if isinstance(conn, db.Oracle):
-        already_exists = conn.query(f"SELECT COUNT(*) from tab where tname='{table_name.upper()}'", int) > 0
-        if not already_exists:
-            conn.query(tbl.create())
-    elif isinstance(conn, db.Clickhouse):
+    if isinstance(conn, db.Clickhouse):
         conn.query(f"CREATE TABLE {table_name}(id int, col {type_}) engine = Memory;", None)
     else:
-        conn.query(tbl.create(if_not_exists=True))
+        conn.query(tbl.create())
+
+    if conn.dialect.SUPPORTS_INDEXES:
+        (index_id,) = table_path
+        conn.query(f"CREATE INDEX xa_{index_id} ON {table_name} ({quote('id')}, {quote('col')})")
+        conn.query(f"CREATE INDEX xb_{index_id} ON {table_name} ({quote('id')})")
 
-    _create_indexes(conn, table_name)
     conn.query(commit)
 
 
diff --git a/tests/test_diff_tables.py b/tests/test_diff_tables.py
diff --git a/tests/test_joindiff.py b/tests/test_joindiff.py

Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@`
`13`	`13`	`DbPath,`
`14`	`14`	`ColType,`
`15`	`15`	`UnknownColType,`
`16`		`- Boolean`
	`16`	`+ Boolean,`
`17`	`17`	`)`
`18`	`18`	`from ..abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue`
`19`	`19`	`from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, BaseDialect, ThreadedDatabase, import_helper, parse_table_name`