Tests refactor: Convert more code to use query-builder

Author: erezsh

data_diff/databases/base.py

@@ -251,7 +251,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(16), list)
  if not samples_by_row:
  raise ValueError(f"Table {table_path} is empty.")

data_diff/queries/ast_classes.py

@@ -1,6 +1,7 @@
 from dataclasses import field
 from datetime import datetime
 from typing import Any, Generator, List, Optional, Sequence, Tuple, Union
+from uuid import UUID
 
 from runtype import dataclass
 
@@ -298,22 +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_rows(self, rows):
+ def truncate(self):
+ return TruncateTable(self)
+
+ def insert_rows(self, rows, *, columns=None):
  rows = list(rows)
- return InsertToTable(self, ConstantTable(rows))
+ return InsertToTable(self, ConstantTable(rows), columns=columns)
 
- def insert_row(self, *values):
- return InsertToTable(self, ConstantTable([values]))
+ 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
@@ -598,17 +606,21 @@ def compile(self, c: Compiler) -> str:
 
 @dataclass
 class ConstantTable(ExprNode):
- rows: List[Tuple]
+ rows: Sequence[Sequence]
 
  def compile(self, c: Compiler) -> str:
  raise NotImplementedError()
 
  def _value(self, v):
- if isinstance(v, str):
+ if v is None:
+ return "NULL"
+ elif isinstance(v, str):
  return f"'{v}'"
  elif isinstance(v, datetime):
  return f"timestamp '{v}'"
- return str(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)
@@ -633,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})"
 
 
@@ -651,19 +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:
  if isinstance(self.expr, ConstantTable):
  expr = self.expr.compile_for_insert(c)
  else:
  expr = c.compile(self.expr)
 
- return f"INSERT INTO {c.compile(self.path)} {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