Refactor tests oct2022

.github/workflows/ci.yml

@@ -46,7 +46,6 @@ jobs:
         env:
             DATADIFF_SNOWFLAKE_URI: '${{ secrets.DATADIFF_SNOWFLAKE_URI }}'
             DATADIFF_PRESTO_URI: '${{ secrets.DATADIFF_PRESTO_URI }}'
-            DATADIFF_TRINO_URI: '${{ secrets.DATADIFF_TRINO_URI }}'
             DATADIFF_CLICKHOUSE_URI: 'clickhouse://clickhouse:Password1@localhost:9000/clickhouse'
             DATADIFF_VERTICA_URI: 'vertica://vertica:Password1@localhost:5433/vertica'
         run: |

data_diff/databases/base.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 import math
 import sys
 import logging
@@ -120,6 +121,10 @@ def query(self, sql_ast: Union[Expr, Generator], res_type: type = list):
         compiler = Compiler(self)
         if isinstance(sql_ast, Generator):
             sql_code = ThreadLocalInterpreter(compiler, sql_ast)
+        elif isinstance(sql_ast, list):
+            for i in sql_ast[:-1]:
+                self.query(i)
+            return self.query(sql_ast[-1], res_type)
         else:
             sql_code = compiler.compile(sql_ast)
             if sql_code is SKIP:
@@ -249,7 +254,7 @@ def _refine_coltypes(self, table_path: DbPath, col_dict: Dict[str, ColType], whe
         if not text_columns:
             return
 
-        fields = [self.normalize_uuid(c, String_UUID()) for c in text_columns]
+        fields = [self.normalize_uuid(self.quote(c), String_UUID()) for c in text_columns]
         samples_by_row = self.query(table(*table_path).select(*fields).where(where or SKIP).limit(sample_size), list)
         if not samples_by_row:
             raise ValueError(f"Table {table_path} is empty.")
@@ -329,6 +334,7 @@ def type_repr(self, t) -> str:
             str: "VARCHAR",
             bool: "BOOLEAN",
             float: "FLOAT",
+            datetime: "TIMESTAMP",
         }[t]
 
     def _query_cursor(self, c, sql_code: str):

data_diff/databases/bigquery.py

@@ -109,8 +109,6 @@ def is_autocommit(self) -> bool:
 
     def type_repr(self, t) -> str:
         try:
-            return {
-                str: "STRING",
-            }[t]
+            return {str: "STRING", float: "FLOAT64"}[t]
         except KeyError:
             return super().type_repr(t)

data_diff/databases/clickhouse.py

@@ -150,3 +150,7 @@ def normalize_number(self, value: str, coltype: FractionalType) -> str:
             )
         """
         return value
+
+    @property
+    def is_autocommit(self) -> bool:
+        return True

data_diff/databases/databricks.py

@@ -151,3 +151,7 @@ def parse_table_name(self, name: str) -> DbPath:
 
     def close(self):
         self._conn.close()
+
+    @property
+    def is_autocommit(self) -> bool:
+        return True

data_diff/databases/presto.py

@@ -153,3 +153,9 @@ def is_autocommit(self) -> bool:
 
     def explain_as_text(self, query: str) -> str:
         return f"EXPLAIN (FORMAT TEXT) {query}"
+
+    def type_repr(self, t) -> str:
+        try:
+            return {float: "REAL"}[t]
+        except KeyError:
+            return super().type_repr(t)

data_diff/joindiff_tables.py

@@ -10,8 +10,8 @@
 
 from runtype import dataclass
 
-from data_diff.databases.database_types import DbPath, NumericType
-from data_diff.databases.base import QueryError
+from .databases.database_types import DbPath, NumericType
+from .query_utils import append_to_table, drop_table
 
 
 from .utils import safezip
@@ -48,7 +48,7 @@ def sample(table_expr):
     return table_expr.order_by(Random()).limit(10)
 
 
-def create_temp_table(c: Compiler, path: TablePath, expr: Expr):
+def create_temp_table(c: Compiler, path: TablePath, expr: Expr) -> str:
     db = c.database
     if isinstance(db, BigQuery):
         return f"create table {c.compile(path)} OPTIONS(expiration_timestamp=TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL 1 DAY)) as {c.compile(expr)}"
@@ -60,42 +60,6 @@ def create_temp_table(c: Compiler, path: TablePath, expr: Expr):
         return f"create temporary table {c.compile(path)} as {c.compile(expr)}"
 
 
-def drop_table_oracle(name: DbPath):
-    t = table(name)
-    # Experience shows double drop is necessary
-    with suppress(QueryError):
-        yield t.drop()
-        yield t.drop()
-    yield commit
-
-
-def drop_table(name: DbPath):
-    t = table(name)
-    yield t.drop(if_exists=True)
-    yield commit
-
-
-def append_to_table_oracle(path: DbPath, expr: Expr):
-    """See append_to_table"""
-    assert expr.schema, expr
-    t = table(path, schema=expr.schema)
-    with suppress(QueryError):
-        yield t.create()  # uses expr.schema
-        yield commit
-    yield t.insert_expr(expr)
-    yield commit
-
-
-def append_to_table(path: DbPath, expr: Expr):
-    """Append to table"""
-    assert expr.schema, expr
-    t = table(path, schema=expr.schema)
-    yield t.create(if_not_exists=True)  # uses expr.schema
-    yield commit
-    yield t.insert_expr(expr)
-    yield commit
-
-
 def bool_to_int(x):
     return if_(x, 1, 0)
 
@@ -170,10 +134,7 @@ def _diff_tables(self, table1: TableSegment, table2: TableSegment) -> DiffResult
 
         bg_funcs = [partial(self._test_duplicate_keys, table1, table2)] if self.validate_unique_key else []
         if self.materialize_to_table:
-            if isinstance(db, Oracle):
-                db.query(drop_table_oracle(self.materialize_to_table))
-            else:
-                db.query(drop_table(self.materialize_to_table))
+            drop_table(db, self.materialize_to_table)
 
         with self._run_in_background(*bg_funcs):
 
@@ -348,6 +309,5 @@ def exclusive_rows(expr):
     def _materialize_diff(self, db, diff_rows, segment_index=None):
         assert self.materialize_to_table
 
-        f = append_to_table_oracle if isinstance(db, Oracle) else append_to_table
-        db.query(f(self.materialize_to_table, diff_rows.limit(self.write_limit)))
+        append_to_table(db, self.materialize_to_table, diff_rows.limit(self.write_limit))
         logger.info("Materialized diff to table '%s'.", ".".join(self.materialize_to_table))

data_diff/queries/api.py

@@ -1,4 +1,6 @@
 from typing import Optional
+
+from data_diff.utils import CaseAwareMapping, CaseSensitiveDict
 from .ast_classes import *
 from .base import args_as_tuple
 
@@ -30,11 +32,14 @@ def cte(expr: Expr, *, name: Optional[str] = None, params: Sequence[str] = None)
     return Cte(expr, name, params)
 
 
-def table(*path: str, schema: Schema = None) -> TablePath:
+def table(*path: str, schema: Union[dict, CaseAwareMapping] = None) -> TablePath:
     if len(path) == 1 and isinstance(path[0], tuple):
         (path,) = path
     if not all(isinstance(i, str) for i in path):
         raise TypeError(f"All elements of table path must be of type 'str'. Got: {path}")
+    if schema and not isinstance(schema, CaseAwareMapping):
+        assert isinstance(schema, dict)
+        schema = CaseSensitiveDict(schema)
     return TablePath(path, schema)
 
 

data_diff/queries/ast_classes.py

@@ -1,6 +1,7 @@
 from dataclasses import field
 from datetime import datetime
-from typing import Any, Generator, Optional, Sequence, Tuple, Union
+from typing import Any, Generator, List, Optional, Sequence, Tuple, Union
+from uuid import UUID
 
 from runtype import dataclass
 
@@ -298,18 +299,29 @@ class TablePath(ExprNode, ITable):
     path: DbPath
     schema: Optional[Schema] = field(default=None, repr=False)
 
-    def create(self, if_not_exists=False):
-        if not self.schema:
-            raise ValueError("Schema must have a value to create table")
-        return CreateTable(self, if_not_exists=if_not_exists)
+    def create(self, source_table: ITable = None, *, if_not_exists=False):
+        if source_table is None and not self.schema:
+            raise ValueError("Either schema or source table needed to create table")
+        if isinstance(source_table, TablePath):
+            source_table = source_table.select()
+        return CreateTable(self, source_table, if_not_exists=if_not_exists)
 
     def drop(self, if_exists=False):
         return DropTable(self, if_exists=if_exists)
 
-    def insert_values(self, rows):
-        raise NotImplementedError()
+    def truncate(self):
+        return TruncateTable(self)
+
+    def insert_rows(self, rows, *, columns=None):
+        rows = list(rows)
+        return InsertToTable(self, ConstantTable(rows), columns=columns)
+
+    def insert_row(self, *values, columns=None):
+        return InsertToTable(self, ConstantTable([values]), columns=columns)
 
     def insert_expr(self, expr: Expr):
+        if isinstance(expr, TablePath):
+            expr = expr.select()
         return InsertToTable(self, expr)
 
     @property
@@ -592,6 +604,29 @@ def compile(self, c: Compiler) -> str:
         return c.database.random()
 
 
+@dataclass
+class ConstantTable(ExprNode):
+    rows: Sequence[Sequence]
+
+    def compile(self, c: Compiler) -> str:
+        raise NotImplementedError()
+
+    def _value(self, v):
+        if v is None:
+            return "NULL"
+        elif isinstance(v, str):
+            return f"'{v}'"
+        elif isinstance(v, datetime):
+            return f"timestamp '{v}'"
+        elif isinstance(v, UUID):
+            return f"'{v}'"
+        return repr(v)
+
+    def compile_for_insert(self, c: Compiler):
+        values = ", ".join("(%s)" % ", ".join(self._value(v) for v in row) for row in self.rows)
+        return f"VALUES {values}"
+
+
 @dataclass
 class Explain(ExprNode):
     select: Select
@@ -610,11 +645,15 @@ class Statement(Compilable):
 @dataclass
 class CreateTable(Statement):
     path: TablePath
+    source_table: Expr = None
     if_not_exists: bool = False
 
     def compile(self, c: Compiler) -> str:
-        schema = ", ".join(f"{k} {c.database.type_repr(v)}" for k, v in self.path.schema.items())
         ne = "IF NOT EXISTS " if self.if_not_exists else ""
+        if self.source_table:
+            return f"CREATE TABLE {ne}{c.compile(self.path)} AS {c.compile(self.source_table)}"
+
+        schema = ", ".join(f"{c.database.quote(k)} {c.database.type_repr(v)}" for k, v in self.path.schema.items())
         return f"CREATE TABLE {ne}{c.compile(self.path)}({schema})"
 
 
@@ -628,14 +667,30 @@ def compile(self, c: Compiler) -> str:
         return f"DROP TABLE {ie}{c.compile(self.path)}"
 
 
+@dataclass
+class TruncateTable(Statement):
+    path: TablePath
+
+    def compile(self, c: Compiler) -> str:
+        return f"TRUNCATE TABLE {c.compile(self.path)}"
+
+
 @dataclass
 class InsertToTable(Statement):
     # TODO Support insert for only some columns
     path: TablePath
     expr: Expr
+    columns: List[str] = None
 
     def compile(self, c: Compiler) -> str:
-        return f"INSERT INTO {c.compile(self.path)} {c.compile(self.expr)}"
+        if isinstance(self.expr, ConstantTable):
+            expr = self.expr.compile_for_insert(c)
+        else:
+            expr = c.compile(self.expr)
+
+        columns = f"(%s)" % ", ".join(map(c.quote, self.columns)) if self.columns is not None else ""
+
+        return f"INSERT INTO {c.compile(self.path)}{columns} {expr}"
 
 
 @dataclass

data_diff/query_utils.py

@@ -0,0 +1,55 @@
+"Module for query utilities that didn't make it into the query-builder (yet)"
+
+from contextlib import suppress
+
+from data_diff.databases.database_types import DbPath
+from data_diff.databases.base import QueryError
+
+from .databases import Oracle
+from .queries import table, commit, Expr
+
+def _drop_table_oracle(name: DbPath):
+    t = table(name)
+    # Experience shows double drop is necessary
+    with suppress(QueryError):
+        yield t.drop()
+        yield t.drop()
+    yield commit
+
+
+def _drop_table(name: DbPath):
+    t = table(name)
+    yield t.drop(if_exists=True)
+    yield commit
+
+
+def drop_table(db, tbl):
+    if isinstance(db, Oracle):
+        db.query(_drop_table_oracle(tbl))
+    else:
+        db.query(_drop_table(tbl))
+
+
+def _append_to_table_oracle(path: DbPath, expr: Expr):
+    """See append_to_table"""
+    assert expr.schema, expr
+    t = table(path, schema=expr.schema)
+    with suppress(QueryError):
+        yield t.create()  # uses expr.schema
+        yield commit
+    yield t.insert_expr(expr)
+    yield commit
+
+
+def _append_to_table(path: DbPath, expr: Expr):
+    """Append to table"""
+    assert expr.schema, expr
+    t = table(path, schema=expr.schema)
+    yield t.create(if_not_exists=True)  # uses expr.schema
+    yield commit
+    yield t.insert_expr(expr)
+    yield commit
+
+def append_to_table(db, path, expr):
+    f = _append_to_table_oracle if isinstance(db, Oracle) else _append_to_table
+    db.query(f(path, expr))