From 47a92f7c1ca3bd21e9c92247fac0803b3d4c89b8 Mon Sep 17 00:00:00 2001 From: vmatt Date: Fri, 13 Feb 2026 11:50:11 +0100 Subject: [PATCH 01/11] add sqeleton to reladiff --- .gitignore | 1 + pyproject.toml | 1 - reladiff/sqeleton/__init__.py | 4 + reladiff/sqeleton/__main__.py | 28 + reladiff/sqeleton/abcs/__init__.py | 15 + reladiff/sqeleton/abcs/compiler.py | 12 + reladiff/sqeleton/abcs/database_types.py | 425 +++++++ reladiff/sqeleton/abcs/mixins.py | 158 +++ reladiff/sqeleton/bound_exprs.py | 81 ++ reladiff/sqeleton/conn_editor.py | 292 +++++ reladiff/sqeleton/databases/__init__.py | 19 + reladiff/sqeleton/databases/_connect.py | 276 +++++ reladiff/sqeleton/databases/base.py | 659 +++++++++++ reladiff/sqeleton/databases/bigquery.py | 213 ++++ reladiff/sqeleton/databases/clickhouse.py | 197 ++++ reladiff/sqeleton/databases/databricks.py | 204 ++++ reladiff/sqeleton/databases/dremio.py | 207 ++++ reladiff/sqeleton/databases/duckdb.py | 193 ++++ reladiff/sqeleton/databases/mssql.py | 25 + reladiff/sqeleton/databases/mysql.py | 150 +++ reladiff/sqeleton/databases/oracle.py | 220 ++++ reladiff/sqeleton/databases/postgresql.py | 186 ++++ reladiff/sqeleton/databases/presto.py | 208 ++++ reladiff/sqeleton/databases/redshift.py | 128 +++ reladiff/sqeleton/databases/snowflake.py | 228 ++++ reladiff/sqeleton/databases/trino.py | 71 ++ reladiff/sqeleton/databases/vertica.py | 181 +++ reladiff/sqeleton/queries/__init__.py | 25 + reladiff/sqeleton/queries/api.py | 213 ++++ reladiff/sqeleton/queries/ast_classes.py | 1223 +++++++++++++++++++++ reladiff/sqeleton/queries/base.py | 27 + reladiff/sqeleton/queries/compiler.py | 535 +++++++++ reladiff/sqeleton/queries/extras.py | 71 ++ reladiff/sqeleton/query_utils.py | 54 + reladiff/sqeleton/repl.py | 284 +++++ reladiff/sqeleton/schema.py | 73 ++ reladiff/sqeleton/utils.py | 345 ++++++ sqeleton-project/pyproject.toml | 114 ++ 38 files changed, 7345 insertions(+), 1 deletion(-) create mode 100644 reladiff/sqeleton/__init__.py create mode 100644 reladiff/sqeleton/__main__.py create mode 100644 reladiff/sqeleton/abcs/__init__.py create mode 100644 reladiff/sqeleton/abcs/compiler.py create mode 100644 reladiff/sqeleton/abcs/database_types.py create mode 100644 reladiff/sqeleton/abcs/mixins.py create mode 100644 reladiff/sqeleton/bound_exprs.py create mode 100644 reladiff/sqeleton/conn_editor.py create mode 100644 reladiff/sqeleton/databases/__init__.py create mode 100644 reladiff/sqeleton/databases/_connect.py create mode 100644 reladiff/sqeleton/databases/base.py create mode 100644 reladiff/sqeleton/databases/bigquery.py create mode 100644 reladiff/sqeleton/databases/clickhouse.py create mode 100644 reladiff/sqeleton/databases/databricks.py create mode 100644 reladiff/sqeleton/databases/dremio.py create mode 100644 reladiff/sqeleton/databases/duckdb.py create mode 100644 reladiff/sqeleton/databases/mssql.py create mode 100644 reladiff/sqeleton/databases/mysql.py create mode 100644 reladiff/sqeleton/databases/oracle.py create mode 100644 reladiff/sqeleton/databases/postgresql.py create mode 100644 reladiff/sqeleton/databases/presto.py create mode 100644 reladiff/sqeleton/databases/redshift.py create mode 100644 reladiff/sqeleton/databases/snowflake.py create mode 100644 reladiff/sqeleton/databases/trino.py create mode 100644 reladiff/sqeleton/databases/vertica.py create mode 100644 reladiff/sqeleton/queries/__init__.py create mode 100644 reladiff/sqeleton/queries/api.py create mode 100644 reladiff/sqeleton/queries/ast_classes.py create mode 100644 reladiff/sqeleton/queries/base.py create mode 100644 reladiff/sqeleton/queries/compiler.py create mode 100644 reladiff/sqeleton/queries/extras.py create mode 100644 reladiff/sqeleton/query_utils.py create mode 100644 reladiff/sqeleton/repl.py create mode 100644 reladiff/sqeleton/schema.py create mode 100644 reladiff/sqeleton/utils.py create mode 100644 sqeleton-project/pyproject.toml diff --git a/.gitignore b/.gitignore index 681f689e..b772a1de 100644 --- a/.gitignore +++ b/.gitignore @@ -144,3 +144,4 @@ benchmark_*.png # VSCode .vscode +.cecli* diff --git a/pyproject.toml b/pyproject.toml index 951f6eef..40fc72f0 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,6 @@ dsnparse = "*" click = ">=8.1" rich = "*" toml = ">=0.10.2" -sqeleton = "^0.1.7" mysql-connector-python = {version=">=8.0.29", optional=true} psycopg2-binary = {version="*", optional=true} snowflake-connector-python = {version=">=2.7.2", optional=true} diff --git a/reladiff/sqeleton/__init__.py b/reladiff/sqeleton/__init__.py new file mode 100644 index 00000000..40f43bb8 --- /dev/null +++ b/reladiff/sqeleton/__init__.py @@ -0,0 +1,4 @@ +from .databases import connect +from .queries import table, this, SKIP, code, commit + +__version__ = "0.1.8" diff --git a/reladiff/sqeleton/__main__.py b/reladiff/sqeleton/__main__.py new file mode 100644 index 00000000..a159d8e7 --- /dev/null +++ b/reladiff/sqeleton/__main__.py @@ -0,0 +1,28 @@ +import click +from .repl import repl as repl_main + + +@click.group(no_args_is_help=True) +def main(): + pass + + +@main.command(no_args_is_help=True) +@click.argument("database", required=True) +def repl(database): + return repl_main(database) + + +CONN_EDITOR_HELP = """CONFIG_PATH - Path to a TOML config file of db connections, new or existing.""" + + +@main.command(no_args_is_help=True, help=CONN_EDITOR_HELP) +@click.argument("config_path", required=True) +def conn_editor(config_path): + from .conn_editor import main + + return main(config_path) + + +if __name__ == "__main__": + main() diff --git a/reladiff/sqeleton/abcs/__init__.py b/reladiff/sqeleton/abcs/__init__.py new file mode 100644 index 00000000..c56885ff --- /dev/null +++ b/reladiff/sqeleton/abcs/__init__.py @@ -0,0 +1,15 @@ +from .database_types import ( + AbstractDatabase, + AbstractDialect, + DbKey, + DbPath, + DbTime, + IKey, + ColType_UUID, + NumericType, + PrecisionType, + StringType, + Boolean, + Text, +) +from .compiler import AbstractCompiler, Compilable diff --git a/reladiff/sqeleton/abcs/compiler.py b/reladiff/sqeleton/abcs/compiler.py new file mode 100644 index 00000000..dee9f69d --- /dev/null +++ b/reladiff/sqeleton/abcs/compiler.py @@ -0,0 +1,12 @@ +from typing import Any, Dict +from abc import ABC, abstractmethod + + +class AbstractCompiler(ABC): + @abstractmethod + def compile(self, elem: Any, params: Dict[str, Any] = None) -> str: + ... + + +class Compilable(ABC): + "Marks an item as compilable. Needs to be implemented through multiple-dispatch." diff --git a/reladiff/sqeleton/abcs/database_types.py b/reladiff/sqeleton/abcs/database_types.py new file mode 100644 index 00000000..f5aefebe --- /dev/null +++ b/reladiff/sqeleton/abcs/database_types.py @@ -0,0 +1,425 @@ +import uuid +import decimal +from abc import ABC, abstractmethod +from typing import Sequence, Optional, Tuple, Union, Dict, List, TypeVar, Generic, Any +from datetime import datetime + +from runtype import dataclass + +from ..utils import ArithAlphanumeric, ArithUUID, Self, Unknown + + +DbPath = Tuple[str, ...] +DbKey = Union[int, str, bytes, ArithUUID, ArithAlphanumeric] +DbTime = datetime + + +class ColType: + supported = True + + +@dataclass +class PrecisionType(ColType): + precision: int + rounds: Union[bool, Unknown] = Unknown + + +class TemporalType(PrecisionType): + pass + + +class Timestamp(TemporalType): + pass + + +class TimestampTZ(TemporalType): + pass + + +class Datetime(TemporalType): + pass + + +class Date(TemporalType): + pass + + +@dataclass +class NumericType(ColType): + # 'precision' signifies how many fractional digits (after the dot) we want to compare + precision: int + + +class FractionalType(NumericType): + pass + + +class IKey(ABC): + "Interface for ColType, for using a column as a key in table." + + @property + @abstractmethod + def python_type(self) -> type: + "Return the equivalent Python type of the key" + + def make_value(self, value): + return self.python_type(value) + + +class Float(FractionalType, IKey): + @property + def python_type(self) -> type: + if self.precision == 0: + return int + return float + + +class Decimal(FractionalType, IKey): # Snowflake may use Decimal as a key + @property + def python_type(self) -> type: + if self.precision == 0: + return int + return decimal.Decimal + + +class Boolean(ColType, IKey): + precision = 0 + python_type = bool + +@dataclass +class StringType(ColType): + python_type = str + + +class ColType_UUID(ColType, IKey): + python_type = ArithUUID + + +class ColType_Alphanum(ColType, IKey): + python_type = ArithAlphanumeric + + +class Native_UUID(ColType_UUID): + pass + + +class String_UUID(ColType_UUID, StringType): + pass + + +class String_Alphanum(ColType_Alphanum, StringType): + @staticmethod + def test_value(value: str) -> bool: + try: + ArithAlphanumeric(value) + return True + except ValueError: + return False + + def make_value(self, value): + return self.python_type(value) + + +class String_VaryingAlphanum(String_Alphanum): + pass + + +@dataclass +class String_FixedAlphanum(String_Alphanum): + length: int + + def make_value(self, value): + if len(value) != self.length: + raise ValueError( + f"Expected alphanumeric value of length {self.length}, but got '{value}'." + ) + return self.python_type(value, max_len=self.length) + + +@dataclass +class Text(StringType): + supported = False + + +@dataclass +class Integer(NumericType, IKey): + precision: int = 0 + python_type: type = int + + def __post_init__(self): + assert self.precision == 0 + + +@dataclass +class UnknownColType(ColType): + text: str + + supported = False + + +class AbstractDialect(ABC): + """Dialect-dependent query expressions""" + + @property + @abstractmethod + def name(self) -> str: + "Name of the dialect" + + @classmethod + @abstractmethod + def load_mixins(cls, *abstract_mixins) -> Self: + "Load a list of mixins that implement the given abstract mixins" + + @property + @abstractmethod + def ROUNDS_ON_PREC_LOSS(self) -> bool: + "True if db rounds real values when losing precision, False if it truncates." + + @abstractmethod + def quote(self, s: str): + "Quote SQL name" + + @abstractmethod + def concat(self, items: List[str]) -> str: + "Provide SQL for concatenating a bunch of columns into a string" + + @abstractmethod + def is_distinct_from(self, a: str, b: str) -> str: + "Provide SQL for a comparison where NULL = NULL is true" + + @abstractmethod + def to_string(self, s: str) -> str: + # TODO rewrite using cast_to(x, str) + "Provide SQL for casting a column to string" + + @abstractmethod + def random(self) -> str: + "Provide SQL for generating a random number betweein 0..1" + + @abstractmethod + def current_timestamp(self) -> str: + "Provide SQL for returning the current timestamp, aka now" + + @abstractmethod + def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None): + "Provide SQL fragment for limit and offset inside a select" + + @abstractmethod + def explain_as_text(self, query: str) -> str: + "Provide SQL for explaining a query, returned as table(varchar)" + + @abstractmethod + def timestamp_value(self, t: datetime) -> str: + "Provide SQL for the given timestamp value" + + @abstractmethod + def uuid_value(self, t: uuid.UUID) -> str: + "Provide SQL for the given UUID value" + + @abstractmethod + def set_timezone_to_utc(self) -> str: + "Provide SQL for setting the session timezone to UTC" + + @abstractmethod + def parse_type( + self, + table_path: DbPath, + col_name: str, + type_repr: str, + datetime_precision: int = None, + numeric_precision: int = None, + numeric_scale: int = None, + ) -> ColType: + "Parse type info as returned by the database" + + +T_Dialect = TypeVar("T_Dialect", bound=AbstractDialect) + + +class AbstractDatabase(Generic[T_Dialect]): + @property + @abstractmethod + def name(self) -> str: + "The name of the database" + + @property + @abstractmethod + def dialect(self) -> T_Dialect: + "The dialect of the database. Used internally by Database, and also available publicly." + + @classmethod + @abstractmethod + def load_mixins(cls, *abstract_mixins) -> type: + "Extend the dialect with a list of mixins that implement the given abstract mixins." + + @property + @abstractmethod + def CONNECT_URI_HELP(self) -> str: + "Example URI to show the user in help and error messages" + + @property + @abstractmethod + def CONNECT_URI_PARAMS(self) -> List[str]: + "List of parameters given in the path of the URI" + + @abstractmethod + def _query(self, sql_code: Any) -> list: + "Send query to database and return result" + + @abstractmethod + def query_table_schema(self, path: DbPath) -> Dict[str, tuple]: + """Query the database for the schema of the table in 'path', and return {column: tuple, ...} + where the tuple is (table_name, col_name, type_repr, datetime_precision?, numeric_precision?, numeric_scale?) + + Use the method .process_query_table_schema() to convert the tuples into types. + + Note: This method is used instead of select_table_schema(), + because not all databases support accessing the schema using a SQL query. + """ + + @abstractmethod + def select_table_unique_columns(self, path: DbPath) -> str: + "Provide SQL for selecting the names of unique columns in the table" + + @abstractmethod + def query_table_unique_columns(self, path: DbPath) -> List[str]: + """Query the table for its unique columns for table in 'path', and return {column}""" + + @abstractmethod + def _process_table_schema( + self, + path: DbPath, + raw_schema: Dict[str, tuple], + filter_columns: Sequence[str], + where: str = None, + ): + """Process the result of query_table_schema(). + + Done in a separate step, to minimize the amount of processed columns. + Needed because processing each column may: + * throw errors and warnings + * query the database to sample values + + """ + + @abstractmethod + def process_query_table_schema( + self, + path: DbPath, + raw_schema: Dict[str, tuple], + refine: bool = True, + refine_where: Optional[str] = None, + ) -> Tuple[Dict[str, ColType], Optional[list]]: + """Process the result of query_table_schema(). + + We are doing it here, separately, because: + - parse_type() may throw an exception. Therefor, some users may want to query + each column separately. + - Refining the column types may query the database to sample values, and some + users may want to do so in separate threads. + + If refine is True, it samples the table data to refine the column types when possible. + Use refine_where to filter the rows that may be sampled. + """ + + @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." + + @property + @abstractmethod + def is_autocommit(self) -> bool: + "Return whether the database autocommits changes. When false, COMMIT statements are skipped." + + +class AbstractTable(ABC): + @abstractmethod + def select(self, *exprs, distinct=False, **named_exprs) -> "AbstractTable": + """Choose new columns, based on the old ones. (aka Projection) + + Parameters: + exprs: List of expressions to constitute the columns of the new table. + If not provided, returns all columns in source table (i.e. ``select *``) + distinct: 'select' or 'select distinct' + named_exprs: More expressions to constitute the columns of the new table, aliased to keyword name. + + """ + # XXX distinct=SKIP + + @abstractmethod + def where(self, *exprs) -> "AbstractTable": + """Filter the rows, based on the given predicates. (aka Selection)""" + + @abstractmethod + def order_by(self, *exprs) -> "AbstractTable": + """Order the rows lexicographically, according to the given expressions.""" + + @abstractmethod + def limit(self, limit: int) -> "AbstractTable": + """Stop yielding rows after the given limit. i.e. take the first 'n=limit' rows""" + + @abstractmethod + def join(self, target) -> "AbstractTable": + """Join the current table with the target table, returning a new table containing both side-by-side. + + When joining, it's recommended to use explicit tables names, instead of `this`, + in order to avoid potential name collisions. + + Example: + :: + + person = table('person') + city = table('city') + + name_and_city = ( + person + .join(city) + .on(person['city_id'] == city['id']) + .select(person['id'], city['name']) + ) + """ + + @abstractmethod + def group_by(self, *keys): + """Behaves like in SQL, except for a small change in syntax: + + A call to `.agg()` must follow every call to `.group_by()`. + + Example: + :: + + # SELECT a, sum(b) FROM tmp GROUP BY 1 + table('tmp').group_by(this.a).agg(this.b.sum()) + + # SELECT a, sum(b) FROM a GROUP BY 1 HAVING (b > 10) + (table('tmp') + .group_by(this.a) + .agg(this.b.sum()) + .having(this.b > 10) + ) + + """ + + @abstractmethod + def count(self) -> int: + """SELECT count() FROM self""" + + @abstractmethod + def union(self, other: "AbstractTable"): + """SELECT * FROM self UNION other""" + + @abstractmethod + def union_all(self, other: "AbstractTable"): + """SELECT * FROM self UNION ALL other""" + + @abstractmethod + def minus(self, other: "AbstractTable"): + """SELECT * FROM self EXCEPT other""" + + @abstractmethod + def intersect(self, other: "AbstractTable"): + """SELECT * FROM self INTERSECT other""" diff --git a/reladiff/sqeleton/abcs/mixins.py b/reladiff/sqeleton/abcs/mixins.py new file mode 100644 index 00000000..15102888 --- /dev/null +++ b/reladiff/sqeleton/abcs/mixins.py @@ -0,0 +1,158 @@ +from abc import ABC, abstractmethod +from .database_types import TemporalType, FractionalType, ColType_UUID, Boolean, ColType, String_UUID +from .compiler import Compilable + + +class AbstractMixin(ABC): + "A mixin for a database dialect" + + +class AbstractMixin_NormalizeValue(AbstractMixin): + @abstractmethod + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + """Creates an SQL expression, that converts 'value' to a normalized timestamp. + + The returned expression must accept any SQL datetime/timestamp, and return a string. + + Date format: ``YYYY-MM-DD HH:mm:SS.FFFFFF`` + + Precision of dates should be rounded up/down according to coltype.rounds + """ + + @abstractmethod + def normalize_number(self, value: str, coltype: FractionalType) -> str: + """Creates an SQL expression, that converts 'value' to a normalized number. + + The returned expression must accept any SQL int/numeric/float, and return a string. + + Floats/Decimals are expected in the format + "I.P" + + Where I is the integer part of the number (as many digits as necessary), + and must be at least one digit (0). + P is the fractional digits, the amount of which is specified with + coltype.precision. Trailing zeroes may be necessary. + If P is 0, the dot is omitted. + + Note: We use 'precision' differently than most databases. For decimals, + it's the same as ``numeric_scale``, and for floats, who use binary precision, + it can be calculated as ``log10(2**numeric_precision)``. + """ + + def normalize_boolean(self, value: str, _coltype: Boolean) -> str: + """Creates an SQL expression, that converts 'value' to either '0' or '1'.""" + return self.to_string(value) + + def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str: + """Creates an SQL expression, that strips uuids of artifacts like whitespace.""" + if isinstance(coltype, String_UUID): + return f"TRIM({value})" + return self.to_string(value) + + def normalize_value_by_type(self, value: str, coltype: ColType) -> str: + """Creates an SQL expression, that converts 'value' to a normalized representation. + + The returned expression must accept any SQL value, and return a string. + + The default implementation dispatches to a method according to `coltype`: + + :: + + TemporalType -> normalize_timestamp() + FractionalType -> normalize_number() + *else* -> to_string() + + (`Integer` falls in the *else* category) + + """ + if isinstance(coltype, TemporalType): + return self.normalize_timestamp(value, coltype) + elif isinstance(coltype, FractionalType): + return self.normalize_number(value, coltype) + elif isinstance(coltype, ColType_UUID): + return self.normalize_uuid(value, coltype) + elif isinstance(coltype, Boolean): + return self.normalize_boolean(value, coltype) + return self.to_string(value) + + +class AbstractMixin_MD5(AbstractMixin): + """Methods for calculating an MD6 hash as an integer.""" + + @abstractmethod + def md5_as_int(self, s: str) -> str: + "Provide SQL for computing md5 and returning an int" + + +class AbstractMixin_Schema(AbstractMixin): + """Methods for querying the database schema + + TODO: Move AbstractDatabase.query_table_schema() and friends over here + """ + + def table_information(self) -> Compilable: + "Query to return a table of schema information about existing tables" + raise NotImplementedError() + + @abstractmethod + def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable: + """Query to select the list of tables in the schema. (query return type: table[str]) + + If 'like' is specified, the value is applied to the table name, using the 'like' operator. + """ + + +class AbstractMixin_Regex(AbstractMixin): + @abstractmethod + def test_regex(self, string: Compilable, pattern: Compilable) -> Compilable: + """Tests whether the regex pattern matches the string. Returns a bool expression.""" + + +class AbstractMixin_RandomSample(AbstractMixin): + @abstractmethod + def random_sample_n(self, tbl: str, size: int) -> str: + """Take a random sample of the given size, i.e. return 'size' amount of rows""" + + @abstractmethod + def random_sample_ratio_approx(self, tbl: str, ratio: float) -> str: + """Take a random sample of the approximate size determined by the ratio (0..1), + ratio of 0 means no rows, and 1 means all rows + + i.e. the actual mount of rows returned may vary by standard deviation. + """ + + # def random_sample_ratio(self, table: AbstractTable, ratio: float): + # """Take a random sample of the size determined by the ratio (0..1), where 0 means no rows, and 1 means all rows + # """ + + +class AbstractMixin_TimeTravel(AbstractMixin): + @abstractmethod + def time_travel( + self, + table: Compilable, + before: bool = False, + timestamp: Compilable = None, + offset: Compilable = None, + statement: Compilable = None, + ) -> Compilable: + """Selects historical data from a table + + Parameters: + table - The name of the table whose history we're querying + timestamp - A constant timestamp + offset - the time 'offset' seconds before now + statement - identifier for statement, e.g. query ID + + Must specify exactly one of `timestamp`, `offset` or `statement`. + """ + + +class AbstractMixin_OptimizerHints(AbstractMixin): + @abstractmethod + def optimizer_hints(self, optimizer_hints: str) -> str: + """Creates a compatible optimizer_hints string + + Parameters: + optimizer_hints - string of optimizer hints + """ diff --git a/reladiff/sqeleton/bound_exprs.py b/reladiff/sqeleton/bound_exprs.py new file mode 100644 index 00000000..b193de2d --- /dev/null +++ b/reladiff/sqeleton/bound_exprs.py @@ -0,0 +1,81 @@ +"""Expressions bound to a specific database""" + +import inspect +from functools import wraps +from typing import Union, TYPE_CHECKING + +from runtype import dataclass + +from .abcs import AbstractDatabase, AbstractCompiler +from .queries.ast_classes import ExprNode, TablePath, Compilable +from .queries.api import table +from .schema import create_schema + + +@dataclass +class BoundNode(ExprNode): + database: AbstractDatabase + node: Compilable + + def __getattr__(self, attr): + value = getattr(self.node, attr) + if inspect.ismethod(value): + + @wraps(value) + def bound_method(*args, **kw): + return BoundNode(self.database, value(*args, **kw)) + + return bound_method + return value + + def query(self, res_type=list): + return self.database.query(self.node, res_type=res_type) + + @property + def type(self): + return self.node.type + + def compile(self, c: AbstractCompiler) -> str: + assert c.database is self.database + return c.compile_elem(self.node) + + +def bind_node(node, database): + return BoundNode(database, node) + + +ExprNode.bind = bind_node + + +@dataclass +class BoundTable(BoundNode): # ITable + database: AbstractDatabase + node: TablePath + + def with_schema(self, schema): + table_path = self.node.replace(schema=schema) + return self.replace(node=table_path) + + def query_schema(self, *, refine: bool = True, refine_where = None, case_sensitive=True): + table_path = self.node + + if table_path.schema: + return self + + raw_schema = self.database.query_table_schema(table_path.path) + schema, _samples = self.database.process_query_table_schema(table_path.path, raw_schema, refine, refine_where) + schema = create_schema(self.database, table_path.path, schema, case_sensitive) + return self.with_schema(schema) + + @property + def schema(self): + return self.node.schema + + +def bound_table(database: AbstractDatabase, table_path: Union[TablePath, str, tuple], **kw): + return BoundTable(database, table(table_path, **kw)) + + +if TYPE_CHECKING: + class BoundTable(BoundTable, TablePath): + pass diff --git a/reladiff/sqeleton/conn_editor.py b/reladiff/sqeleton/conn_editor.py new file mode 100644 index 00000000..31ca66e9 --- /dev/null +++ b/reladiff/sqeleton/conn_editor.py @@ -0,0 +1,292 @@ +import sys +import re +from concurrent.futures import ThreadPoolExecutor +import toml + +from runtype import dataclass + +try: + import textual +except ModuleNotFoundError: + raise ModuleNotFoundError( + "Error: Cannot find the TUI library 'textual'. Please install it using `pip install sqeleton[tui]`" + ) + +from textual.app import App, ComposeResult +from textual.containers import Vertical, Container, Horizontal +from textual.widgets import Header, Footer, Input, Label, Button, ListView, ListItem + +from textual_select import Select + +from .databases._connect import DATABASE_BY_SCHEME +from . import connect + + +class ContentSwapper(Container): + def __init__(self, *initial_widgets, **kw): + self.container = Container(*initial_widgets) + super().__init__(**kw) + + def compose(self): + yield self.container + + def new_content(self, *new_widgets): + for c in self.container.children: + c.remove() + + self.container.mount(*new_widgets) + + +def test_connect(connect_dict): + conn = connect(connect_dict) + assert conn.query("select 1+1", int) == 2 + + +@dataclass +class Config: + config: dict + + @property + def databases(self): + if "database" not in self.config: + self.config["database"] = {} + assert isinstance(self.config["database"], dict) + return self.config["database"] + + def add_database(self, name: str, db: dict): + assert isinstance(db, dict) + self.config["database"][name] = db + + def remove_database(self, name): + del self.config["database"][name] + + +class EditConnection(Vertical): + def __init__(self, db_name: str, config: Config) -> None: + self._db_name = db_name + self._config = config + super().__init__() + + def compose(self): + self.params_container = ContentSwapper(id="params") + self.driver_select = Select( + placeholder="Select a database driver", + search=True, + items=[{"value": k, "text": k} for k in DATABASE_BY_SCHEME], + list_mount="#driver_container", + value=self._config.databases.get(self._db_name, {}).get("driver"), + id="driver", + ) + + yield Vertical( + Label("Connection name:"), + Input(id="conn_name", value=self._db_name, classes="margin-bottom-1"), + Label("Database Driver:"), + self.driver_select, + self.params_container, + id="driver_container", + ) + + self.create_params() + + def create_params(self): + driver = self.driver_select.value + if not driver: + return + db_config = self._config.databases.get(self._db_name, {}) + + db_cls = DATABASE_BY_SCHEME[driver] + # Filter out repetitions, but keep order + base_params = re.findall("<(.*?)>", db_cls.CONNECT_URI_HELP) + params = dict.fromkeys([p.lower() for p in base_params] + [k for k in db_config if k != "driver"]) + + widgets = [] + for p in params: + label = p + if p in ("user", "pass", "port", "dbname"): + label += " (optional)" + widgets.append(Label(f"{label}:")) + p = p.lower() + widgets.append(Input(id="input_" + p, name=p, classes="param", value=db_config.get(p))) + + self.params_container.new_content( + Vertical( + *widgets, + Horizontal( + Button("Test & Save Connection", id="test_and_save"), + Button("Save Connection Without Testing", id="save_without_test"), + ), + Vertical(id="result"), + ) + ) + + def on_select_changed(self, event: Select.Changed) -> None: + driver = str(event.value) + self.driver_select.text = driver + self.driver_select.value = driver + self.create_params() + + def _get_connect_dict(self): + connect_dict = {"driver": self.driver_select.value} + for p in self.query(".param"): + if p.value: + connect_dict[p.name] = p.value + + return connect_dict + + async def on_button_pressed(self, event: Button.Pressed) -> None: + button_id = event.button.id + if button_id == "test_and_save": + connect_dict = self._get_connect_dict() + + result_container = self.query_one("#result") + result_container.mount(Label(f"Trying to connect to {connect_dict}")) + + try: + test_connect(connect_dict) + except Exception as e: + error = str(e) + result_container.mount(Label(f"[red]Error: {error}[/red]")) + else: + result_container.mount(Label(f"[green]Success![green]")) + self.save(connect_dict) + + elif button_id == "save_without_test": + connect_dict = self._get_connect_dict() + self.save(connect_dict) + + def save(self, connect_dict): + assert isinstance(connect_dict, dict) + result_container = self.query_one("#result") + result_container.mount(Label(f"Database saved")) + + name = self.query_one("#conn_name").value + self._config.add_database(name, connect_dict) + self.app.config_changed() + + +class ListOfConnections(Vertical): + def __init__(self, config: Config, **kw): + self.config = config + super().__init__(**kw) + + def compose(self) -> ComposeResult: + list_items = [ + ListItem(Label(name, id="list_label_" + name), name=name) for name in self.config.databases.keys() + ] + + self.lv = lv = ListView(*list_items, id="connection_list") + yield lv + + lv.focus() + + def on_list_view_highlighted(self, h: ListView.Highlighted): + name = h.item.name + self.app.query_one("#edit_container").new_content(EditConnection(name, self.config)) + + +class ConnectionEditor(App): + CSS = """ + #conn_list { + display: block; + dock: left; + height: 100%; + max-width: 30%; + margin: 1 + } + + #edit { + margin: 1 + } + + #test_and_save { + margin: 1 + } + #save_without_test { + margin: 1 + } + + .margin-bottom-1 { + margin-bottom: 1 + } + """ + + BINDINGS = [ + ("q", "quit", "Quit"), + ("d", "toggle_dark", "Toggle dark mode"), + ("insert", "add_conn", "Add new connection"), + ("delete", "del_conn", "Delete selected connection"), + ("t", "test_conn", "Test selected connection"), + ("a", "test_all_conns", "Test all connections"), + ] + + def run(self, toml_path, **kw): + self.toml_path = toml_path + try: + with open(self.toml_path) as f: + self.config = Config(toml.load(f)) + except FileNotFoundError: + self.config = Config({}) + + return super().run(**kw) + + def config_changed(self): + self.list_swapper.new_content(ListOfConnections(self.config)) + with open(self.toml_path, "w") as f: + toml.dump(self.config.config, f) + + def compose(self) -> ComposeResult: + """Create child widgets for the app.""" + self.list_swapper = ContentSwapper(ListOfConnections(self.config), id="conn_list") + self.edit_swapper = ContentSwapper(id="edit_container") + self.edit_swapper.new_content(EditConnection("New", self.config)) + + yield Header() + yield Container(self.list_swapper, self.edit_swapper) + yield Footer() + + def action_toggle_dark(self) -> None: + """An action to toggle dark mode.""" + self.dark = not self.dark + + def action_add_conn(self): + self.edit_swapper.new_content(EditConnection("New", self.config)) + + def _selected_connection(self): + connection_list: ListView = self.query_one("#connection_list") + return connection_list.children[connection_list.index] + + def action_del_conn(self): + name = self._selected_connection().name + self.config.remove_database(name) + self.config_changed() + + def _test_existing_db(self, name): + label: Label = self.query_one("#list_label_" + name) + label.update(f"{name}🔃") + connect_dict = self.config.databases[name] + try: + test_connect(connect_dict) + except Exception as e: + label.update(f"{name}❌ {str(e)[:16]}") + else: + label.update(f"{name}✅") + + def action_test_conn(self): + name = self._selected_connection().name + t = ThreadPoolExecutor() + t.submit(self._test_existing_db, name) + + def action_test_all_conns(self): + t = ThreadPoolExecutor() + for name in self.config.databases: + t.submit(self._test_existing_db, name) + + +def main(toml_path: str): + app = ConnectionEditor() + app.run(toml_path) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/reladiff/sqeleton/databases/__init__.py b/reladiff/sqeleton/databases/__init__.py new file mode 100644 index 00000000..6da5dd64 --- /dev/null +++ b/reladiff/sqeleton/databases/__init__.py @@ -0,0 +1,19 @@ +from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, QueryError, ConnectError, BaseDialect, Database, logger +from ..abcs import DbPath, DbKey, DbTime +from ._connect import Connect + +from .postgresql import PostgreSQL +from .mysql import MySQL +from .oracle import Oracle +from .snowflake import Snowflake +from .bigquery import BigQuery +from .redshift import Redshift +from .presto import Presto +from .databricks import Databricks +from .trino import Trino +from .clickhouse import Clickhouse +from .vertica import Vertica +from .duckdb import DuckDB +from .dremio import Dremio + +connect = Connect() diff --git a/reladiff/sqeleton/databases/_connect.py b/reladiff/sqeleton/databases/_connect.py new file mode 100644 index 00000000..e1af62c7 --- /dev/null +++ b/reladiff/sqeleton/databases/_connect.py @@ -0,0 +1,276 @@ +import urllib +from typing import Type, Optional, Union, Dict +from itertools import zip_longest +from contextlib import suppress +import dsnparse +import toml + +from runtype import dataclass + +from ..abcs.mixins import AbstractMixin +from ..utils import WeakCache, Self +from .base import Database, ThreadedDatabase +from .postgresql import PostgreSQL +from .mysql import MySQL +from .oracle import Oracle +from .snowflake import Snowflake +from .bigquery import BigQuery +from .redshift import Redshift +from .presto import Presto +from .databricks import Databricks +from .trino import Trino +from .clickhouse import Clickhouse +from .vertica import Vertica +from .duckdb import DuckDB +from .dremio import Dremio + + +@dataclass +class MatchUriPath: + database_cls: Type[Database] + + def match_path(self, dsn): + help_str = self.database_cls.CONNECT_URI_HELP + params = self.database_cls.CONNECT_URI_PARAMS + kwparams = self.database_cls.CONNECT_URI_KWPARAMS + + dsn_dict = dict(urllib.parse.parse_qsl(dsn.query)) + matches = {} + for param, arg in zip_longest(params, dsn.paths): + if param is None: + raise ValueError(f"Too many parts to path. Expected format: {help_str}") + + optional = param.endswith("?") + param = param.rstrip("?") + + if arg is None: + try: + arg = dsn_dict.pop(param) + except KeyError: + if not optional: + raise ValueError(f"URI must specify '{param}'. Expected format: {help_str}") + + arg = None + + assert param and param not in matches + matches[param] = arg + + for param in kwparams: + try: + arg = dsn_dict.pop(param) + except KeyError: + raise ValueError(f"URI must specify '{param}'. Expected format: {help_str}") + + assert param and arg and param not in matches, (param, arg, matches.keys()) + matches[param] = arg + + for param, value in dsn_dict.items(): + if param in matches: + raise ValueError( + f"Parameter '{param}' already provided as positional argument. Expected format: {help_str}" + ) + + matches[param] = value + + return matches + + +DATABASE_BY_SCHEME = { + "postgresql": PostgreSQL, + "mysql": MySQL, + "oracle": Oracle, + "redshift": Redshift, + "snowflake": Snowflake, + "presto": Presto, + "bigquery": BigQuery, + "databricks": Databricks, + "duckdb": DuckDB, + "trino": Trino, + "clickhouse": Clickhouse, + "vertica": Vertica, + "dremio": Dremio, +} + + +class Connect: + """Provides methods for connecting to a supported database using a URL or connection dict.""" + + def __init__(self, database_by_scheme: Dict[str, Database] = DATABASE_BY_SCHEME): + self.database_by_scheme = database_by_scheme + self.match_uri_path = {name: MatchUriPath(cls) for name, cls in database_by_scheme.items()} + self.conn_cache = WeakCache() + + def for_databases(self, *dbs): + database_by_scheme = {k: db for k, db in self.database_by_scheme.items() if k in dbs} + return type(self)(database_by_scheme) + + def load_mixins(self, *abstract_mixins: AbstractMixin) -> Self: + "Extend all the databases with a list of mixins that implement the given abstract mixins." + database_by_scheme = {k: db.load_mixins(*abstract_mixins) for k, db in self.database_by_scheme.items()} + return type(self)(database_by_scheme) + + def connect_to_uri(self, db_uri: str, thread_count: Optional[int] = 1) -> Database: + """Connect to the given database uri + + thread_count determines the max number of worker threads per database, + if relevant. None means no limit. + + Parameters: + db_uri (str): The URI for the database to connect + thread_count (int, optional): Size of the threadpool. Ignored by cloud databases. (default: 1) + + Note: For non-cloud databases, a low thread-pool size may be a performance bottleneck. + + Supported schemes: + - postgresql + - mysql + - oracle + - snowflake + - bigquery + - redshift + - presto + - databricks + - trino + - clickhouse + - vertica + - duckdb + - dremio + """ + + dsn = dsnparse.parse(db_uri) + if len(dsn.schemes) > 1: + raise NotImplementedError("No support for multiple schemes") + (scheme,) = dsn.schemes + + if scheme == "toml": + toml_path = dsn.path or dsn.host + database = dsn.fragment + if not database: + raise ValueError("Must specify a database name, e.g. 'toml://path#database'. ") + with open(toml_path) as f: + config = toml.load(f) + try: + conn_dict = config["database"][database] + except KeyError: + raise ValueError(f"Cannot find database config named '{database}'.") + return self.connect_with_dict(conn_dict, thread_count) + + try: + matcher = self.match_uri_path[scheme] + except KeyError: + raise NotImplementedError(f"Scheme '{scheme}' currently not supported") + + cls = matcher.database_cls + + if scheme == "databricks": + assert not dsn.user + kw = {} + kw["access_token"] = dsn.password + kw["http_path"] = dsn.path + kw["server_hostname"] = dsn.host + kw.update(dsn.query) + elif scheme == "duckdb": + kw = {} + kw["filepath"] = dsn.dbname + kw["dbname"] = dsn.user + else: + kw = matcher.match_path(dsn) + + if scheme == "bigquery": + kw["project"] = dsn.host + return cls(**kw) + + if scheme == "snowflake": + kw["account"] = dsn.host + assert not dsn.port + kw["user"] = dsn.user + kw["password"] = dsn.password + else: + kw["host"] = dsn.host + kw["port"] = dsn.port + kw["user"] = dsn.user + if dsn.password: + kw["password"] = dsn.password + + kw = {k: v for k, v in kw.items() if v is not None} + + if issubclass(cls, ThreadedDatabase): + db = cls(thread_count=thread_count, **kw) + else: + db = cls(**kw) + + return self._connection_created(db) + + def connect_with_dict(self, d, thread_count): + d = dict(d) + driver = d.pop("driver") + try: + matcher = self.match_uri_path[driver] + except KeyError: + raise NotImplementedError(f"Driver '{driver}' currently not supported") + + cls = matcher.database_cls + if issubclass(cls, ThreadedDatabase): + db = cls(thread_count=thread_count, **d) + else: + db = cls(**d) + + return self._connection_created(db) + + def _connection_created(self, db): + "Nop function to be overridden by subclasses." + return db + + def __call__(self, db_conf: Union[str, dict], thread_count: Optional[int] = 1, shared: bool = True) -> Database: + """Connect to a database using the given database configuration. + + Configuration can be given either as a URI string, or as a dict of {option: value}. + + The dictionary configuration uses the same keys as the TOML 'database' definition given with --conf. + + thread_count determines the max number of worker threads per database, + if relevant. None means no limit. + + Parameters: + db_conf (str | dict): The configuration for the database to connect. URI or dict. + thread_count (int, optional): Size of the threadpool. Ignored by cloud databases. (default: 1) + shared (bool): Whether to cache and return the same connection for the same db_conf. (default: True) + + Note: For non-cloud databases, a low thread-pool size may be a performance bottleneck. + + Supported drivers: + - postgresql + - mysql + - oracle + - snowflake + - bigquery + - redshift + - presto + - databricks + - trino + - clickhouse + - vertica + - dremio + + Example: + >>> connect("mysql://localhost/db") + + >>> connect({"driver": "mysql", "host": "localhost", "database": "db"}) + + """ + if shared: + with suppress(KeyError): + conn = self.conn_cache.get(db_conf) + if not conn.is_closed: + return conn + + if isinstance(db_conf, str): + conn = self.connect_to_uri(db_conf, thread_count) + elif isinstance(db_conf, dict): + conn = self.connect_with_dict(db_conf, thread_count) + else: + raise TypeError(f"db configuration must be a URI string or a dictionary. Instead got '{db_conf}'.") + + if shared: + self.conn_cache.add(db_conf, conn) + return conn diff --git a/reladiff/sqeleton/databases/base.py b/reladiff/sqeleton/databases/base.py new file mode 100644 index 00000000..100c8ea7 --- /dev/null +++ b/reladiff/sqeleton/databases/base.py @@ -0,0 +1,659 @@ +import uuid +from datetime import datetime +import math +import sys +import logging +from typing import Any, Callable, Dict, Generator, Tuple, Optional, Sequence, List, Union, TypeVar, Type, overload +from functools import partial, wraps +from concurrent.futures import ThreadPoolExecutor +import threading +from abc import abstractmethod + +from runtype import dataclass, issubclass, pytypes + +from sqeleton.queries.compiler import CompiledCode + +from ..utils import is_uuid, safezip, Self +from ..queries import ExprNode, Compiler, table, Select, SKIP, T_SKIP, Explain, Code, this, commit +from ..queries.ast_classes import ForeignKey, Random, CompilableNode, TablePath +from ..abcs.database_types import ( + AbstractDatabase, + AbstractDialect, + AbstractTable, + ColType, + Integer, + Decimal, + Float, + Native_UUID, + String_UUID, + String_Alphanum, + String_VaryingAlphanum, + TemporalType, + UnknownColType, + TimestampTZ, + Text, + DbTime, + DbPath, + Boolean, +) +from ..abcs.mixins import Compilable +from ..abcs.mixins import ( + AbstractMixin_Schema, + AbstractMixin_RandomSample, + AbstractMixin_NormalizeValue, + AbstractMixin_OptimizerHints, +) +from ..bound_exprs import bound_table + +logger = logging.getLogger("database") + + +def parse_table_name(t): + return tuple(t.split(".")) + + +def import_helper(package: str = None, text=""): + def dec(f): + @wraps(f) + def _inner(): + try: + return f() + except ModuleNotFoundError as e: + s = text + if package: + s += f"You can install it using 'pip install sqeleton[{package}]'." + raise ModuleNotFoundError(f"{e}\n\n{s}\n") + + return _inner + + return dec + + +class ConnectError(Exception): + pass + + +class QueryError(Exception): + pass + + +def _one(seq): + (x,) = seq + return x + + +@dataclass +class QueryResult: + rows: list + columns: list = None + + def __iter__(self): + return iter(self.rows) + + def __len__(self): + return len(self.rows) + + def __getitem__(self, i): + return self.rows[i] + + +class ThreadLocalInterpreter: + """An interpeter used to execute a sequence of queries within the same thread and cursor. + + Useful for cursor-sensitive operations, such as creating a temporary table. + """ + + def __init__(self, compiler: Compiler, gen: Generator): + self.gen = gen + self.compiler = compiler + + def apply_queries(self, callback: Callable[[CompiledCode], Any]) -> None: + q: ExprNode = next(self.gen) + while True: + sql = self.compiler.compile_with_args(q) + try: + try: + res = callback(sql) if sql is not SKIP else SKIP + except Exception as e: + q = self.gen.throw(type(e), e) + else: + q = self.gen.send(res) + except StopIteration: + break + + +SqlCode = Union[str, CompiledCode, ThreadLocalInterpreter] + + +def apply_query(callback: Callable[[CompiledCode], Any], sql_code: SqlCode) -> Optional[QueryResult]: + if isinstance(sql_code, ThreadLocalInterpreter): + return sql_code.apply_queries(callback) + elif isinstance(sql_code, str): + sql_code = CompiledCode(sql_code, [], None) # Unknown type. #TODO: Should we guess? + + return callback(sql_code) + + +class Mixin_Schema(AbstractMixin_Schema): + def table_information(self) -> TablePath: + return table("information_schema", "tables") + + def list_tables(self, table_schema: str, like: Compilable = None) -> Select: + return ( + self.table_information() + .where( + this.table_schema == table_schema, + this.table_name.like(like) if like is not None else SKIP, + this.table_type == "BASE TABLE", + ) + .select(this.table_name) + ) + + +class Mixin_RandomSample(AbstractMixin_RandomSample): + def random_sample_n(self, tbl: AbstractTable, size: int) -> AbstractTable: + # TODO use a more efficient algorithm, when the table count is known + return tbl.order_by(Random()).limit(size) + + def random_sample_ratio_approx(self, tbl: AbstractTable, ratio: float) -> AbstractTable: + return tbl.where(Random() < ratio) + + +class Mixin_OptimizerHints(AbstractMixin_OptimizerHints): + def optimizer_hints(self, hints: str) -> str: + return f"/*+ {hints} */ " + + +class BaseDialect(AbstractDialect): + SUPPORTS_PRIMARY_KEY = False + SUPPORTS_INDEXES = False + TYPE_CLASSES: Dict[str, type] = {} + MIXINS = frozenset() + ARG_SYMBOL = "%s" + PLACEHOLDER_TABLE = None # Used for Oracle + + def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None): + if offset: + raise NotImplementedError("No support for OFFSET in query") + + return f"LIMIT {limit}" + + def concat(self, items: List[str]) -> str: + assert len(items) > 1 + joined_exprs = ", ".join(items) + return f"concat({joined_exprs})" + + def is_distinct_from(self, a: str, b: str) -> str: + return f"{a} is distinct from {b}" + + def timestamp_value(self, t: DbTime) -> str: + return f"'{t.isoformat()}'" + + def uuid_value(self, u: uuid.UUID) -> str: + return f"'{u}'" + + def random(self) -> str: + return "random()" + + def current_timestamp(self) -> str: + return "current_timestamp()" + + def explain_as_text(self, query: str) -> str: + return f"EXPLAIN {query}" + + def immediate_values(self, rows) -> str: + values = ", ".join("(%s)" % ", ".join(row) for row in rows) + return f"VALUES {values}" + + def type_repr(self, t) -> str: + if isinstance(t, str): + return t + elif isinstance(t, TimestampTZ): + return f"TIMESTAMP({min(t.precision, DEFAULT_DATETIME_PRECISION)})" + elif isinstance(t, ForeignKey): + return self.type_repr(t.type) + return { + int: "INT", + str: "VARCHAR", + bytes: "BYTEA", + bool: "BOOLEAN", + float: "FLOAT", + datetime: "TIMESTAMP", + }[t] + + # def decl_repr(self, name, type_): + # if isinstance(type_, ForeignKey): + # return f"FOREIGN KEY ({name}) REFERENCES Persons(PersonID)" + + def _parse_type_repr(self, type_repr: str) -> Optional[Type[ColType]]: + return self.TYPE_CLASSES.get(type_repr) + + def parse_type( + self, + table_path: DbPath, + col_name: str, + type_repr: str, + datetime_precision: int = None, + numeric_precision: int = None, + numeric_scale: int = None, + ) -> ColType: + """ """ + + cls = self._parse_type_repr(type_repr) + if not cls: + return UnknownColType(type_repr) + + if issubclass(cls, TemporalType): + return cls( + precision=datetime_precision if datetime_precision is not None else DEFAULT_DATETIME_PRECISION, + rounds=self.ROUNDS_ON_PREC_LOSS, + ) + + elif issubclass(cls, Integer): + return cls() + + elif issubclass(cls, Boolean): + return cls() + + elif issubclass(cls, Decimal): + if numeric_scale is None: + numeric_scale = 0 # Needed for Oracle. + return cls(precision=numeric_scale) + + elif issubclass(cls, Float): + # assert numeric_scale is None + return cls( + precision=self._convert_db_precision_to_digits( + numeric_precision if numeric_precision is not None else DEFAULT_NUMERIC_PRECISION + ) + ) + + elif issubclass(cls, (Text, Native_UUID)): + return cls() + + raise TypeError(f"Parsing {type_repr} returned an unknown type '{cls}'.") + + def _convert_db_precision_to_digits(self, p: int) -> int: + """Convert from binary precision, used by floats, to decimal precision.""" + # See: https://en.wikipedia.org/wiki/Single-precision_floating-point_format + return math.floor(math.log(2**p, 10)) + + @classmethod + def load_mixins(cls, *abstract_mixins) -> "Self": + mixins = {m for m in cls.MIXINS if issubclass(m, abstract_mixins)} + + class _DialectWithMixins(cls, *mixins, *abstract_mixins): + pass + + _DialectWithMixins.__name__ = cls.__name__ + return _DialectWithMixins() + + +T = TypeVar("T", bound=BaseDialect) +TRes = TypeVar("TRes") + + +QueryInputItem = Union[CompilableNode, T_SKIP] +QueryInput = Union[str, QueryInputItem, Generator, List[QueryInputItem]] + + +class Database(AbstractDatabase[T]): + """Base abstract class for databases. + + Used for providing connection code and implementation specific SQL utilities. + + Instanciated using :meth:`~sqeleton.connect` + """ + + default_schema: str = None + SUPPORTS_ALPHANUMS = True + SUPPORTS_UNIQUE_CONSTAINT = False + + CONNECT_URI_KWPARAMS = [] + + _interactive = False + is_closed = False + + dialect: AbstractDialect + + @property + def name(self): + return type(self).__name__ + + def compile(self, sql_ast): + compiler = Compiler(self) + return compiler.compile(sql_ast) + + # def set_logger_level(self, level: Union[str, int]): + # if isinstance(level, str): + # level = getattr(logging, level) + + # logger.setLevel(level) + + @overload + def query(self, query_input: QueryInput) -> Any: + ... + + @overload + def query(self, query_input: QueryInput, res_type: None) -> Any: + ... + + @overload + def query(self, query_input: QueryInput, res_type: Type[TRes]) -> TRes: + ... + + def query(self, query_input, res_type=None): + """Query the given SQL code/AST, and attempt to convert the result to type 'res_type' + + If given a generator: + It will execute all the yielded sql queries with the same thread and cursor. + The results of the queries are returned by the `yield` stmt (using the .send() mechanism). + It's a cleaner approach than exposing cursors, but may not be enough in all cases. + """ + if query_input is SKIP: + return + + compiler = Compiler(self) + if isinstance(query_input, Generator): + sql_code = ThreadLocalInterpreter(compiler, query_input) + elif isinstance(query_input, list): + for i in query_input[:-1]: + self.query(i) + return self.query(query_input[-1], res_type) + else: + if isinstance(query_input, str): + sql_code = query_input + else: + if res_type is None: + res_type = query_input.type + sql_code = compiler.compile_with_args(query_input) + if sql_code is SKIP: + return SKIP + + if self._interactive and isinstance(query_input, Select): + explained_sql = compiler.compile_with_args(Explain(query_input)) + explain = self._query(explained_sql) + for row in explain: + # Most returned a 1-tuple. Presto returns a string + if isinstance(row, tuple): + (row,) = row + logger.debug("EXPLAIN: %s", row) + answer = input("Continue? [y/n] ") + if answer.lower() not in ["y", "yes"]: + sys.exit(1) + + res = self._query(sql_code) + + if res_type == None: + pass # Do no casting + elif res is None: + assert res_type is not None + raise ValueError(f"Query returned NULL, but query() is expecting type {res_type}") + elif res_type is list: + return list(res) + elif res_type in (int, str): + if not res: + raise ValueError("Query returned 0 rows, expected 1") + row = _one(res) + if not row: + raise ValueError("Row is empty, expected 1 column") + res = _one(row) + if res is None: # May happen due to sum() of 0 items + return None + return res_type(res) + elif res_type is datetime: + res = _one(_one(res)) + if isinstance(res, str): + res = datetime.fromisoformat(res[:23]) # TODO use a better parsing method + return res + elif res_type is tuple: + assert len(res) == 1, (sql_code, res) + return tuple(res[0]) + else: + # TODO fix this API from runtype side + res_type = pytypes.type_caster.to_canon(res_type) + if isinstance(res_type, pytypes.SequenceType): + item_type = res_type.item + if issubclass(item_type, Union[int, str, bytes, float]): + return [_one(row) for row in res] + elif issubclass(item_type, tuple): + return [tuple(row) for row in res] + elif issubclass(item_type, dict): + return [dict(safezip(res.columns, row)) for row in res] + elif issubclass(item_type, type): + (elem_type,) = res_type.__args__ + return [elem_type(**dict(safezip(res.columns, row))) for row in res] + + if len(res) == 0: + return None # TODO: Only allow if res_type is Optional + assert len(res) == 1, len(res) + d = dict(safezip(res.columns, res[0])) + return res_type(**d) + return res + + def enable_interactive(self): + self._interactive = True + + def select_table_schema(self, path: DbPath) -> str: + """Provide SQL for selecting the table schema as (name, type, date_prec, num_prec)""" + schema, name = self._normalize_table_path(path) + + return ( + "SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale " + "FROM information_schema.columns " + f"WHERE table_name = '{name}' AND table_schema = '{schema}'" + ) + + def query_table_schema(self, path: DbPath) -> Dict[str, tuple]: + rows = self.query(self.select_table_schema(path), list) + if not rows: + raise RuntimeError(f"{self.name}: Table '{'.'.join(path)}' does not exist, or has no columns") + + d = {r[0]: r for r in rows} + assert len(d) == len(rows) + return d + + def select_table_unique_columns(self, path: DbPath) -> str: + schema, name = self._normalize_table_path(path) + + return ( + "SELECT column_name " + "FROM information_schema.key_column_usage " + f"WHERE table_name = '{name}' AND table_schema = '{schema}'" + ) + + def query_table_unique_columns(self, path: DbPath) -> List[str]: + if not self.SUPPORTS_UNIQUE_CONSTAINT: + raise NotImplementedError("This database doesn't support 'unique' constraints") + res = self.query(self.select_table_unique_columns(path), List[str]) + return list(res) + + def _process_table_schema( + self, path: DbPath, raw_schema: Dict[str, tuple], filter_columns: Sequence[str] = None, where: str = None + ): + if filter_columns is None: + filtered_schema = raw_schema + else: + accept = {i.lower() for i in filter_columns} + filtered_schema = {name: row for name, row in raw_schema.items() if name.lower() in accept} + + col_dict = {row[0]: self.dialect.parse_type(path, *row) for _name, row in filtered_schema.items()} + + samples = self._refine_coltypes(path, col_dict, where) + if samples is not None and not samples: + raise ValueError(f"Table {path} appears to be empty") + + # Return a dict of form {name: type} after normalization + return col_dict + + def process_query_table_schema( + self, path: Tuple[str], raw_schema: Dict[str, Tuple], refine: bool = True, refine_where: Optional[str] = None + ) -> Tuple[Dict[str, ColType], Optional[list]]: + col_dict = {name: self.dialect.parse_type(path, *row) for name, row in raw_schema.items()} + + samples = self._refine_coltypes(path, col_dict, refine_where) if refine else None + + return col_dict, samples + + def _refine_coltypes( + self, table_path: DbPath, col_dict: Dict[str, ColType], where: Optional[str] = None, sample_size=64 + ): + """Refine the types in the column dict, by querying the database for a sample of their values + + 'where' restricts the rows to be sampled. + """ + + text_columns = [k for k, v in col_dict.items() if isinstance(v, Text)] + if not text_columns: + return None + + if isinstance(self.dialect, AbstractMixin_NormalizeValue): + fields = [Code(self.dialect.normalize_uuid(self.dialect.quote(c), String_UUID())) for c in text_columns] + else: + fields = this[text_columns] + + samples_by_row = self.query( + table(*table_path).select(*fields).where(Code(where) if where else SKIP).limit(sample_size), list + ) + if not samples_by_row: + return [] + + samples_by_col = list(zip(*samples_by_row)) + + for col_name, samples in safezip(text_columns, samples_by_col): + uuid_samples = [s for s in samples if s and is_uuid(s)] + + if uuid_samples: + if len(uuid_samples) != len(samples): + logger.warning( + f"Mixed UUID/Non-UUID values detected in column {'.'.join(table_path)}.{col_name}, disabling UUID support." + ) + else: + assert col_name in col_dict + col_dict[col_name] = String_UUID() + continue + + if self.SUPPORTS_ALPHANUMS: # Anything but MySQL (so far) + alphanum_samples = [s for s in samples if String_Alphanum.test_value(s)] + if alphanum_samples: + if len(alphanum_samples) != len(samples): + logger.debug( + f"Mixed Alphanum/Non-Alphanum values detected in column {'.'.join(table_path)}.{col_name}. It cannot be used as a key." + ) + else: + assert col_name in col_dict + col_dict[col_name] = String_VaryingAlphanum() + + return samples_by_row + + # @lru_cache() + # def get_table_schema(self, path: DbPath) -> Dict[str, ColType]: + # return self.query_table_schema(path) + + def _normalize_table_path(self, path: DbPath) -> DbPath: + if len(path) == 1: + return self.default_schema, path[0] + elif len(path) == 2: + return path + + 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: CompiledCode) -> Optional[QueryResult]: + assert isinstance(sql_code, CompiledCode), sql_code + try: + logger.debug(f"{self.name} Executing SQL: {sql_code.code} || {sql_code.args}") + c.execute(sql_code.code, sql_code.args or ()) + # insert, delete and update may return values if they have the "returning" clause. + if sql_code.type is not None or sql_code.code.lstrip().lower().startswith( + ("select", "explain", "show", "with") + ): + columns = c.description and [col[0] for col in c.description] + return QueryResult(c.fetchall(), columns) + except Exception as _e: + # logger.exception(e) + # logger.error(f'Caused by SQL: {sql_code}') + raise + + def _query_conn(self, conn, sql_code: SqlCode) -> Optional[QueryResult]: + c = conn.cursor() + callback = partial(self._query_cursor, c) + return apply_query(callback, sql_code) + + def close(self): + self.is_closed = True + return super().close() + + def list_tables(self, tables_like, schema=None): + assert isinstance(self.dialect, Mixin_Schema) + return self.query(self.dialect.list_tables(schema or self.default_schema, tables_like)) + + def table(self, *path, **kw): + return bound_table(self, path, **kw) + + @classmethod + def load_mixins(cls, *abstract_mixins) -> type: + class _DatabaseWithMixins(cls): + dialect = cls.dialect.load_mixins(*abstract_mixins) + + _DatabaseWithMixins.__name__ = cls.__name__ + return _DatabaseWithMixins + + def commit(self): + return self.query(commit) + + +class ThreadedDatabase(Database): + """Access the database through singleton threads. + + Used for database connectors that do not support sharing their connection between different threads. + """ + + def __init__(self, thread_count=1): + self._init_error = None + self._queue = ThreadPoolExecutor(thread_count, initializer=self.set_conn) + self.thread_local = threading.local() + logger.info(f"[{self.name}] Starting a threadpool, size={thread_count}.") + + def set_conn(self): + assert not hasattr(self.thread_local, "conn") + try: + self.thread_local.conn = self.create_connection() + except Exception as e: + self._init_error = e + + def _query(self, sql_code: SqlCode) -> QueryResult: + r = self._queue.submit(self._query_in_worker, sql_code) + return r.result() + + def _query_in_worker(self, sql_code: SqlCode): + "This method runs in a worker thread" + if self._init_error: + raise self._init_error + return self._query_conn(self.thread_local.conn, sql_code) + + @abstractmethod + def create_connection(self): + "Return a connection instance, that supports the .cursor() method." + + def close(self): + super().close() + self._queue.shutdown() + + @property + def is_autocommit(self) -> bool: + return False + + +CHECKSUM_HEXDIGITS = 15 # Must be 15 or lower, otherwise SUM() overflows +MD5_HEXDIGITS = 32 + +_CHECKSUM_BITSIZE = CHECKSUM_HEXDIGITS << 2 +CHECKSUM_MASK = (2**_CHECKSUM_BITSIZE) - 1 + +DEFAULT_DATETIME_PRECISION = 6 +DEFAULT_NUMERIC_PRECISION = 24 + +TIMESTAMP_PRECISION_POS = 20 # len("2022-06-03 12:24:35.") == 20 diff --git a/reladiff/sqeleton/databases/bigquery.py b/reladiff/sqeleton/databases/bigquery.py new file mode 100644 index 00000000..0b4dc66c --- /dev/null +++ b/reladiff/sqeleton/databases/bigquery.py @@ -0,0 +1,213 @@ +from typing import List, Union +from ..abcs.database_types import ( + Timestamp, + Datetime, + Integer, + Decimal, + Float, + Text, + DbPath, + FractionalType, + TemporalType, + Boolean, +) +from ..abcs.mixins import ( + AbstractMixin_MD5, + AbstractMixin_NormalizeValue, + AbstractMixin_Schema, + AbstractMixin_TimeTravel, +) +from ..abcs import Compilable +from ..queries import this, table, SKIP, code +from .base import BaseDialect, Database, import_helper, parse_table_name, ConnectError, apply_query, QueryResult +from .base import TIMESTAMP_PRECISION_POS, ThreadLocalInterpreter, Mixin_RandomSample + + +@import_helper(text="Please install BigQuery and configure your google-cloud access.") +def import_bigquery(): + from google.cloud import bigquery + + return bigquery + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"cast(cast( ('0x' || substr(TO_HEX(md5({s})), 18)) as int64) as numeric)" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + if coltype.rounds: + timestamp = f"timestamp_micros(cast(round(unix_micros(cast({value} as timestamp))/1000000, {coltype.precision})*1000000 as int))" + return f"FORMAT_TIMESTAMP('%F %H:%M:%E6S', {timestamp})" + + if coltype.precision == 0: + return f"FORMAT_TIMESTAMP('%F %H:%M:%S.000000', {value})" + elif coltype.precision == 6: + return f"FORMAT_TIMESTAMP('%F %H:%M:%E6S', {value})" + + timestamp6 = f"FORMAT_TIMESTAMP('%F %H:%M:%E6S', {value})" + return ( + f"RPAD(LEFT({timestamp6}, {TIMESTAMP_PRECISION_POS+coltype.precision}), {TIMESTAMP_PRECISION_POS+6}, '0')" + ) + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + return f"format('%.{coltype.precision}f', {value})" + + def normalize_boolean(self, value: str, _coltype: Boolean) -> str: + return self.to_string(f"cast({value} as int)") + + +class Mixin_Schema(AbstractMixin_Schema): + def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable: + return ( + table(table_schema, "INFORMATION_SCHEMA", "TABLES") + .where( + this.table_schema == table_schema, + this.table_name.like(like) if like is not None else SKIP, + this.table_type == "BASE TABLE", + ) + .select(this.table_name) + ) + + +class Mixin_TimeTravel(AbstractMixin_TimeTravel): + def time_travel( + self, + table: Compilable, + before: bool = False, + timestamp: Compilable = None, + offset: Compilable = None, + statement: Compilable = None, + ) -> Compilable: + if before: + raise NotImplementedError("before=True not supported for BigQuery time-travel") + + if statement is not None: + raise NotImplementedError("BigQuery time-travel doesn't support querying by statement id") + + if timestamp is not None: + assert offset is None + return code("{table} FOR SYSTEM_TIME AS OF {timestamp}", table=table, timestamp=timestamp) + + assert offset is not None + return code( + "{table} FOR SYSTEM_TIME AS OF TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL {offset} HOUR);", + table=table, + offset=offset, + ) + + +class Dialect(BaseDialect, Mixin_Schema): + name = "BigQuery" + ROUNDS_ON_PREC_LOSS = False # Technically BigQuery doesn't allow implicit rounding or truncation + TYPE_CLASSES = { + # Dates + "TIMESTAMP": Timestamp, + "DATETIME": Datetime, + # Numbers + "INT64": Integer, + "INT32": Integer, + "NUMERIC": Decimal, + "BIGNUMERIC": Decimal, + "FLOAT64": Float, + "FLOAT32": Float, + # Text + "STRING": Text, + # Boolean + "BOOL": Boolean, + } + MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_TimeTravel, Mixin_RandomSample} + + def random(self) -> str: + return "RAND()" + + def quote(self, s: str): + return f"`{s}`" + + def to_string(self, s: str): + return f"cast({s} as string)" + + def type_repr(self, t) -> str: + try: + return {str: "STRING", float: "FLOAT64"}[t] + except KeyError: + return super().type_repr(t) + + def set_timezone_to_utc(self) -> str: + raise NotImplementedError() + + +class BigQuery(Database): + CONNECT_URI_HELP = "bigquery:///" + CONNECT_URI_PARAMS = ["dataset"] + dialect = Dialect() + + def __init__(self, project, *, dataset, **kw): + bigquery = import_bigquery() + + self._client = bigquery.Client(project, **kw) + self.project = project + self.dataset = dataset + + self.default_schema = dataset + + def _normalize_returned_value(self, value): + if isinstance(value, bytes): + return value.decode() + return value + + def _query_atom(self, sql_code: str): + from google.cloud import bigquery + + try: + result = self._client.query(sql_code).result() + columns = [c.name for c in result.schema] + rows = list(result) + except Exception as e: + msg = "Exception when trying to execute SQL code:\n %s\n\nGot error: %s" + raise ConnectError(msg % (sql_code, e)) + + if rows and isinstance(rows[0], bigquery.table.Row): + rows = [tuple(self._normalize_returned_value(v) for v in row.values()) for row in rows] + return QueryResult(rows, columns) + + def _query(self, sql_code: Union[str, ThreadLocalInterpreter]) -> QueryResult: + return apply_query(self._query_atom, sql_code) + + def close(self): + super().close() + self._client.close() + + def select_table_schema(self, path: DbPath) -> str: + project, schema, name = self._normalize_table_path(path) + return ( + "SELECT column_name, data_type, 6 as datetime_precision, 38 as numeric_precision, 9 as numeric_scale " + f"FROM `{project}`.`{schema}`.INFORMATION_SCHEMA.COLUMNS " + f"WHERE table_name = '{name}' AND table_schema = '{schema}'" + ) + + def query_table_unique_columns(self, path: DbPath) -> List[str]: + return [] + + def _normalize_table_path(self, path: DbPath) -> DbPath: + if len(path) == 0: + raise ValueError(f"{self.name}: Bad table path for {self}: ()") + elif len(path) == 1: + return (self.project, self.default_schema, path[0]) + elif len(path) == 2: + return (self.project,) + path + elif len(path) == 3: + return path + else: + raise ValueError( + 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 diff --git a/reladiff/sqeleton/databases/clickhouse.py b/reladiff/sqeleton/databases/clickhouse.py new file mode 100644 index 00000000..f5cf6201 --- /dev/null +++ b/reladiff/sqeleton/databases/clickhouse.py @@ -0,0 +1,197 @@ +from typing import Optional, Type + +from .base import ( + MD5_HEXDIGITS, + CHECKSUM_HEXDIGITS, + TIMESTAMP_PRECISION_POS, + BaseDialect, + ThreadedDatabase, + import_helper, + ConnectError, + Mixin_RandomSample, +) +from ..abcs.database_types import ( + ColType, + Decimal, + Float, + Integer, + FractionalType, + Native_UUID, + TemporalType, + Text, + Timestamp, + Boolean, +) +from ..abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue + +# https://clickhouse.com/docs/en/operations/server-configuration-parameters/settings/#default-database +DEFAULT_DATABASE = "default" + + +@import_helper("clickhouse") +def import_clickhouse(): + import clickhouse_driver + + return clickhouse_driver + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + substr_idx = 1 + MD5_HEXDIGITS - CHECKSUM_HEXDIGITS + return f"reinterpretAsUInt128(reverse(unhex(lowerUTF8(substr(hex(MD5({s})), {substr_idx})))))" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_number(self, value: str, coltype: FractionalType) -> str: + # If a decimal value has trailing zeros in a fractional part, when casting to string they are dropped. + # For example: + # select toString(toDecimal128(1.10, 2)); -- the result is 1.1 + # select toString(toDecimal128(1.00, 2)); -- the result is 1 + # So, we should use some custom approach to save these trailing zeros. + # To avoid it, we can add a small value like 0.000001 to prevent dropping of zeros from the end when casting. + # For examples above it looks like: + # select toString(toDecimal128(1.10, 2 + 1) + toDecimal128(0.001, 3)); -- the result is 1.101 + # After that, cut an extra symbol from the string, i.e. 1.101 -> 1.10 + # So, the algorithm is: + # 1. Cast to decimal with precision + 1 + # 2. Add a small value 10^(-precision-1) + # 3. Cast the result to string + # 4. Drop the extra digit from the string. To do that, we need to slice the string + # with length = digits in an integer part + 1 (symbol of ".") + precision + + if coltype.precision == 0: + return self.to_string(f"round({value})") + + precision = coltype.precision + # TODO: too complex, is there better performance way? + value = f""" + if({value} >= 0, '', '-') || left( + toString( + toDecimal128( + round(abs({value}), {precision}), + {precision} + 1 + ) + + + toDecimal128( + exp10(-{precision + 1}), + {precision} + 1 + ) + ), + toUInt8( + greatest( + floor(log10(abs({value}))) + 1, + 1 + ) + ) + 1 + {precision} + ) + """ + return value + + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + prec = coltype.precision + if coltype.rounds: + timestamp = f"toDateTime64(round(toUnixTimestamp64Micro(toDateTime64({value}, 6)) / 1000000, {prec}), 6)" + return self.to_string(timestamp) + + fractional = f"toUnixTimestamp64Micro(toDateTime64({value}, {prec})) % 1000000" + fractional = f"lpad({self.to_string(fractional)}, 6, '0')" + value = f"formatDateTime({value}, '%Y-%m-%d %H:%M:%S') || '.' || {self.to_string(fractional)}" + return f"rpad({value}, {TIMESTAMP_PRECISION_POS + 6}, '0')" + + +class Dialect(BaseDialect): + name = "Clickhouse" + ROUNDS_ON_PREC_LOSS = False + ARG_SYMBOL = None # TODO Clickhouse only supports named parameters, not positional + TYPE_CLASSES = { + "Int8": Integer, + "Int16": Integer, + "Int32": Integer, + "Int64": Integer, + "Int128": Integer, + "Int256": Integer, + "UInt8": Integer, + "UInt16": Integer, + "UInt32": Integer, + "UInt64": Integer, + "UInt128": Integer, + "UInt256": Integer, + "Float32": Float, + "Float64": Float, + "Decimal": Decimal, + "UUID": Native_UUID, + "String": Text, + "FixedString": Text, + "DateTime": Timestamp, + "DateTime64": Timestamp, + "Bool": Boolean, + } + MIXINS = {Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample} + + def quote(self, s: str) -> str: + return f'"{s}"' + + def to_string(self, s: str) -> str: + return f"toString({s})" + + def _convert_db_precision_to_digits(self, p: int) -> int: + # Done the same as for PostgreSQL but need to rewrite in another way + # because it does not help for float with a big integer part. + return super()._convert_db_precision_to_digits(p) - 2 + + def _parse_type_repr(self, type_repr: str) -> Optional[Type[ColType]]: + nullable_prefix = "Nullable(" + if type_repr.startswith(nullable_prefix): + type_repr = type_repr[len(nullable_prefix) :].rstrip(")") + + if type_repr.startswith("Decimal"): + type_repr = "Decimal" + elif type_repr.startswith("FixedString"): + type_repr = "FixedString" + elif type_repr.startswith("DateTime64"): + type_repr = "DateTime64" + + return self.TYPE_CLASSES.get(type_repr) + + # def timestamp_value(self, t: DbTime) -> str: + # # return f"'{t}'" + # return f"'{str(t)[:19]}'" + + def set_timezone_to_utc(self) -> str: + raise NotImplementedError() + + def current_timestamp(self) -> str: + return "now()" + + +class Clickhouse(ThreadedDatabase): + dialect = Dialect() + CONNECT_URI_HELP = "clickhouse://:@/" + CONNECT_URI_PARAMS = ["database?"] + + def __init__(self, *, thread_count: int, **kw): + super().__init__(thread_count=thread_count) + + self._args = kw + # In Clickhouse database and schema are the same + self.default_schema = kw.get("database", DEFAULT_DATABASE) + + def create_connection(self): + clickhouse = import_clickhouse() + + class SingleConnection(clickhouse.dbapi.connection.Connection): + """Not thread-safe connection to Clickhouse""" + + def cursor(self, cursor_factory=None): + if not len(self.cursors): + _ = super().cursor() + return self.cursors[0] + + try: + return SingleConnection(**self._args) + except clickhouse.OperationError as e: + raise ConnectError(*e.args) from e + + @property + def is_autocommit(self) -> bool: + return True diff --git a/reladiff/sqeleton/databases/databricks.py b/reladiff/sqeleton/databases/databricks.py new file mode 100644 index 00000000..9d11dd6d --- /dev/null +++ b/reladiff/sqeleton/databases/databricks.py @@ -0,0 +1,204 @@ +import math +from typing import Dict, Sequence, Optional, Tuple +import logging + +from ..abcs.database_types import ( + Integer, + Float, + Decimal, + Timestamp, + Text, + TemporalType, + NumericType, + DbPath, + ColType, + UnknownColType, + Boolean, +) +from ..abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue +from .base import ( + MD5_HEXDIGITS, + CHECKSUM_HEXDIGITS, + BaseDialect, + ThreadedDatabase, + import_helper, + parse_table_name, + Mixin_RandomSample, +) + + +@import_helper(text="You can install it using 'pip install databricks-sql-connector'") +def import_databricks(): + import databricks.sql + + return databricks + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"cast(conv(substr(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16, 10) as decimal(38, 0))" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + """Databricks timestamp contains no more than 6 digits in precision""" + + if coltype.rounds: + timestamp = f"cast(round(unix_micros({value}) / 1000000, {coltype.precision}) * 1000000 as bigint)" + return f"date_format(timestamp_micros({timestamp}), 'yyyy-MM-dd HH:mm:ss.SSSSSS')" + + precision_format = "S" * coltype.precision + "0" * (6 - coltype.precision) + return f"date_format({value}, 'yyyy-MM-dd HH:mm:ss.{precision_format}')" + + def normalize_number(self, value: str, coltype: NumericType) -> str: + value = f"cast({value} as decimal(38, {coltype.precision}))" + if coltype.precision > 0: + value = f"format_number({value}, {coltype.precision})" + return f"replace({self.to_string(value)}, ',', '')" + + def normalize_boolean(self, value: str, _coltype: Boolean) -> str: + return self.to_string(f"cast ({value} as int)") + + +class Dialect(BaseDialect): + name = "Databricks" + ROUNDS_ON_PREC_LOSS = True + TYPE_CLASSES = { + # Numbers + "INT": Integer, + "SMALLINT": Integer, + "TINYINT": Integer, + "BIGINT": Integer, + "FLOAT": Float, + "DOUBLE": Float, + "DECIMAL": Decimal, + # Timestamps + "TIMESTAMP": Timestamp, + # Text + "STRING": Text, + # Boolean + "BOOLEAN": Boolean, + } + MIXINS = {Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample} + + def quote(self, s: str): + return f"`{s}`" + + def to_string(self, s: str) -> str: + return f"cast({s} as string)" + + def _convert_db_precision_to_digits(self, p: int) -> int: + # Subtracting 2 due to wierd precision issues + return max(super()._convert_db_precision_to_digits(p) - 2, 0) + + def set_timezone_to_utc(self) -> str: + return "SET TIME ZONE 'UTC'" + + +class Databricks(ThreadedDatabase): + dialect = Dialect() + CONNECT_URI_HELP = "databricks://:@/" + CONNECT_URI_PARAMS = ["catalog", "schema"] + + def __init__(self, *, thread_count, **kw): + logging.getLogger("databricks.sql").setLevel(logging.WARNING) + + self._args = kw + self.default_schema = kw.get("schema", "default") + self.catalog = self._args.get("catalog", "hive_metastore") + super().__init__(thread_count=thread_count) + + def create_connection(self): + databricks = import_databricks() + + try: + return databricks.sql.connect( + server_hostname=self._args["server_hostname"], + http_path=self._args["http_path"], + access_token=self._args["access_token"], + catalog=self.catalog, + ) + except databricks.sql.exc.Error as e: + raise ConnectionError(*e.args) from e + + def query_table_schema(self, path: DbPath) -> Dict[str, tuple]: + # Databricks has INFORMATION_SCHEMA only for Databricks Runtime, not for Databricks SQL. + # https://docs.databricks.com/spark/latest/spark-sql/language-manual/information-schema/columns.html + # So, to obtain information about schema, we should use another approach. + + conn = self.create_connection() + + catalog, schema, table = self._normalize_table_path(path) + with conn.cursor() as cursor: + cursor.columns(catalog_name=catalog, schema_name=schema, table_name=table) + try: + rows = cursor.fetchall() + finally: + conn.close() + if not rows: + raise RuntimeError(f"{self.name}: Table '{'.'.join(path)}' does not exist, or has no columns") + + d = {r.COLUMN_NAME: (r.COLUMN_NAME, r.TYPE_NAME, r.DECIMAL_DIGITS, None, None) for r in rows} + assert len(d) == len(rows) + return d + + def _process_table_schema( + self, path: DbPath, raw_schema: Dict[str, tuple], filter_columns: Sequence[str], where: str = None + ): + accept = {i.lower() for i in filter_columns} + rows = [row for name, row in raw_schema.items() if name.lower() in accept] + + resulted_rows = [] + for row in rows: + row_type = "DECIMAL" if row[1].startswith("DECIMAL") else row[1] + type_cls = self.dialect.TYPE_CLASSES.get(row_type, UnknownColType) + + if issubclass(type_cls, Integer): + row = (row[0], row_type, None, None, 0) + + elif issubclass(type_cls, Float): + numeric_precision = math.ceil(row[2] / math.log(2, 10)) + row = (row[0], row_type, None, numeric_precision, None) + + elif issubclass(type_cls, Decimal): + items = row[1][8:].rstrip(")").split(",") + numeric_precision, numeric_scale = int(items[0]), int(items[1]) + row = (row[0], row_type, None, numeric_precision, numeric_scale) + + elif issubclass(type_cls, Timestamp): + row = (row[0], row_type, row[2], None, None) + + else: + row = (row[0], row_type, None, None, None) + + resulted_rows.append(row) + + col_dict: Dict[str, ColType] = {row[0]: self.dialect.parse_type(path, *row) for row in resulted_rows} + + self._refine_coltypes(path, col_dict, where) + return col_dict + + def process_query_table_schema(self, path: DbPath, raw_schema: Dict[str, Tuple], refine: bool = True, refine_where: Optional[str] = None) -> Tuple[Dict[str, ColType], Optional[list]]: + if not refine: + raise NotImplementedError() + return self._process_table_schema(path, raw_schema, list(raw_schema), refine_where), None + + 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 + + def _normalize_table_path(self, path: DbPath) -> DbPath: + if len(path) == 1: + return self.catalog, self.default_schema, path[0] + elif len(path) == 2: + return self.catalog, path[0], path[1] + elif len(path) == 3: + return path + + raise ValueError( + f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected format: table, schema.table, or catalog.schema.table" + ) diff --git a/reladiff/sqeleton/databases/dremio.py b/reladiff/sqeleton/databases/dremio.py new file mode 100644 index 00000000..8e6829b9 --- /dev/null +++ b/reladiff/sqeleton/databases/dremio.py @@ -0,0 +1,207 @@ +from math import isnan +from typing import Optional, Union + +from sqeleton.queries import this, SKIP +from .base import BaseDialect, QueryResult, import_helper, parse_table_name, ThreadLocalInterpreter, \ + Mixin_RandomSample, logger, ThreadedDatabase, Mixin_Schema, TIMESTAMP_PRECISION_POS +from ..abcs import Compilable +from ..abcs.database_types import ( + Timestamp, + Integer, + Float, + Text, + FractionalType, + Date, + DbPath, + DbTime, + Decimal, + ColType, + TemporalType, + Boolean, + ColType_UUID +) +from ..abcs.mixins import AbstractMixin_NormalizeValue, AbstractMixin_MD5 +from ..queries.compiler import CompiledCode + + +def query_cursor(c, sql_code: CompiledCode) -> Optional[QueryResult]: + logger.debug(f"[Dremio] Executing SQL: {sql_code.code} || {sql_code.args}") + c.execute(sql_code.code, sql_code.args) + + columns = c.description and [col[0] for col in c.description] + return QueryResult(c.fetchall(), columns) + + +@import_helper("dremio") +def import_dremio(): + import sqlalchemy + + return sqlalchemy + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"CAST(CONV(SUBSTR(MD5({s}), 18), 16, 10) AS BIGINT)" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + + def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str: + # Trim doesn't work on CHAR type + return f"TRIM(CAST({value} AS VARCHAR))" + + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + """Dremio timestamp contains no more than 6 digits in precision""" + precision_format = ('F' * coltype.precision) if coltype.precision > 0 else '' + return f"rpad(TO_CHAR(CAST({value} AS TIMESTAMP), 'YYYY-MM-DD HH24:MI:SS.{precision_format}'), {TIMESTAMP_PRECISION_POS + 6}, '0')" + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + if coltype.precision > 0: + value = f"format_number({value}, {coltype.precision})" + return f"replace({self.to_string(value)}, ',', '')" + + def normalize_boolean(self, value: str, _coltype: Boolean) -> str: + return self.to_string(f"cast({value} as int)") + +class Dialect(BaseDialect, Mixin_Schema): + name = "Dremio" + ROUNDS_ON_PREC_LOSS = False # False if it truncates. + ARG_SYMBOL = None # Not implemented by Dremio + TYPE_CLASSES = { + # https://docs.dremio.com/current/reference/sql/data-types/ + # com.dremio.dac.explore.DataTypeUtil + "BOOLEAN": Boolean, + "INT": Integer, + "INTEGER": Integer, + "TINYINT": Integer, # An alias for Integer. + "SMALLINT": Integer, # An alias for Integer. + "BIGINT": Integer, + "DECIMAL": Decimal, + "DEC": Decimal, # An alias for Decimal. + "NUMERIC": Decimal, # An alias for Decimal. + "FLOAT": Float, + "DOUBLE": Float, + "REAL": Float, + "DATE": Date, + "TIME": Timestamp, + "TIMESTAMP": Timestamp, # Dremio does not support timezone-aware timestamps. + "CHAR": Text, + "CHARACTER": Text, + "VARCHAR": Text, + "CHARACTER VARYING": Text, # An alias for VARCHAR. + "VARBINARY": Text, + "BINARY VARYING": Text, # An alias for VARBINARY. + } + MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample} + + def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable: + """Query to select the list of tables in the schema. (query return type: table[str]) + + If 'like' is specified, the value is applied to the table name, using the 'like' operator. + """ + return ( + self.table_information() + .where( + this.table_schema == table_schema, + this.table_name.like(like) if like is not None else SKIP, + ) + .select(this.table_name) + ) + + def timestamp_value(self, t: DbTime) -> str: + s = t.isoformat(' ', 'milliseconds') + return f"cast('{s}' as timestamp)" + + def quote(self, s: str): + return f'"{s}"' + + def to_string(self, s: str): + return f"cast({s} as varchar)" + + @staticmethod + def nan_to_none(value: Optional[object]) -> Optional[int]: + if value is not None and not isinstance(value, int): + return None if isnan(value) else int(value) + else: + return value + + def parse_type( + self, + table_path: DbPath, + col_name: str, + type_repr: str, + datetime_precision: int = None, + numeric_precision: int = None, + numeric_scale: int = None, + ) -> ColType: + datetime_precision = self.nan_to_none(datetime_precision) + numeric_precision = self.nan_to_none(numeric_precision) + numeric_scale = self.nan_to_none(numeric_scale) + + return super().parse_type(table_path, col_name, type_repr, datetime_precision, numeric_precision, numeric_scale) + + def set_timezone_to_utc(self) -> str: + """Dremio does not support setting session timezone to UTC. Dremio retrieves the time or timestamp value + with the assumption that the time zone is in Coordinated Universal Time (UTC). Use CONVERT_TIMEZONE in + queries instead. + """ + pass + + + +class Dremio(ThreadedDatabase): + dialect = Dialect() + CONNECT_URI_HELP = ("dremio://[:]@:/?" + "Token=&" + "UseEncryption=&" + "DisableCertificateVerification=") + CONNECT_URI_PARAMS = ['space?', 'Token?', 'UseEncryption?', 'DisableCertificateVerification?'] + + def __init__(self, thread_count, **kw): + self._args = kw + if 'space' in self._args: + self.default_schema = self._args['space'] + super().__init__(thread_count=thread_count) + + def create_connection(self): + # Dremio PAT token is preferred over password. It must be URL-encoded as it will contain special characters. + password_or_token = self._args.get('Token') if self._args.get('Token') else self._args.get('password') + user_pw_str = f"{self._args.get('user')}:{password_or_token}@" + keywords = [] + if 'UseEncryption' in self._args: + keywords.append(f"UseEncryption={self._args['UseEncryption']}") + if 'DisableCertificateVerification' in self._args: + keywords.append(f"DisableCertificateVerification={self._args['DisableCertificateVerification']}") + keyword_str = '?' + '&'.join(keywords) if keywords else '' + + connection_string = (f"dremio+flight://{user_pw_str}{self._args['host']}:{self._args['port']}/" + f"{self.default_schema}{keyword_str}") + + dremiodb = import_dremio() + engine: dremiodb.Engine = dremiodb.create_engine(url=connection_string, echo=True) + + try: + connection: dremiodb.engine.Connection = engine.connect() + except dremiodb.exc.SQLAlchemyError as e: + raise ConnectionError(*e.args) from e + + return connection.connection + + def _query(self, sql_code: Union[str, ThreadLocalInterpreter]): + "Uses the standard SQL cursor interface" + from pyarrow._flight import FlightUnauthenticatedError + try: + return super()._query(sql_code=sql_code) + except FlightUnauthenticatedError as e: + raise ValueError( + 'Possible reasons - Your TOKEN may have expired in the Dremio -> Account Settings -> Personal Access ' + 'Token page. Or you have not provided a username that matches the Personal Access Token: {e}' + ) from e + + 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 diff --git a/reladiff/sqeleton/databases/duckdb.py b/reladiff/sqeleton/databases/duckdb.py new file mode 100644 index 00000000..154997db --- /dev/null +++ b/reladiff/sqeleton/databases/duckdb.py @@ -0,0 +1,193 @@ +from typing import Union + +from ..utils import match_regexps +from ..abcs.database_types import ( + Timestamp, + TimestampTZ, + DbPath, + ColType, + Float, + Decimal, + Integer, + TemporalType, + Native_UUID, + Text, + FractionalType, + Boolean, + AbstractTable, +) +from ..abcs.mixins import ( + AbstractMixin_MD5, + AbstractMixin_NormalizeValue, + AbstractMixin_RandomSample, + AbstractMixin_Regex, +) +from .base import ( + Database, + BaseDialect, + import_helper, + ConnectError, + ThreadLocalInterpreter, + TIMESTAMP_PRECISION_POS, +) +from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, Mixin_Schema +from ..queries.ast_classes import Func, Compilable +from ..queries.api import code + + +@import_helper("duckdb") +def import_duckdb(): + import duckdb + + return duckdb + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"('0x' || SUBSTRING(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS},{CHECKSUM_HEXDIGITS}))::BIGINT" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + # It's precision 6 by default. If precision is less than 6 -> we remove the trailing numbers. + if coltype.rounds and coltype.precision > 0: + return f"CONCAT(SUBSTRING(STRFTIME({value}::TIMESTAMP, '%Y-%m-%d %H:%M:%S.'),1,23), LPAD(((ROUND(strftime({value}::timestamp, '%f')::DECIMAL(15,7)/100000,{coltype.precision-1})*100000)::INT)::VARCHAR,6,'0'))" + + return f"rpad(substring(strftime({value}::timestamp, '%Y-%m-%d %H:%M:%S.%f'),1,{TIMESTAMP_PRECISION_POS+coltype.precision}),26,'0')" + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + return self.to_string(f"{value}::DECIMAL(38, {coltype.precision})") + + def normalize_boolean(self, value: str, _coltype: Boolean) -> str: + return self.to_string(f"{value}::INTEGER") + + +class Mixin_RandomSample(AbstractMixin_RandomSample): + def random_sample_n(self, tbl: AbstractTable, size: int) -> AbstractTable: + return code("SELECT * FROM ({tbl}) USING SAMPLE {size};", tbl=tbl, size=size) + + def random_sample_ratio_approx(self, tbl: AbstractTable, ratio: float) -> AbstractTable: + return code("SELECT * FROM ({tbl}) USING SAMPLE {percent}%;", tbl=tbl, percent=int(100 * ratio)) + + +class Mixin_Regex(AbstractMixin_Regex): + def test_regex(self, string: Compilable, pattern: Compilable) -> Compilable: + return Func("regexp_matches", [string, pattern]) + + +class Dialect(BaseDialect, Mixin_Schema): + name = "DuckDB" + ROUNDS_ON_PREC_LOSS = False + SUPPORTS_PRIMARY_KEY = True + SUPPORTS_INDEXES = True + MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample} + ARG_SYMBOL = "?" + + TYPE_CLASSES = { + # Timestamps + "TIMESTAMP WITH TIME ZONE": TimestampTZ, + "TIMESTAMP": Timestamp, + # Numbers + "DOUBLE": Float, + "FLOAT": Float, + "DECIMAL": Decimal, + "INTEGER": Integer, + "BIGINT": Integer, + # Text + "VARCHAR": Text, + "TEXT": Text, + # UUID + "UUID": Native_UUID, + # Bool + "BOOLEAN": Boolean, + } + + def quote(self, s: str): + return f'"{s}"' + + def to_string(self, s: str): + return f"{s}::VARCHAR" + + def _convert_db_precision_to_digits(self, p: int) -> int: + # Subtracting 2 due to wierd precision issues in PostgreSQL + return super()._convert_db_precision_to_digits(p) - 2 + + def parse_type( + self, + table_path: DbPath, + col_name: str, + type_repr: str, + datetime_precision: int = None, + numeric_precision: int = None, + numeric_scale: int = None, + ) -> ColType: + regexps = { + r"DECIMAL\((\d+),(\d+)\)": Decimal, + } + + for m, t_cls in match_regexps(regexps, type_repr): + precision = int(m.group(2)) + return t_cls(precision=precision) + + return super().parse_type(table_path, col_name, type_repr, datetime_precision, numeric_precision, numeric_scale) + + def set_timezone_to_utc(self) -> str: + return "SET GLOBAL TimeZone='UTC'" + + def current_timestamp(self) -> str: + return "current_timestamp" + + +class DuckDB(Database): + dialect = Dialect() + SUPPORTS_UNIQUE_CONSTAINT = False # Temporary, until we implement it + default_schema = "main" + CONNECT_URI_HELP = "duckdb://@" + CONNECT_URI_PARAMS = ["database", "dbpath"] + + def __init__(self, **kw): + self._args = kw + self._conn = self.create_connection() + + @property + def is_autocommit(self) -> bool: + return True + + def _query(self, sql_code: Union[str, ThreadLocalInterpreter]): + "Uses the standard SQL cursor interface" + return self._query_conn(self._conn, sql_code) + + def close(self): + super().close() + self._conn.close() + + def create_connection(self): + ddb = import_duckdb() + try: + return ddb.connect(self._args["filepath"]) + except ddb.OperationalError as e: + raise ConnectError(*e.args) from e + + def select_table_schema(self, path: DbPath) -> str: + database, schema, table = self._normalize_table_path(path) + + info_schema_path = ["information_schema", "columns"] + if database: + info_schema_path.insert(0, database) + + return ( + f"SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale FROM {'.'.join(info_schema_path)} " + f"WHERE table_name = '{table}' AND table_schema = '{schema}'" + ) + + def _normalize_table_path(self, path: DbPath) -> DbPath: + if len(path) == 1: + return None, self.default_schema, path[0] + elif len(path) == 2: + return None, path[0], path[1] + elif len(path) == 3: + return path + + raise ValueError( + f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected format: table, schema.table, or database.schema.table" + ) diff --git a/reladiff/sqeleton/databases/mssql.py b/reladiff/sqeleton/databases/mssql.py new file mode 100644 index 00000000..8d394e3c --- /dev/null +++ b/reladiff/sqeleton/databases/mssql.py @@ -0,0 +1,25 @@ +# class MsSQL(ThreadedDatabase): +# "AKA sql-server" + +# def __init__(self, host, port, user, password, *, database, thread_count, **kw): +# args = dict(server=host, port=port, database=database, user=user, password=password, **kw) +# self._args = {k: v for k, v in args.items() if v is not None} + +# super().__init__(thread_count=thread_count) + +# def create_connection(self): +# mssql = import_mssql() +# try: +# return mssql.connect(**self._args) +# except mssql.Error as e: +# raise ConnectError(*e.args) from e + +# def quote(self, s: str): +# return f"[{s}]" + +# def md5_as_int(self, s: str) -> str: +# return f"CONVERT(decimal(38,0), CONVERT(bigint, HashBytes('MD5', {s}), 2))" +# # return f"CONVERT(bigint, (CHECKSUM({s})))" + +# def to_string(self, s: str): +# return f"CONVERT(varchar, {s})" diff --git a/reladiff/sqeleton/databases/mysql.py b/reladiff/sqeleton/databases/mysql.py new file mode 100644 index 00000000..75522fec --- /dev/null +++ b/reladiff/sqeleton/databases/mysql.py @@ -0,0 +1,150 @@ +from ..abcs.database_types import ( + Datetime, + Timestamp, + Float, + Decimal, + Integer, + Text, + TemporalType, + FractionalType, + ColType_UUID, + Boolean, + Date, +) +from ..abcs.mixins import ( + AbstractMixin_MD5, + AbstractMixin_NormalizeValue, + AbstractMixin_Regex, + AbstractMixin_RandomSample, +) +from .base import Mixin_OptimizerHints, ThreadedDatabase, import_helper, ConnectError, BaseDialect, Compilable +from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, TIMESTAMP_PRECISION_POS, Mixin_Schema, Mixin_RandomSample +from ..queries.ast_classes import BinBoolOp + + +@import_helper("mysql") +def import_mysql(): + import mysql.connector + + return mysql.connector + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"cast(conv(substring(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16, 10) as unsigned)" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + if coltype.rounds: + return self.to_string(f"cast( cast({value} as datetime({coltype.precision})) as datetime(6))") + + s = self.to_string(f"cast({value} as datetime(6))") + return f"RPAD(RPAD({s}, {TIMESTAMP_PRECISION_POS+coltype.precision}, '.'), {TIMESTAMP_PRECISION_POS+6}, '0')" + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + return self.to_string(f"cast({value} as decimal(38, {coltype.precision}))") + + def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str: + return f"TRIM(CAST({value} AS char))" + + +class Mixin_Regex(AbstractMixin_Regex): + def test_regex(self, string: Compilable, pattern: Compilable) -> Compilable: + return BinBoolOp("REGEXP", [string, pattern]) + + +class Dialect(BaseDialect, Mixin_Schema, Mixin_OptimizerHints): + name = "MySQL" + ROUNDS_ON_PREC_LOSS = True + SUPPORTS_PRIMARY_KEY = True + SUPPORTS_INDEXES = True + + TYPE_CLASSES = { + # Dates + "datetime": Datetime, + "timestamp": Timestamp, + "date": Date, + # Numbers + "double": Float, + "float": Float, + "decimal": Decimal, + "int": Integer, + "bigint": Integer, + "smallint": Integer, + "tinyint": Integer, + "mediumint": Integer, + # Text + "varchar": Text, + "char": Text, + "varbinary": Text, + "binary": Text, + "text": Text, + "mediumtext": Text, + "longtext": Text, + "tinytext": Text, + # Boolean + "boolean": Boolean, + } + MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample} + + def quote(self, s: str): + return f"`{s}`" + + def to_string(self, s: str): + return f"cast({s} as char)" + + def is_distinct_from(self, a: str, b: str) -> str: + return f"not ({a} <=> {b})" + + def random(self) -> str: + return "RAND()" + + def type_repr(self, t) -> str: + if isinstance(t, type): + try: + return { + str: "VARCHAR(1024)", + }[t] + except KeyError: + pass + return super().type_repr(t) + + def explain_as_text(self, query: str) -> str: + return f"EXPLAIN FORMAT=TREE {query}" + + def optimizer_hints(self, s: str): + return f"/*+ {s} */ " + + def set_timezone_to_utc(self) -> str: + return "SET @@session.time_zone='+00:00'" + + +class MySQL(ThreadedDatabase): + dialect = Dialect() + SUPPORTS_ALPHANUMS = False + SUPPORTS_UNIQUE_CONSTAINT = True + CONNECT_URI_HELP = "mysql://:@/" + CONNECT_URI_PARAMS = ["database?"] + + def __init__(self, *, thread_count, **kw): + self._args = kw + + super().__init__(thread_count=thread_count) + + # In MySQL schema and database are synonymous + try: + self.default_schema = kw["database"] + except KeyError: + raise ValueError("MySQL URL must specify a database") + + def create_connection(self): + mysql = import_mysql() + try: + return mysql.connect(charset="utf8", use_unicode=True, **self._args) + except mysql.Error as e: + if e.errno == mysql.errorcode.ER_ACCESS_DENIED_ERROR: + raise ConnectError("Bad user name or password") from e + elif e.errno == mysql.errorcode.ER_BAD_DB_ERROR: + raise ConnectError("Database does not exist") from e + raise ConnectError(*e.args) from e diff --git a/reladiff/sqeleton/databases/oracle.py b/reladiff/sqeleton/databases/oracle.py new file mode 100644 index 00000000..913b393a --- /dev/null +++ b/reladiff/sqeleton/databases/oracle.py @@ -0,0 +1,220 @@ +from typing import Dict, List, Optional + +from ..queries.compiler import CompiledCode +from ..utils import match_regexps +from ..abcs.database_types import ( + Decimal, + Float, + Text, + DbPath, + TemporalType, + ColType, + DbTime, + ColType_UUID, + Timestamp, + TimestampTZ, + FractionalType, +) +from ..queries.ast_classes import ForeignKey +from ..abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue, AbstractMixin_Schema +from ..abcs import Compilable +from ..queries import this, table, SKIP +from .base import ( + BaseDialect, + Mixin_OptimizerHints, + ThreadedDatabase, + import_helper, + ConnectError, + QueryError, + Mixin_RandomSample, +) +from .base import TIMESTAMP_PRECISION_POS + +SESSION_TIME_ZONE = None # Changed by the tests + + +@import_helper("oracle") +def import_oracle(): + import cx_Oracle + + return cx_Oracle + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + # standard_hash is faster than DBMS_CRYPTO.Hash + # TODO: Find a way to use UTL_RAW.CAST_TO_BINARY_INTEGER ? + return f"to_number(substr(standard_hash({s}, 'MD5'), 18), 'xxxxxxxxxxxxxxx')" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str: + # Cast is necessary for correct MD5 (trimming not enough) + return f"CAST(TRIM({value}) AS VARCHAR(36))" + + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + if coltype.rounds: + return f"to_char(cast({value} as timestamp({coltype.precision})), 'YYYY-MM-DD HH24:MI:SS.FF6')" + + if coltype.precision > 0: + truncated = f"to_char({value}, 'YYYY-MM-DD HH24:MI:SS.FF{coltype.precision}')" + else: + truncated = f"to_char({value}, 'YYYY-MM-DD HH24:MI:SS.')" + return f"RPAD({truncated}, {TIMESTAMP_PRECISION_POS+6}, '0')" + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + # FM999.9990 + format_str = "FM" + "9" * (38 - coltype.precision) + if coltype.precision: + 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") + .where( + this.OWNER == table_schema, + this.TABLE_NAME.like(like) if like is not None else SKIP, + ) + .select(table_name=this.TABLE_NAME) + ) + + +class Dialect(BaseDialect, Mixin_Schema, Mixin_OptimizerHints): + name = "Oracle" + SUPPORTS_PRIMARY_KEY = True + SUPPORTS_INDEXES = True + TYPE_CLASSES: Dict[str, type] = { + "NUMBER": Decimal, + "FLOAT": Float, + # Text + "CHAR": Text, + "NCHAR": Text, + "NVARCHAR2": Text, + "VARCHAR2": Text, + "DATE": Timestamp, + } + ROUNDS_ON_PREC_LOSS = True + PLACEHOLDER_TABLE = "DUAL" + MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample} + + def quote(self, s: str): + return f'"{s}"' + + def to_string(self, s: str): + return f"cast({s} as varchar(1024))" + + def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None): + if offset: + raise NotImplementedError("No support for OFFSET in query") + + return f"FETCH NEXT {limit} ROWS ONLY" + + def concat(self, items: List[str]) -> str: + joined_exprs = " || ".join(items) + return f"({joined_exprs})" + + def timestamp_value(self, t: DbTime) -> str: + return "timestamp '%s'" % t.isoformat(" ") + + def random(self) -> str: + return "dbms_random.value" + + def is_distinct_from(self, a: str, b: str) -> str: + return f"DECODE({a}, {b}, 1, 0) = 0" + + def type_repr(self, t) -> str: + if isinstance(t, ForeignKey): + return self.type_repr(t.type) + try: + return { + str: "VARCHAR(1024)", + }[t] + except KeyError: + return super().type_repr(t) + + def constant_values(self, rows) -> str: + return " UNION ALL ".join("SELECT %s FROM DUAL" % ", ".join(row) for row in rows) + + def explain_as_text(self, query: str) -> str: + raise NotImplementedError("Explain not yet implemented in Oracle") + + def parse_type( + self, + table_path: DbPath, + col_name: str, + type_repr: str, + datetime_precision: int = None, + numeric_precision: int = None, + numeric_scale: int = None, + ) -> ColType: + regexps = { + r"TIMESTAMP\((\d)\) WITH LOCAL TIME ZONE": Timestamp, + r"TIMESTAMP\((\d)\) WITH TIME ZONE": TimestampTZ, + r"TIMESTAMP\((\d)\)": Timestamp, + } + + for m, t_cls in match_regexps(regexps, type_repr): + precision = int(m.group(1)) + return t_cls(precision=precision, rounds=self.ROUNDS_ON_PREC_LOSS) + + return super().parse_type(table_path, col_name, type_repr, datetime_precision, numeric_precision, numeric_scale) + + def set_timezone_to_utc(self) -> str: + return "ALTER SESSION SET TIME_ZONE = 'UTC'" + + def current_timestamp(self) -> str: + return "LOCALTIMESTAMP" + + +class Oracle(ThreadedDatabase): + dialect = Dialect() + CONNECT_URI_HELP = "oracle://:@:port/" + CONNECT_URI_PARAMS = ["database?"] + + def __init__(self, *, host, database, thread_count, port=None, **kw): + self.kwargs = kw + + # Build dsn if not present + if "dsn" not in kw: + # Support for different ports + port = port or 1521 + self.kwargs["dsn"] = f"{host}:{port}/{database}" if database else f"{host}:{port}" + + self.default_schema = kw.get("user").upper() + + super().__init__(thread_count=thread_count) + + def create_connection(self): + self._oracle = import_oracle() + try: + c = self._oracle.connect(**self.kwargs) + if SESSION_TIME_ZONE: + c.cursor().execute(f"ALTER SESSION SET TIME_ZONE = '{SESSION_TIME_ZONE}'") + return c + except Exception as e: + raise ConnectError(*e.args) from e + + def _query_cursor(self, c, sql_code: CompiledCode): + # Convert %s style parameters to :1, :2, :3 style for Oracle to support queries built by sqeleton (tbl.create()) + if sql_code.args: + # Replace %s with :1, :2, :3, etc. + oracle_sql = sql_code.code + for i in range(len(sql_code.args)): + oracle_sql = oracle_sql.replace('%s', f':{i+1}', 1) + sql_code = CompiledCode(oracle_sql, sql_code.args, sql_code.type) + + try: + return super()._query_cursor(c, sql_code) + except self._oracle.DatabaseError as e: + raise QueryError(e) + + def select_table_schema(self, path: DbPath) -> str: + schema, name = self._normalize_table_path(path) + + return ( + f"SELECT column_name, data_type, 6 as datetime_precision, data_precision as numeric_precision, data_scale as numeric_scale" + f" FROM ALL_TAB_COLUMNS WHERE table_name = '{name}' AND owner = '{schema}'" + ) diff --git a/reladiff/sqeleton/databases/postgresql.py b/reladiff/sqeleton/databases/postgresql.py new file mode 100644 index 00000000..154cf651 --- /dev/null +++ b/reladiff/sqeleton/databases/postgresql.py @@ -0,0 +1,186 @@ +from typing import List +from ..abcs.database_types import ( + DbPath, + Timestamp, + TimestampTZ, + Float, + Decimal, + Integer, + TemporalType, + Native_UUID, + Text, + FractionalType, + Boolean, + Date, +) +from ..abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue +from .base import BaseDialect, ThreadedDatabase, import_helper, ConnectError, Mixin_Schema +from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, _CHECKSUM_BITSIZE, TIMESTAMP_PRECISION_POS, Mixin_RandomSample +from ..schema import _Field + +SESSION_TIME_ZONE = None # Changed by the tests + + +@import_helper("postgresql") +def import_postgresql(): + import psycopg2 + import psycopg2.extras + + psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select) + return psycopg2 + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"('x' || substring(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}))::bit({_CHECKSUM_BITSIZE})::bigint" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + if coltype.rounds: + return f"to_char({value}::timestamp({coltype.precision}), 'YYYY-mm-dd HH24:MI:SS.US')" + + timestamp6 = f"to_char({value}::timestamp(6), 'YYYY-mm-dd HH24:MI:SS.US')" + return ( + f"RPAD(LEFT({timestamp6}, {TIMESTAMP_PRECISION_POS+coltype.precision}), {TIMESTAMP_PRECISION_POS+6}, '0')" + ) + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + return self.to_string(f"{value}::decimal(38, {coltype.precision})") + + def normalize_boolean(self, value: str, _coltype: Boolean) -> str: + return self.to_string(f"{value}::int") + + +class PostgresqlDialect(BaseDialect, Mixin_Schema): + name = "PostgreSQL" + ROUNDS_ON_PREC_LOSS = True + SUPPORTS_PRIMARY_KEY = True + SUPPORTS_INDEXES = True + MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample} + + TYPE_CLASSES = { + # Timestamps + "timestamp with time zone": TimestampTZ, + "timestamp without time zone": Timestamp, + "timestamp": Timestamp, + "date": Date, + # Numbers + "double precision": Float, + "real": Float, + "decimal": Decimal, + "smallint": Integer, + "integer": Integer, + "numeric": Decimal, + "bigint": Integer, + # Text + "character": Text, + "character varying": Text, + "varchar": Text, + "text": Text, + # UUID + "uuid": Native_UUID, + # Boolean + "boolean": Boolean, + } + + def quote(self, s: str): + return f'"{s}"' + + def to_string(self, s: str): + return f"{s}::varchar" + + def concat(self, items: List[str]) -> str: + # Postgres does not allow functions that have more that have more than 100 arguments. + # When using the concat function, this limits comparisons to less than 50 fields. + # Using || for concat fixes this. + joined_exprs = " || ".join(items) + return f"({joined_exprs})" + + def _convert_db_precision_to_digits(self, p: int) -> int: + # Subtracting 2 due to wierd precision issues in PostgreSQL + return super()._convert_db_precision_to_digits(p) - 2 + + def set_timezone_to_utc(self) -> str: + return "SET TIME ZONE 'UTC'" + + def current_timestamp(self) -> str: + return "current_timestamp" + + def type_repr(self, t) -> str: + if isinstance(t, _Field): + if t.options.auto: + assert t.type is int + return "SERIAL" + else: + t = t.type + + if isinstance(t, TimestampTZ): + return f"timestamp ({t.precision}) with time zone" + elif t in (dict, list): + return "json" + return super().type_repr(t) + + +class PostgreSQL(ThreadedDatabase): + dialect = PostgresqlDialect() + SUPPORTS_UNIQUE_CONSTAINT = True + CONNECT_URI_HELP = "postgresql://:@/" + CONNECT_URI_PARAMS = ["database?"] + + default_schema = "public" + + def __init__(self, *, thread_count, **kw): + self._args = kw + + super().__init__(thread_count=thread_count) + + def create_connection(self): + if not self._args: + self._args["host"] = None # psycopg2 requires 1+ arguments + + pg = import_postgresql() + try: + c = pg.connect(**self._args) + if SESSION_TIME_ZONE: + c.cursor().execute(f"SET TIME ZONE '{SESSION_TIME_ZONE}'") + return c + except pg.OperationalError as e: + raise ConnectError(*e.args) from e + + def select_table_schema(self, path: DbPath) -> str: + database, schema, table = self._normalize_table_path(path) + + info_schema_path = ["information_schema", "columns"] + if database: + info_schema_path.insert(0, database) + + return ( + f"SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale FROM {'.'.join(info_schema_path)} " + f"WHERE table_name = '{table}' AND table_schema = '{schema}'" + ) + + def select_table_unique_columns(self, path: DbPath) -> str: + database, schema, table = self._normalize_table_path(path) + + info_schema_path = ["information_schema", "key_column_usage"] + if database: + info_schema_path.insert(0, database) + + return ( + "SELECT column_name " + f"FROM {'.'.join(info_schema_path)} " + f"WHERE table_name = '{table}' AND table_schema = '{schema}'" + ) + + def _normalize_table_path(self, path: DbPath) -> DbPath: + if len(path) == 1: + return None, self.default_schema, path[0] + elif len(path) == 2: + return None, path[0], path[1] + elif len(path) == 3: + return path + + raise ValueError( + f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected format: table, schema.table, or database.schema.table" + ) diff --git a/reladiff/sqeleton/databases/presto.py b/reladiff/sqeleton/databases/presto.py new file mode 100644 index 00000000..606169de --- /dev/null +++ b/reladiff/sqeleton/databases/presto.py @@ -0,0 +1,208 @@ +from functools import partial +import re +from typing import Optional + +from sqeleton.queries.ast_classes import ForeignKey + +from ..utils import match_regexps + +from ..abcs.database_types import ( + Timestamp, + TimestampTZ, + Integer, + Float, + Text, + FractionalType, + DbPath, + DbTime, + Decimal, + ColType, + ColType_UUID, + TemporalType, + Boolean, + Native_UUID, +) +from ..abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue +from .base import BaseDialect, Database, QueryResult, import_helper, ThreadLocalInterpreter, Mixin_Schema, Mixin_RandomSample, SqlCode, logger +from .base import ( + MD5_HEXDIGITS, + CHECKSUM_HEXDIGITS, + TIMESTAMP_PRECISION_POS, +) +from ..queries.compiler import CompiledCode + + +def query_cursor(c, sql_code: CompiledCode) -> Optional[QueryResult]: + logger.debug(f"[Presto] Executing SQL: {sql_code.code} || {sql_code.args}") + c.execute(sql_code.code, sql_code.args) + + columns = c.description and [col[0] for col in c.description] + return QueryResult(c.fetchall(), columns) + + +@import_helper("presto") +def import_presto(): + import prestodb + + return prestodb + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"cast(from_base(substr(to_hex(md5(to_utf8({s}))), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16) as decimal(38, 0))" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str: + # Trim doesn't work on CHAR type + return f"TRIM(CAST({value} AS VARCHAR))" + + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + # TODO rounds + if coltype.rounds: + s = f"date_format(cast({value} as timestamp(6)), '%Y-%m-%d %H:%i:%S.%f')" + else: + s = f"date_format(cast({value} as timestamp(6)), '%Y-%m-%d %H:%i:%S.%f')" + + return f"RPAD(RPAD({s}, {TIMESTAMP_PRECISION_POS+coltype.precision}, '.'), {TIMESTAMP_PRECISION_POS+6}, '0')" + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + return self.to_string(f"cast({value} as decimal(38,{coltype.precision}))") + + def normalize_boolean(self, value: str, _coltype: Boolean) -> str: + return self.to_string(f"cast ({value} as int)") + + +class Dialect(BaseDialect, Mixin_Schema): + name = "Presto" + ROUNDS_ON_PREC_LOSS = True + ARG_SYMBOL = None # Not implemented by Presto + TYPE_CLASSES = { + # Timestamps + "timestamp with time zone": TimestampTZ, + "timestamp without time zone": Timestamp, + "timestamp": Timestamp, + # Numbers + "integer": Integer, + "bigint": Integer, + "real": Float, + "double": Float, + # Text + "varchar": Text, + # Boolean + "boolean": Boolean, + # UUID + "uuid": Native_UUID, + } + MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample} + + def explain_as_text(self, query: str) -> str: + return f"EXPLAIN (FORMAT TEXT) {query}" + + def timestamp_value(self, t: DbTime) -> str: + return f"timestamp '{t.isoformat(' ')}'" + + def quote(self, s: str): + return f'"{s}"' + + def to_string(self, s: str): + return f"cast({s} as varchar)" + + def parse_type( + self, + table_path: DbPath, + col_name: str, + type_repr: str, + datetime_precision: int = None, + numeric_precision: int = None, + _numeric_scale: int = None, + ) -> ColType: + timestamp_regexps = { + r"timestamp\((\d)\)": Timestamp, + r"timestamp\((\d)\) with time zone": TimestampTZ, + } + for m, t_cls in match_regexps(timestamp_regexps, type_repr): + precision = int(m.group(1)) + return t_cls(precision=precision, rounds=self.ROUNDS_ON_PREC_LOSS) + + number_regexps = {r"decimal\((\d+),(\d+)\)": Decimal} + for m, n_cls in match_regexps(number_regexps, type_repr): + _prec, scale = map(int, m.groups()) + return n_cls(scale) + + string_regexps = {r"varchar\((\d+)\)": Text, r"char\((\d+)\)": Text} + for m, n_cls in match_regexps(string_regexps, type_repr): + return n_cls() + + return super().parse_type(table_path, col_name, type_repr, datetime_precision, numeric_precision) + + def set_timezone_to_utc(self) -> str: + # return "SET TIME ZONE '+00:00'" + raise NotImplementedError("TODO") + + def current_timestamp(self) -> str: + return "current_timestamp" + + def type_repr(self, t) -> str: + if isinstance(t, TimestampTZ): + return f"timestamp with time zone" + elif isinstance(t, ForeignKey): + return self.type_repr(t.type) + try: + return {float: "REAL"}[t] + except KeyError: + pass + + return super().type_repr(t) + + +class Presto(Database): + dialect = Dialect() + CONNECT_URI_HELP = "presto://@//" + CONNECT_URI_PARAMS = ["catalog", "schema"] + + default_schema = "public" + + def __init__(self, **kw): + prestodb = import_presto() + + if kw.get("schema"): + self.default_schema = kw.get("schema") + + if kw.get("auth") == "basic": # if auth=basic, add basic authenticator for Presto + kw["auth"] = prestodb.auth.BasicAuthentication(kw.pop("user"), kw.pop("password")) + + if "cert" in kw: # if a certificate was specified in URI, verify session with cert + cert = kw.pop("cert") + self._conn = prestodb.dbapi.connect(**kw) + self._conn._http_session.verify = cert + else: + self._conn = prestodb.dbapi.connect(**kw) + + def _query(self, sql_code: SqlCode) -> Optional[QueryResult]: + "Uses the standard SQL cursor interface" + c = self._conn.cursor() + + if isinstance(sql_code, ThreadLocalInterpreter): + return sql_code.apply_queries(partial(query_cursor, c)) + elif isinstance(sql_code, str): + sql_code = CompiledCode(sql_code, [], None) # Unknown type. #TODO: Should we guess? + + return query_cursor(c, sql_code) + + def close(self): + super().close() + self._conn.close() + + def select_table_schema(self, path: DbPath) -> str: + schema, table = self._normalize_table_path(path) + + return ( + "SELECT column_name, data_type, 3 as datetime_precision, 3 as numeric_precision, NULL as numeric_scale " + "FROM INFORMATION_SCHEMA.COLUMNS " + f"WHERE table_name = '{table}' AND table_schema = '{schema}'" + ) + + @property + def is_autocommit(self) -> bool: + return False diff --git a/reladiff/sqeleton/databases/redshift.py b/reladiff/sqeleton/databases/redshift.py new file mode 100644 index 00000000..4b389b02 --- /dev/null +++ b/reladiff/sqeleton/databases/redshift.py @@ -0,0 +1,128 @@ +from typing import List, Dict +from ..abcs.database_types import Float, TemporalType, FractionalType, DbPath, TimestampTZ +from ..abcs.mixins import AbstractMixin_MD5 +from .postgresql import ( + PostgreSQL, + MD5_HEXDIGITS, + CHECKSUM_HEXDIGITS, + TIMESTAMP_PRECISION_POS, + PostgresqlDialect, + Mixin_NormalizeValue, +) + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"strtol(substring(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16)::decimal(38)" + + +class Mixin_NormalizeValue(Mixin_NormalizeValue): + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + if coltype.rounds: + timestamp = f"{value}::timestamp(6)" + # Get seconds since epoch. Redshift doesn't support milli- or micro-seconds. + secs = f"timestamp 'epoch' + round(extract(epoch from {timestamp})::decimal(38)" + # Get the milliseconds from timestamp. + ms = f"extract(ms from {timestamp})" + # Get the microseconds from timestamp, without the milliseconds! + us = f"extract(us from {timestamp})" + # epoch = Total time since epoch in microseconds. + epoch = f"{secs}*1000000 + {ms}*1000 + {us}" + timestamp6 = ( + f"to_char({epoch}, -6+{coltype.precision}) * interval '0.000001 seconds', 'YYYY-mm-dd HH24:MI:SS.US')" + ) + else: + timestamp6 = f"to_char({value}::timestamp(6), 'YYYY-mm-dd HH24:MI:SS.US')" + return ( + f"RPAD(LEFT({timestamp6}, {TIMESTAMP_PRECISION_POS+coltype.precision}), {TIMESTAMP_PRECISION_POS+6}, '0')" + ) + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + return self.to_string(f"{value}::decimal(38,{coltype.precision})") + + +class Dialect(PostgresqlDialect): + name = "Redshift" + TYPE_CLASSES = { + **PostgresqlDialect.TYPE_CLASSES, + "double": Float, + "real": Float, + } + SUPPORTS_INDEXES = False + + def concat(self, items: List[str]) -> str: + joined_exprs = " || ".join(items) + return f"({joined_exprs})" + + def is_distinct_from(self, a: str, b: str) -> str: + return f"({a} IS NULL != {b} IS NULL) OR ({a}!={b})" + + def type_repr(self, t) -> str: + if isinstance(t, TimestampTZ): + return f"timestamptz" + return super().type_repr(t) + + +class Redshift(PostgreSQL): + dialect = Dialect() + CONNECT_URI_HELP = "redshift://:@/" + CONNECT_URI_PARAMS = ["database?"] + + def select_table_schema(self, path: DbPath) -> str: + database, schema, table = self._normalize_table_path(path) + + info_schema_path = ["information_schema", "columns"] + if database: + info_schema_path.insert(0, database) + + return ( + f"SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale FROM {'.'.join(info_schema_path)} " + f"WHERE table_name = '{table.lower()}' AND table_schema = '{schema.lower()}'" + ) + + def select_external_table_schema(self, path: DbPath) -> str: + database, schema, table = self._normalize_table_path(path) + + db_clause = "" + if database: + db_clause = f" AND redshift_database_name = '{database.lower()}'" + + return ( + f"""SELECT + columnname AS column_name + , CASE WHEN external_type = 'string' THEN 'varchar' ELSE external_type END AS data_type + , NULL AS datetime_precision + , NULL AS numeric_precision + , NULL AS numeric_scale + FROM svv_external_columns + WHERE tablename = '{table.lower()}' AND schemaname = '{schema.lower()}' + """ + + db_clause + ) + + def query_external_table_schema(self, path: DbPath) -> Dict[str, tuple]: + rows = self.query(self.select_external_table_schema(path), list) + if not rows: + raise RuntimeError(f"{self.name}: Table '{'.'.join(path)}' does not exist, or has no columns") + + d = {r[0]: r for r in rows} + assert len(d) == len(rows) + return d + + def query_table_schema(self, path: DbPath) -> Dict[str, tuple]: + try: + return super().query_table_schema(path) + except RuntimeError: + return self.query_external_table_schema(path) + + def _normalize_table_path(self, path: DbPath) -> DbPath: + if len(path) == 1: + return None, self.default_schema, path[0] + elif len(path) == 2: + return None, path[0], path[1] + elif len(path) == 3: + return path + + raise ValueError( + f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected format: table, schema.table, or database.schema.table" + ) diff --git a/reladiff/sqeleton/databases/snowflake.py b/reladiff/sqeleton/databases/snowflake.py new file mode 100644 index 00000000..c78a26e7 --- /dev/null +++ b/reladiff/sqeleton/databases/snowflake.py @@ -0,0 +1,228 @@ +from typing import Union, List +import logging + +from ..abcs.database_types import ( + Timestamp, + TimestampTZ, + Decimal, + Float, + Text, + FractionalType, + TemporalType, + DbPath, + Boolean, + Date, +) +from ..abcs.mixins import ( + AbstractMixin_MD5, + AbstractMixin_NormalizeValue, + AbstractMixin_Schema, + AbstractMixin_TimeTravel, +) +from ..abcs import Compilable +from sqeleton.queries import table, this, SKIP, code +from .base import ( + BaseDialect, + ConnectError, + Database, + import_helper, + CHECKSUM_MASK, + ThreadLocalInterpreter, + Mixin_RandomSample, +) + + +@import_helper("snowflake") +def import_snowflake(): + import snowflake.connector + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.backends import default_backend + + return snowflake, serialization, default_backend + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"BITAND(md5_number_lower64({s}), {CHECKSUM_MASK})" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + if coltype.rounds: + timestamp = f"to_timestamp(round(date_part(epoch_nanosecond, convert_timezone('UTC', {value})::timestamp(9))/1000000000, {coltype.precision}))" + else: + timestamp = f"cast(convert_timezone('UTC', {value}) as timestamp({coltype.precision}))" + + return f"to_char({timestamp}, 'YYYY-MM-DD HH24:MI:SS.FF6')" + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + return self.to_string(f"cast({value} as decimal(38, {coltype.precision}))") + + def normalize_boolean(self, value: str, _coltype: Boolean) -> str: + return self.to_string(f"{value}::int") + + +class Mixin_Schema(AbstractMixin_Schema): + def table_information(self) -> Compilable: + return table("INFORMATION_SCHEMA", "TABLES") + + def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable: + return ( + self.table_information() + .where( + this.TABLE_SCHEMA == table_schema, + this.TABLE_NAME.like(like) if like is not None else SKIP, + this.TABLE_TYPE == "BASE TABLE", + ) + .select(table_name=this.TABLE_NAME) + ) + + +class Mixin_TimeTravel(AbstractMixin_TimeTravel): + def time_travel( + self, + table: Compilable, + before: bool = False, + timestamp: Compilable = None, + offset: Compilable = None, + statement: Compilable = None, + ) -> Compilable: + at_or_before = "AT" if before else "BEFORE" + if timestamp is not None: + assert offset is None and statement is None + key = "timestamp" + value = timestamp + elif offset is not None: + assert statement is None + key = "offset" + value = offset + else: + assert statement is not None + key = "statement" + value = statement + + return code(f"{{table}} {at_or_before}({key} => {{value}})", table=table, value=value) + + +class Dialect(BaseDialect, Mixin_Schema): + name = "Snowflake" + ROUNDS_ON_PREC_LOSS = False + TYPE_CLASSES = { + # Timestamps + "TIMESTAMP_NTZ": Timestamp, + "TIMESTAMP_LTZ": Timestamp, + "TIMESTAMP_TZ": TimestampTZ, + "DATE": Date, + # Numbers + "NUMBER": Decimal, + "FLOAT": Float, + # Text + "TEXT": Text, + # Boolean + "BOOLEAN": Boolean, + } + MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_TimeTravel, Mixin_RandomSample} + + def explain_as_text(self, query: str) -> str: + return f"EXPLAIN USING TEXT {query}" + + def quote(self, s: str): + return f'"{s}"' + + def to_string(self, s: str): + return f"cast({s} as string)" + + def table_information(self) -> Compilable: + return table("INFORMATION_SCHEMA", "TABLES") + + def set_timezone_to_utc(self) -> str: + return "ALTER SESSION SET TIMEZONE = 'UTC'" + + def optimizer_hints(self, hints: str) -> str: + raise NotImplementedError("Optimizer hints not yet implemented in snowflake") + + def type_repr(self, t) -> str: + if isinstance(t, TimestampTZ): + return f"timestamp_tz({t.precision})" + return super().type_repr(t) + + +class Snowflake(Database): + dialect = Dialect() + CONNECT_URI_HELP = "snowflake://:@//?warehouse=" + CONNECT_URI_PARAMS = ["database", "schema"] + CONNECT_URI_KWPARAMS = ["warehouse"] + + def __init__(self, *, schema: str, **kw): + snowflake, serialization, default_backend = import_snowflake() + logging.getLogger("snowflake.connector").setLevel(logging.WARNING) + + # Ignore the error: snowflake.connector.network.RetryRequest: could not find io module state + # It's a known issue: https://github.com/snowflakedb/snowflake-connector-python/issues/145 + logging.getLogger("snowflake.connector.network").disabled = True + + assert '"' not in schema, "Schema name should not contain quotes!" + # If a private key is used, read it from the specified path and pass it as "private_key" to the connector. + if "key" in kw: + with open(kw.get("key"), "rb") as key: + if "password" in kw: + raise ConnectError("Cannot use password and key at the same time") + if kw.get("private_key_passphrase"): + encoded_passphrase = kw.get("private_key_passphrase").encode() + else: + encoded_passphrase = None + p_key = serialization.load_pem_private_key( + key.read(), + password=encoded_passphrase, + backend=default_backend(), + ) + + kw["private_key"] = p_key.private_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + + self._conn = snowflake.connector.connect(schema=f'"{schema}"', **kw) + self._conn.telemetry_enabled = False + self.default_schema = schema + + def close(self): + super().close() + self._conn.close() + + def _query(self, sql_code: Union[str, ThreadLocalInterpreter]): + "Uses the standard SQL cursor interface" + return self._query_conn(self._conn, sql_code) + + def select_table_schema(self, path: DbPath) -> str: + """Provide SQL for selecting the table schema as (name, type, date_prec, num_prec)""" + database, schema, name = self._normalize_table_path(path) + info_schema_path = ["information_schema", "columns"] + if database: + info_schema_path.insert(0, database) + + return ( + "SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale " + f"FROM {'.'.join(info_schema_path)} " + f"WHERE table_name = '{name}' AND table_schema = '{schema}'" + ) + + def _normalize_table_path(self, path: DbPath) -> DbPath: + if len(path) == 1: + return None, self.default_schema, path[0] + elif len(path) == 2: + return None, path[0], path[1] + elif len(path) == 3: + return path + + raise ValueError( + f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected format: table, schema.table, or database.schema.table" + ) + + @property + def is_autocommit(self) -> bool: + return True + + def query_table_unique_columns(self, path: DbPath) -> List[str]: + return [] diff --git a/reladiff/sqeleton/databases/trino.py b/reladiff/sqeleton/databases/trino.py new file mode 100644 index 00000000..f07b149d --- /dev/null +++ b/reladiff/sqeleton/databases/trino.py @@ -0,0 +1,71 @@ +import uuid + +from ..abcs.database_types import TemporalType, ColType_UUID, String_UUID +from . import presto +from .base import import_helper +from .base import TIMESTAMP_PRECISION_POS + + +@import_helper("trino") +def import_trino(): + import trino + + return trino + + +Mixin_MD5 = presto.Mixin_MD5 + + +class Mixin_NormalizeValue(presto.Mixin_NormalizeValue): + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + if coltype.rounds: + s = f"date_format(cast({value} as timestamp({coltype.precision})), '%Y-%m-%d %H:%i:%S.%f')" + else: + s = f"date_format(cast({value} as timestamp(6)), '%Y-%m-%d %H:%i:%S.%f')" + + return f"RPAD(RPAD({s}, {TIMESTAMP_PRECISION_POS + coltype.precision}, '.'), {TIMESTAMP_PRECISION_POS + 6}, '0')" + + def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str: + if isinstance(coltype, String_UUID): + return f"TRIM({value})" + return f"CAST({value} AS VARCHAR)" + + +class Dialect(presto.Dialect): + name = "Trino" + ARG_SYMBOL = "?" + + def set_timezone_to_utc(self) -> str: + return "SET TIME ZONE '+00:00'" + + def uuid_value(self, u: uuid.UUID) -> str: + return f"CAST('{u}' AS UUID)" + + +class Trino(presto.Presto): + dialect = Dialect() + CONNECT_URI_HELP = "trino://@//" + CONNECT_URI_PARAMS = ["catalog", "schema"] + + def __init__(self, **kw): + + trino = import_trino() + + if kw.get("schema"): + self.default_schema = kw.get("schema") + + if kw.get("password"): + kw["auth"] = trino.auth.BasicAuthentication( + kw.pop("user"), kw.pop("password") + ) + kw["http_scheme"] = "https" + + cert = kw.pop("cert", None) + self._conn = trino.dbapi.connect(**kw) + if cert is not None: + self._conn._http_session.verify = cert + + + @property + def is_autocommit(self) -> bool: + return True \ No newline at end of file diff --git a/reladiff/sqeleton/databases/vertica.py b/reladiff/sqeleton/databases/vertica.py new file mode 100644 index 00000000..3f853eae --- /dev/null +++ b/reladiff/sqeleton/databases/vertica.py @@ -0,0 +1,181 @@ +from typing import List + +from ..utils import match_regexps +from .base import ( + CHECKSUM_HEXDIGITS, + MD5_HEXDIGITS, + TIMESTAMP_PRECISION_POS, + BaseDialect, + ConnectError, + DbPath, + ColType, + ThreadedDatabase, + import_helper, + Mixin_RandomSample, +) +from ..abcs.database_types import ( + Decimal, + Float, + FractionalType, + Integer, + TemporalType, + Text, + Timestamp, + TimestampTZ, + Boolean, + ColType_UUID, +) +from ..abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue, AbstractMixin_Schema +from ..abcs import Compilable +from ..queries import table, this, SKIP + + +@import_helper("vertica") +def import_vertica(): + import vertica_python + + return vertica_python + + +class Mixin_MD5(AbstractMixin_MD5): + def md5_as_int(self, s: str) -> str: + return f"CAST(HEX_TO_INTEGER(SUBSTRING(MD5({s}), {1 + MD5_HEXDIGITS - CHECKSUM_HEXDIGITS})) AS NUMERIC(38, 0))" + + +class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): + def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: + if coltype.rounds: + return f"TO_CHAR({value}::TIMESTAMP({coltype.precision}), 'YYYY-MM-DD HH24:MI:SS.US')" + + timestamp6 = f"TO_CHAR({value}::TIMESTAMP(6), 'YYYY-MM-DD HH24:MI:SS.US')" + return ( + f"RPAD(LEFT({timestamp6}, {TIMESTAMP_PRECISION_POS+coltype.precision}), {TIMESTAMP_PRECISION_POS+6}, '0')" + ) + + def normalize_number(self, value: str, coltype: FractionalType) -> str: + return self.to_string(f"CAST({value} AS NUMERIC(38, {coltype.precision}))") + + def normalize_uuid(self, value: str, _coltype: ColType_UUID) -> str: + # Trim doesn't work on CHAR type + return f"TRIM(CAST({value} AS VARCHAR))" + + def normalize_boolean(self, value: str, _coltype: Boolean) -> str: + return self.to_string(f"cast ({value} as int)") + + +class Mixin_Schema(AbstractMixin_Schema): + def table_information(self) -> Compilable: + return table("v_catalog", "tables") + + def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable: + return ( + self.table_information() + .where( + this.table_schema == table_schema, + this.table_name.like(like) if like is not None else SKIP, + ) + .select(this.table_name) + ) + + +class Dialect(BaseDialect, Mixin_Schema): + name = "Vertica" + ROUNDS_ON_PREC_LOSS = True + + TYPE_CLASSES = { + # Timestamps + "timestamp": Timestamp, + "timestamptz": TimestampTZ, + # Numbers + "numeric": Decimal, + "int": Integer, + "float": Float, + # Text + "char": Text, + "varchar": Text, + # Boolean + "boolean": Boolean, + } + MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample} + + def quote(self, s: str): + return f'"{s}"' + + def concat(self, items: List[str]) -> str: + return " || ".join(items) + + def to_string(self, s: str) -> str: + return f"CAST({s} AS VARCHAR)" + + def is_distinct_from(self, a: str, b: str) -> str: + return f"not ({a} <=> {b})" + + def parse_type( + self, + table_path: DbPath, + col_name: str, + type_repr: str, + datetime_precision: int = None, + numeric_precision: int = None, + numeric_scale: int = None, + ) -> ColType: + timestamp_regexps = { + r"timestamp\(?(\d?)\)?": Timestamp, + r"timestamptz\(?(\d?)\)?": TimestampTZ, + } + for m, t_cls in match_regexps(timestamp_regexps, type_repr): + precision = int(m.group(1)) if m.group(1) else 6 + return t_cls(precision=precision, rounds=self.ROUNDS_ON_PREC_LOSS) + + number_regexps = { + r"numeric\((\d+),(\d+)\)": Decimal, + } + for m, n_cls in match_regexps(number_regexps, type_repr): + _prec, scale = map(int, m.groups()) + return n_cls(scale) + + string_regexps = { + r"varchar\((\d+)\)": Text, + r"char\((\d+)\)": Text, + } + for m, n_cls in match_regexps(string_regexps, type_repr): + return n_cls() + + return super().parse_type(table_path, col_name, type_repr, datetime_precision, numeric_precision) + + def set_timezone_to_utc(self) -> str: + return "SET TIME ZONE TO 'UTC'" + + def current_timestamp(self) -> str: + return "current_timestamp(6)" + + +class Vertica(ThreadedDatabase): + dialect = Dialect() + CONNECT_URI_HELP = "vertica://:@/" + CONNECT_URI_PARAMS = ["database?"] + + default_schema = "public" + + def __init__(self, *, thread_count, **kw): + self._args = kw + self._args["AUTOCOMMIT"] = False + + super().__init__(thread_count=thread_count) + + def create_connection(self): + vertica = import_vertica() + try: + c = vertica.connect(**self._args) + return c + except vertica.errors.ConnectionError as e: + raise ConnectError(*e.args) from e + + def select_table_schema(self, path: DbPath) -> str: + schema, name = self._normalize_table_path(path) + + return ( + "SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale " + "FROM V_CATALOG.COLUMNS " + f"WHERE table_name = '{name}' AND table_schema = '{schema}'" + ) diff --git a/reladiff/sqeleton/queries/__init__.py b/reladiff/sqeleton/queries/__init__.py new file mode 100644 index 00000000..c52133fb --- /dev/null +++ b/reladiff/sqeleton/queries/__init__.py @@ -0,0 +1,25 @@ +from .compiler import CompileError, Compiler +from .base import SKIP, T_SKIP +from .api import ( + this, + join, + outerjoin, + table, + sum_, + avg, + min_, + max_, + cte, + commit, + when, + coalesce, + and_, + if_, + or_, + leftjoin, + rightjoin, + current_timestamp, + code, +) +from .ast_classes import Expr, ExprNode, Select, Count, BinOp, Explain, In, Code, Column, ITable, ForeignKey +from .extras import Checksum, NormalizeAsString, ApplyFuncAndNormalizeAsString diff --git a/reladiff/sqeleton/queries/api.py b/reladiff/sqeleton/queries/api.py new file mode 100644 index 00000000..e2b59028 --- /dev/null +++ b/reladiff/sqeleton/queries/api.py @@ -0,0 +1,213 @@ +from typing import Optional + +from .ast_classes import * +from .base import args_as_tuple +from ..schema import SchemaInput, _Schema + + +this = This() + + +def join(*tables: ITable) -> Join: + """Inner-join a sequence of table expressions" + + When joining, it's recommended to use explicit tables names, instead of `this`, in order to avoid potential name collisions. + + Example: + :: + + person = table('person') + city = table('city') + + name_and_city = ( + join(person, city) + .on(person['city_id'] == city['id']) + .select(person['id'], city['name']) + ) + """ + return Join(tables) + + +def leftjoin(*tables: ITable) -> Join: + """Left-joins a sequence of table expressions. + + See Also: ``join()`` + """ + return Join(tables, "LEFT") + + +def rightjoin(*tables: ITable): + """Right-joins a sequence of table expressions. + + See Also: ``join()`` + """ + return Join(tables, "RIGHT") + + +def outerjoin(*tables: ITable): + """Outer-joins a sequence of table expressions. + + See Also: ``join()`` + """ + return Join(tables, "FULL OUTER") + + +def cte(expr: Expr, *, name: Optional[str] = None, params: Sequence[str] = None): + """Define a CTE""" + return Cte(expr, name, params) + + +def table(*path: str, schema: SchemaInput = None) -> TablePath: + """Defines a table with a path (dotted name), and optionally a schema. + + Parameters: + path: A list of names that make up the path to the table. + schema: a dictionary of {name: type} + """ + if len(path) == 1 and isinstance(path[0], tuple): + (path,) = path + if len(path) == 1 and isinstance(path[0], TablePath): + path = path[0].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: + schema = _Schema.make(schema) + return TablePath(path, schema) + + +def table_from_sqlmodel(cls: type) -> TablePath: + import sqlmodel + + assert issubclass(cls, sqlmodel.SQLModel) + table_name = cls.__name__.lower() + schema = cls.__annotations__ + return table(table_name, schema=schema) + + +def or_(*exprs: Expr): + """Apply OR between a sequence of boolean expressions""" + exprs = args_as_tuple(exprs) + if len(exprs) == 1: + return exprs[0] + return BinBoolOp("OR", exprs) + + +def and_(*exprs: Expr): + """Apply AND between a sequence of boolean expressions""" + exprs = args_as_tuple(exprs) + if len(exprs) == 1: + return exprs[0] + return BinBoolOp("AND", exprs) + + +def sum_(expr: Expr): + """Call SUM(expr)""" + return Func("sum", [expr]) + + +def avg(expr: Expr): + """Call AVG(expr)""" + return Func("avg", [expr]) + + +def min_(expr: Expr): + """Call MIN(expr)""" + return Func("min", [expr]) + + +def max_(expr: Expr): + """Call MAX(expr)""" + return Func("max", [expr]) + + +def exists(expr: Expr): + """Call EXISTS(expr)""" + return Exists(expr) + + +def if_(cond: Expr, then: Expr, else_: Optional[Expr] = None): + """Conditional expression, shortcut to when-then-else. + + Example: + :: + + # SELECT CASE WHEN b THEN c ELSE d END FROM foo + table('foo').select(if_(b, c, d)) + """ + return when(cond).then(then).else_(else_) + + +def when(*when_exprs: Expr): + """Start a when-then expression + + Example: + :: + + # SELECT CASE + # WHEN (type = 'text') THEN text + # WHEN (type = 'number') THEN number + # ELSE 'unknown type' END + # FROM foo + rows = table('foo').select( + when(this.type == 'text').then(this.text) + .when(this.type == 'number').then(this.number) + .else_('unknown type') + ) + """ + return CaseWhen([]).when(*when_exprs) + + +def coalesce(*exprs): + "Returns a call to COALESCE" + exprs = args_as_tuple(exprs) + return Func("COALESCE", exprs) + + +def insert_rows_in_batches(db, tbl: TablePath, rows, *, columns=None, batch_size=1024 * 8): + assert batch_size > 0 + rows = list(rows) + + while rows: + batch, rows = rows[:batch_size], rows[batch_size:] + db.query(tbl.insert_rows(batch, columns=columns)) + + +def current_timestamp(): + """Returns CURRENT_TIMESTAMP() or NOW()""" + return CurrentTimestamp() + + +def code(code: str, **kw: Dict[str, Expr]) -> Code: + """Inline raw SQL code. + + It allows users to use features and syntax that Sqeleton doesn't yet support. + + It's the user's responsibility to make sure the contents of the string given to `code()` are correct and safe for execution. + + Strings given to `code()` are actually templates, and can embed query expressions given as arguments: + + Parameters: + code: template string of SQL code. Templated variables are signified with '{var}'. + kw: optional parameters for SQL template. + + Examples: + :: + + # SELECT b, FROM tmp WHERE + table('tmp').select(this.b, code("")).where(code("")) + + :: + + def tablesample(tbl, size): + return code("SELECT * FROM {tbl} TABLESAMPLE BERNOULLI ({size})", tbl=tbl, size=size) + + nonzero = table('points').where(this.x > 0, this.y > 0) + + # SELECT * FROM points WHERE (x > 0) AND (y > 0) TABLESAMPLE BERNOULLI (10) + sample_expr = tablesample(nonzero) + """ + return Code(code, kw) + + +commit = Commit() diff --git a/reladiff/sqeleton/queries/ast_classes.py b/reladiff/sqeleton/queries/ast_classes.py new file mode 100644 index 00000000..e5bfd47c --- /dev/null +++ b/reladiff/sqeleton/queries/ast_classes.py @@ -0,0 +1,1223 @@ +from dataclasses import field +from datetime import datetime +from typing import Any, Generator, List, Optional, Sequence, Union, Dict, Literal, overload +from collections import ChainMap +from functools import lru_cache + +from runtype import dataclass as _dataclass, cv_type_checking, isa + +from ..utils import ArithString +from ..abcs import Compilable +from ..abcs.database_types import AbstractTable, AbstractDialect +from ..schema import Schema, TableType, Options + +from .base import SKIP, DbPath, args_as_tuple, SqeletonError + + +class Root: + "Nodes inheriting from Root can be used as root statements in SQL (e.g. SELECT yes, RANDOM() no)" + + +class QueryBuilderError(SqeletonError): + pass + + +class QB_TypeError(QueryBuilderError): + pass + + +dataclass = _dataclass(eq=False, order=False) + +ellipsis = type(Ellipsis) + + +def cache(user_function, /): + 'Simple lightweight unbounded cache. Sometimes called "memoize".' + # Taken from https://github.com/python/cpython/blob/3.11/Lib/functools.py + return lru_cache(maxsize=None)(user_function) + + +class CompilableNode(Compilable): + "Base class for query expression nodes" + + type: Any = None + + def _dfs_values(self): + yield self + for k, vs in self.asdict().items(): # __dict__ provided by runtype.dataclass + if k == "source_table": + # Skip data-sources, we're only interested in data-parameters + continue + if not isinstance(vs, (list, tuple)): + vs = [vs] + for v in vs: + if isinstance(v, CompilableNode): + yield from v._dfs_values() + + +class ExprNode(CompilableNode): + "Base class for query expression nodes" + + def cast_to(self, to): + return Cast(self, to) + + +# Query expressions can only interact with objects that are an instance of 'Expr' +Expr = Union[ExprNode, str, bytes, bool, int, float, datetime, ArithString, None, dict, list] + + +@dataclass +class Code(ExprNode, Root): + """Raw SQL code""" + + code: str + args: Dict[str, Expr] = None + + +def _expr_type(e: Expr) -> type: + if isinstance(e, ExprNode): + return e.type + return type(e) + + +@dataclass +class Alias(ExprNode): + """An alias of a column""" + + expr: Expr + name: str + + @property + def type(self): + return _expr_type(self.expr) + + +def _drop_skips(exprs): + return [e for e in exprs if e is not SKIP] + + +def _drop_skips_dict(exprs_dict): + return {k: v for k, v in exprs_dict.items() if v is not SKIP} + + +class ITable(AbstractTable): + """Interface for tabular objects, with focus on SQL operations""" + + source_table: Any + schema: Schema = None + + def select(self, *exprs: Expr, distinct: bool = SKIP, optimizer_hints=SKIP, **named_exprs) -> "Select": + """Create a new table with the specified fields""" + exprs = args_as_tuple(exprs) + exprs = _drop_skips(exprs) + named_exprs = _drop_skips_dict(named_exprs) + exprs += _named_exprs_as_aliases(named_exprs) + resolve_names(self.source_table, exprs) + return Select.make(self, columns=exprs, distinct=distinct, optimizer_hints=optimizer_hints) + + def where(self, *exprs) -> "Select": + """Create a filter for the table, using a WHERE clause. + + Parameters: + exprs: A list of expressions to filter the table with. Uses of ``SKIP`` will be ignored. + + Returns: + A new table expression with the filter applied. + + """ + exprs = args_as_tuple(exprs) + exprs = _drop_skips(exprs) + if not exprs: + return self + + resolve_names(self.source_table, exprs) + return Select.make(self, where_exprs=exprs) + + def order_by(self, *exprs): + """Order the table by the given expressions, using the ORDER BY clause. + + Parameters: + exprs: A list of expressions to order by. Uses of ``SKIP`` will be ignored. + + Returns: + A new table expression with the order applied. + """ + exprs = args_as_tuple(exprs) + exprs = _drop_skips(exprs) + if not exprs: + return self + + resolve_names(self.source_table, exprs) + return Select.make(self, order_by_exprs=exprs) + + def limit(self, limit: int): + """Limit the number of rows returned by the query, using the LIMIT clause. + + Parameters: + limit: The maximum number of rows to return. ``SKIP`` will set no limit. + """ + if limit is SKIP: + return self + + return Select.make(self, limit_expr=limit) + + def join(self, target: "ITable") -> "Join": + """Join this table with the target table. + + Parameters: + target: The table to join with. + + Returns: + A new table expression representing the join. + """ + return Join([self, target]) + + def group_by(self, *keys) -> "GroupBy": + """Group according to the given keys, using the GROUP BY clause. + + Must be followed by a call to :ref:``GroupBy.agg()`` + + Parameters: + keys: A list of expressions to group by. Uses of ``SKIP`` will be ignored. + + Returns: + A new table expression with the group-by applied. + """ + keys = args_as_tuple(keys) + keys = _drop_skips(keys) + resolve_names(self.source_table, keys) + + return GroupBy(self, keys) + + def _get_column(self, name: str): + if self.schema: + name = self.schema.get_key(name) # Get the actual name. Might be case-insensitive. + return Column(self, name) + + # def __getattr__(self, column): + # return self._get_column(column) + + def __getitem__(self, column): + if isinstance(column, (list, tuple)): + return [self[c] for c in column] + elif column is ...: + assert self.schema + return [self[k] for k in self.schema] + if not isinstance(column, str): + raise TypeError(column) + return self._get_column(column) + + def count(self) -> "Select": + """SELECT COUNT(*) FROM self""" + return Select(self, [Count()]) + + def union(self, other: "ITable") -> "TableOp": + """SELECT * FROM self UNION other""" + return TableOp("UNION", self, other) + + def union_all(self, other: "ITable") -> "TableOp": + """SELECT * FROM self UNION ALL other""" + return TableOp("UNION ALL", self, other) + + def minus(self, other: "ITable") -> "TableOp": + """SELECT * FROM self EXCEPT other""" + # aka + return TableOp("EXCEPT", self, other) + + def intersect(self, other: "ITable") -> "TableOp": + """SELECT * FROM self INTERSECT other""" + return TableOp("INTERSECT", self, other) + + def alias(self, name: str) -> "TableAlias": + """Create an alias for the table. If there already is an alias, it will be replaced.""" + if isinstance(self, TableAlias): + return self.replace(name=name) + return TableAlias(self, name) + + +@dataclass +class Concat(ExprNode): + """Concatenate expressions, with an optional seperator.""" + + exprs: list + sep: str = None + + type = str + + +@dataclass +class Count(ExprNode): + """Count the number of rows in a table""" + + expr: Expr = None + distinct: bool = False + + type = int + + +class LazyOps: + def __add__(self, other): + return BinOp("+", [self, other]) + + def __sub__(self, other): + return BinOp("-", [self, other]) + + def __rsub__(self, other): + return BinOp("-", [other, self]) + + def __mul__(self, other): + return BinOp("*", [self, other]) + + __radd__ = __add__ + __rmul__ = __mul__ + + def __truediv__(self, other): + return BinOp("/", [self, other]) + + def __neg__(self): + return UnaryOp("-", self) + + def not_(self): + return UnaryOp("NOT ", self) + + def __gt__(self, other): + return BinBoolOp(">", [self, other]) + + def __ge__(self, other): + return BinBoolOp(">=", [self, other]) + + def __eq__(self, other): + if cv_type_checking.get(): + return super().__eq__(other) + + if other is None: + return BinBoolOp("IS", [self, None]) + + return BinBoolOp("=", [self, other]) + + def __ne__(self, other): + if cv_type_checking.get(): + return super().__ne__(other) + + if other is None: + return BinBoolOp("IS NOT", [self, None]) + + return BinBoolOp("<>", [self, other]) + + def __lt__(self, other): + return BinBoolOp("<", [self, other]) + + def __le__(self, other): + return BinBoolOp("<=", [self, other]) + + def __or__(self, other): + return BinBoolOp("OR", [self, other]) + + def __and__(self, other): + return BinBoolOp("AND", [self, other]) + + def is_distinct_from(self, other: Expr) -> ExprNode: + """Check if the expression is distinct from the other expression (i.e. not equal, treating NULL as a value)""" + return IsDistinctFrom(self, other) + + def like(self, other: Expr) -> ExprNode: + """Check if the expression is like the other expression, using SQL LIKE syntax""" + return BinBoolOp("LIKE", [self, other]) + + def ilike(self, other: Expr) -> ExprNode: + """Check if the expression is like the other expression, case-insensitive, using SQL ILIKE syntax""" + return BinBoolOp("ILIKE", [self, other]) + + def in_(self, *others: Expr) -> ExprNode: + """Check if the expression is in the list or table given + + Parameters: + others: A list of constants, or a single table expression + + Returns: + A new expression representing the 'IN' operation + """ + others = args_as_tuple(others) + assert isinstance(others, tuple), f"Only lists of constants are supported for now, not {others}" + if len(others) == 0: + return False # SQL value + elif len(others) == 1 and isinstance(others[0], ExprTable): + return InTable(self, others[0]) + return In(self, others) + + def test_regex(self, other: Expr) -> ExprNode: + """Check if the expression matches the regex pattern. + + Parameters: + other: A string expression representing the regex pattern. + + Returns: + A new expression representing the regex test. + """ + return TestRegex(self, other) + + def sum(self) -> ExprNode: + """Sum the values of the expression""" + return Func("SUM", [self]) + + def count(self, distinct: bool = False): + """Count the number of rows in the expression + + Parameters: + distinct (bool): If True, only count distinct rows. + """ + return Count(self, distinct=distinct) + + def max(self): + """Get the maximum value of the list expression""" + return Func("MAX", [self]) + + def min(self): + """Get the minimum value of the list expression""" + return Func("MIN", [self]) + + +@dataclass +class TestRegex(ExprNode, LazyOps): + """Check if the expression matches the regex pattern""" + + string: Expr + pattern: Expr + + +@dataclass +class Func(ExprNode, LazyOps): + """Call a function with the given arguments""" + + name: str + args: Sequence[Expr] + ret_type: type = None + + +@dataclass +class WhenThen(ExprNode): + """A 'when/then' clause in a case-when expression""" + + when: Expr + then: Expr + + +@dataclass +class CaseWhen(ExprNode, LazyOps): + """A case-when expression""" + + cases: Sequence[WhenThen] + else_expr: Expr = None + + @property + def type(self): + then_types = {_expr_type(case.then) for case in self.cases} + if self.else_expr: + then_types.add(_expr_type(self.else_expr)) + if len(then_types) > 1: + raise QB_TypeError(f"Non-matching types in when: {then_types}") + (t,) = then_types + return t + + def when(self, *whens: Expr) -> "QB_When": + """Add a new 'when' clause to the case expression + + Must be followed by a call to `.then()` + """ + whens = args_as_tuple(whens) + whens = _drop_skips(whens) + if not whens: + raise QueryBuilderError("Expected valid whens") + + # XXX reimplementing api.and_() + if len(whens) == 1: + return QB_When(self, whens[0]) + return QB_When(self, BinBoolOp("AND", whens)) + + def else_(self, then: Expr) -> "CaseWhen": + """Add an 'else' clause to the case expression. + + Can only be called once per case-when! + + Parameters: + then: The expression to return if none of the 'when' clauses match. + + Returns: + A new case-when expression with the 'else' clause added. + """ + if self.else_expr is not None: + raise QueryBuilderError(f"Else clause already specified in {self}") + + return self.replace(else_expr=then) + + +@dataclass +class QB_When: + "Partial case-when, used for query-building" + casewhen: CaseWhen + when: Expr + + def then(self, then: Expr) -> CaseWhen: + """Add a 'then' clause after a 'when' was added.""" + case = WhenThen(self.when, then) + return self.casewhen.replace(cases=self.casewhen.cases + [case]) + + +@_dataclass(eq=False, order=False) +class IsDistinctFrom(ExprNode, LazyOps): + """Check if two expressions are distinct from each other (i.e. not equal, treating NULL as a value)""" + + a: Expr + b: Expr + type = bool + + +@_dataclass(eq=False, order=False) +class BinOp(ExprNode, LazyOps): + """Binary operation on expressions""" + + op: str + args: Sequence[Expr] + + @property + def type(self): + types = {_expr_type(i) for i in self.args} + if len(types) > 1: + # raise TypeError(f"Expected all args to have the same type, got {types}") + return Union[tuple(types)] + (t,) = types + return t + + +@dataclass +class UnaryOp(ExprNode, LazyOps): + """Unary operation on an expression""" + + op: str + expr: Expr + + @property + def type(self): + return self.expr.type + + +class BinBoolOp(BinOp): + """Binary boolean operation on expressions""" + + type = bool + + +@_dataclass(eq=False, order=False) +class Column(ExprNode, LazyOps): + """A column in a table""" + + source_table: ITable # TODO: TablePath + name: str + + @property + def type(self): + if self.source_table.schema is None: + raise QueryBuilderError(f"Schema required for table {self.source_table}") + return self.source_table.schema[self.name] + + +class ExprTable(ExprNode, ITable): + pass + + +@_dataclass +class TablePath(ExprTable): + """A path to a table in a database""" + + path: DbPath + schema: Optional[Schema] = field(default=None, repr=False) + + @property + def source_table(self): + return self + + @property + def name(self): + return self.path[-1] + + def to_string(self, dialect: AbstractDialect): + return ".".join(map(dialect.quote, self.path)) + + def __repr__(self) -> str: + if self.schema: + return f"TablePath({self.path!r}, schema=<{len(self.schema)} cols>)" + return f"TablePath({self.path!r})" + + # Statement shorthands + def create(self, source_table: ITable = None, *, if_not_exists: bool = False, primary_keys: List[str] = None): + """Returns a query expression to create a new table. + + Parameters: + source_table: a table expression to use for initializing the table. + If not provided, the table must have a schema specified. + if_not_exists: Add a 'if not exists' clause or not. (note: not all dbs support it!) + primary_keys: List of column names which define the primary key + """ + + 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, primary_keys=primary_keys) + + def drop(self, if_exists: bool = False): + """Returns a query expression to delete the table. + + Parameters: + if_not_exists: Add a 'if not exists' clause or not. (note: not all dbs support it!) + """ + return DropTable(self, if_exists=if_exists) + + def truncate(self): + """Returns a query expression to truncate the table. (remove all rows)""" + return TruncateTable(self) + + def delete_rows(self, *where_exprs: Union[Expr, Literal[SKIP]]) -> "DeleteFromTable": + """Returns a query expression to delete rows from the table. + + Parameters: + where_exprs: A list of expressions to filter the rows to delete. Uses of ``SKIP`` will be ignored. + If no where expressions are provided, all rows will be deleted. + + Returns: + A new query expression to delete rows from the table. + """ + where_exprs = args_as_tuple(where_exprs) + where_exprs = _drop_skips(where_exprs) + if not where_exprs: + return self.truncate() + + resolve_names(self.source_table, where_exprs) + return DeleteFromTable(self, where_exprs) + + def update_fields(self, *where_exprs: Expr, **kv) -> "UpdateTable": + """Returns a query expression to update fields in the table. + + Parameters: + where_exprs: A list of expressions to filter the rows to update. Uses of ``SKIP`` will be ignored. + If no where expressions are provided, all rows will be updated. + kv: A dictionary of column names and values to update. The values can be expressions. + + Returns: + A new query expression to update fields in the table. + """ + where_exprs = args_as_tuple(where_exprs) + where_exprs = _drop_skips(where_exprs) + resolve_names(self.source_table, where_exprs) + resolve_names(self.source_table, kv.values()) + return UpdateTable(self, kv, where_exprs) + + def insert_rows(self, rows: Sequence, *, columns: List[str] = None): + """Returns a query expression to insert rows to the table, given as Python values. + + Parameters: + rows: A list of tuples. Must all have the same width. + columns: Names of columns being populated. If specified, must have the same length as the tuples. + """ + # TODO support expressions (now, random, etc.) + rows = list(rows) + if not rows: + return SKIP + + if isinstance(rows[0], dict): + # TODO: Validate all rows are the same? + keys = list(rows[0].keys()) + if not columns: + columns = keys + elif not (set(columns) <= set(rows[0].keys())): + raise ValueError("Keys in dictionary are not a subset of 'columns'") + rows = [[row.get(k) for k in columns] for row in rows] + + return InsertToTable(self, ConstantTable(rows), columns=columns) + + def insert_row(self, *values, columns: List[str] = None, **kw): + """Returns a query expression to insert a single row to the table, given as Python values. + + Parameters: + columns: Names of columns being populated. If specified, must have the same length as 'values' + """ + if (not values) == (not kw): + raise ValueError("Must provide either positional arguments or keyword arguments, but not a mix of both.") + if values: + if len(values) == 1 and isinstance(values[0], TableType): + assert columns is None + kw = { + k: v.default if isinstance(v, Options) else v + for k, v in values[0] + if not (isinstance(v, Options) and v.auto) + } + else: + return InsertToTable(self, ConstantTable([values]), columns=columns) + + assert kw + assert not columns + return InsertToTable(self, ConstantTable([list(kw.values())]), columns=list(kw.keys())) + + def insert_expr(self, expr: Expr): + """Returns a query expression to insert rows to the table, given as a query expression. + + Parameters: + expr: query expression to from which to read the rows + """ + if isinstance(expr, TablePath): + expr = expr.select() + return InsertToTable(self, expr) + + def time_travel( + self, *, before: bool = False, timestamp: datetime = None, offset: int = None, statement: str = None + ) -> Compilable: + """Selects historical data from the table + + Parameters: + before: If false, inclusive of the specified point in time. + If True, only return the time before it. (at/before) + timestamp: A constant timestamp + offset: the time 'offset' seconds before now + statement: identifier for statement, e.g. query ID + + Must specify exactly one of `timestamp`, `offset` or `statement`. + """ + if sum(int(i is not None) for i in (timestamp, offset, statement)) != 1: + raise ValueError("Must specify exactly one of `timestamp`, `offset` or `statement`.") + + if timestamp is not None: + assert offset is None and statement is None + + +@_dataclass +class ForeignKey: + """A foreign key constraint""" + + table: TablePath + field: str + + @property + def type(self): + return self.table.schema[self.field] + + +@dataclass +class TableAlias(ExprTable): + """An alias for a table""" + + source_table: ITable + name: str + + @property + def schema(self): + return self.source_table.schema + + +@dataclass +class Exists(ExprNode, LazyOps): + """Check if a subquery returns any rows""" + + expr: ITable + + type = bool + + +SelectColumns = Sequence[Union[Expr, ellipsis]] + + +def _expand_ellipsis(schema: dict, columns: SelectColumns): + for c in columns: + if c is ...: + # select all, i.e. * + yield from schema.items() + else: + yield c.name, c.type + + +def _union_dicts(*ds): + unioned = {} + for d in ds: + unioned.update(d) + return unioned + + +@dataclass +class Join(ExprNode, ITable, Root): + """A join operation between two tables""" + + source_tables: Sequence[ITable] + op: str = None + on_exprs: Sequence[Expr] = None + columns: SelectColumns = None + + @property + def source_table(self): + return self + + @property + @cache + def schema(self): + if not self.columns: + schemas = [t.schema for t in self.source_tables if t.schema] + assert schemas and all(schemas) + return type(schemas[0])(ChainMap(*schemas)) # TODO merge dictionaries in compliance with SQL dialect! + + # TODO validate types match between both tables + schemas = [s.schema for s in self.source_tables] + d = _union_dicts(*schemas) + return type(schemas[0])(dict(_expand_ellipsis(d, self.columns))) + + def on(self, *exprs) -> "Join": + """Add an ON clause, for filtering the result of the cartesian product (i.e. the JOIN)""" + if len(exprs) == 1: + (e,) = exprs + if isinstance(e, Generator): + exprs = tuple(e) + + exprs = _drop_skips(exprs) + if not exprs: + return self + + exprs = [BinBoolOp("=", [t[e] for t in self.source_tables]) if isinstance(e, str) else e for e in exprs] + + return self.replace(on_exprs=(self.on_exprs or []) + exprs) + + def select(self, *exprs: Expr, **named_exprs) -> "Join": + """Select fields to return from the JOIN operation + + See Also: ``ITable.select()`` + """ + for e in exprs: + if not isa(e, Expr): + raise TypeError(e) + + if self.columns is not None: + # join-select already applied + return ITable.select(self, *exprs, **named_exprs) + + exprs = _drop_skips(exprs) + named_exprs = _drop_skips_dict(named_exprs) + exprs += _named_exprs_as_aliases(named_exprs) + resolve_names(self.source_table, exprs) + # TODO Ensure exprs <= self.columns ? + return self.replace(columns=exprs) + + +@dataclass +class GroupBy(ExprNode, ITable, Root): + """A group-by operation using the GROUP BY clause""" + + table: ITable + keys_: Sequence[Expr] = None # IKey? + values_: Sequence[Expr] = None + having_exprs: Sequence[Expr] = None + + @property + def source_table(self): + return self + + @property + @cache + def schema(self): + s = self.table.schema + if s is None: + return None + return type(s)({c.name: c.type for c in self.keys_ + self.values_}) + + def __post_init__(self): + assert self.keys_ or self.values_ + + def having(self, *exprs) -> "GroupBy": + """Add a 'HAVING' clause to the group-by + + Parameters: + exprs: A list of expressions to filter the groups. Uses of ``SKIP`` will be ignored. + + Returns: + A new group-by expression with the filter applied. + """ + exprs = args_as_tuple(exprs) + exprs = _drop_skips(exprs) + if not exprs: + return self + + resolve_names(self.table, exprs) + return self.replace(having_exprs=(self.having_exprs or []) + exprs) + + def agg(self, *exprs, **named_exprs) -> "GroupBy": + """Select aggregated fields for the group-by. + + Parameters: + exprs: A list of expressions to aggregate. Uses of ``SKIP`` will be ignored. + named_exprs: A dictionary of named expressions to aggregate. Uses of ``SKIP`` will be ignored. + + Returns: + A new group-by expression with the aggregated fields applied. + """ + exprs = args_as_tuple(exprs) + exprs = _drop_skips(exprs) + + named_exprs = _drop_skips_dict(named_exprs) + exprs += _named_exprs_as_aliases(named_exprs) + + resolve_names(self.table, exprs) + return self.replace(values_=(self.values_ or []) + exprs) + + +@dataclass +class TableOp(ExprNode, ITable, Root): + """A binary operation between two tables""" + + op: str + table1: ITable + table2: ITable + + @property + def source_table(self): + return self + + @property + def type(self): + # TODO ensure types of both tables are compatible + return self.table1.type + + @property + def schema(self): + s1 = self.table1.schema + s2 = self.table2.schema + assert len(s1) == len(s2) + return s1 + + +@dataclass +class Desc(ExprNode): + """A descending order""" + + expr: ExprNode + + +@dataclass +class Select(ExprTable, Root): + """A SELECT statement""" + + table: Expr = None + columns: SelectColumns = None + where_exprs: Sequence[Expr] = None + order_by_exprs: Sequence[Expr] = None + group_by_exprs: Sequence[Expr] = None + having_exprs: Sequence[Expr] = None + limit_expr: int = None + distinct: bool = False + optimizer_hints: Sequence[Expr] = None + postfix: str = None + + @property + @cache + def schema(self): + s = self.table.schema + if s is None or not self.columns: + return s + return type(s)(dict(_expand_ellipsis(s, self.columns))) + + @property + def source_table(self): + return self + + @classmethod + def make(cls, table: ITable, distinct: bool = SKIP, optimizer_hints: str = SKIP, **kwargs): + assert "table" not in kwargs + + if optimizer_hints is not SKIP: + kwargs["optimizer_hints"] = optimizer_hints + if distinct is not SKIP: + kwargs["distinct"] = distinct + + # If table is not a select, return a new Select instance + if not isinstance(table, cls): + return cls(table, **kwargs) + + if "columns" in kwargs and table.columns is not None: + return cls(table, **kwargs) + + # We can safely assume isinstance(table, Select) + # We will try to merge them, instead of creating nested instances + if distinct is False and table.distinct: + # Cannot merge the selects + return cls(table, **kwargs) + + if table.limit_expr or table.group_by_exprs: + # Cannot merge the selects + return cls(table, **kwargs) + + # Fill in missing attributes + for k, v in kwargs.items(): + if getattr(table, k) is not None: + if k == "where_exprs": # Additive attribute + kwargs[k] = getattr(table, k) + v + elif k in ["distinct", "optimizer_hints"]: + pass + else: + raise ValueError(k) + + return table.replace(**kwargs) + + +@dataclass +class Cte(ExprNode, ITable): + """A common table expression (i.e. shared subquery)""" + + source_table: Expr + name: str = None + params: Sequence[str] = None + + @property + def schema(self): + # TODO add cte to schema + return self.source_table.schema + + +def _named_exprs_as_aliases(named_exprs): + return [Alias(expr, name) for name, expr in named_exprs.items()] + + +def resolve_names(source_table, exprs): + i = 0 + for expr in exprs: + # Iterate recursively and update _ResolveColumn instances with the right expression + if isinstance(expr, ExprNode): + for v in expr._dfs_values(): + if isinstance(v, _ResolveColumn): + v.resolve(source_table._get_column(v.resolve_name)) + i += 1 + + +@_dataclass(frozen=False, eq=False, order=False) +class _ResolveColumn(ExprNode, LazyOps): + resolve_name: str + resolved: Expr = None + + def resolve(self, expr: Expr): + if self.resolved is not None: + raise QueryBuilderError(f"Column '{self.resolve_name}' Already resolved! To value: {self.resolved}") + self.resolved = expr + + def _get_resolved(self) -> Expr: + if self.resolved is None: + breakpoint() + raise QueryBuilderError(f"Column not resolved: {self.resolve_name}") + return self.resolved + + @property + def type(self): + return self._get_resolved().type + + @property + def name(self): + return self._get_resolved().name + + def __rsub__(self, other): + if other is ...: + # return Wildcard() + pass + + # return super().__rsub__(other) XXX why does it fail? + return LazyOps.__rsub__(self, other) + + +@dataclass +class Wildcard: + """Wildcard for selecting all columns""" + + exclude: List[str] + + +class This: + """Builder object for accessing table attributes. + + Automatically evaluates to the the 'top-most' table during compilation. + """ + + def __getattr__(self, name): + return _ResolveColumn(name) + + @overload + def __getitem__(self, name: str) -> "This": ... + @overload + def __getitem__(self, name: (list, tuple)) -> List["This"]: ... + + def __getitem__(self, name): + if isinstance(name, (list, tuple)): + return [_ResolveColumn(n) for n in name] + return _ResolveColumn(name) + + +@dataclass +class In(ExprNode): + """Check if an expression is in a list of values""" + + expr: Expr + list: Sequence[Expr] + + type = bool + + +@dataclass +class InTable(ExprNode): + """Check if an expression is in a table""" + + expr: Expr + source_table: ExprTable + + type = bool + + +@dataclass +class Cast(ExprNode): + """Cast an expression to a different type""" + + expr: Expr + target_type: Expr + + +@dataclass +class Random(ExprNode, LazyOps): + """A random number""" + + type = float + + +@dataclass +class ConstantTable(ExprNode): + """A table of constant values""" + + rows: Sequence[Sequence] + + +@dataclass +class Explain(ExprNode, Root): + """An EXPLAIN statement""" + + select: Select + + type = str + + +@dataclass +class CurrentTimestamp(ExprNode, LazyOps): + """The current timestamp""" + + type = datetime + + +@dataclass +class TimeTravel(ITable): + """A time-travel operation""" + + table: TablePath + before: bool = False + timestamp: datetime = None + offset: int = None + statement: str = None + + +# DDL + + +class Statement(CompilableNode, Root): + type = None + + +@dataclass +class CreateTable(Statement): + """a CREATE TABLE statement""" + + path: TablePath + source_table: Expr = None + if_not_exists: bool = False + primary_keys: List[str] = None + + +@dataclass +class DropTable(Statement): + """a DROP TABLE statement""" + + path: TablePath + if_exists: bool = False + + +@dataclass +class TruncateTable(Statement): + """a TRUNCATE TABLE statement""" + + path: TablePath + + +class Statement_MaybeReturning(Statement): + returning_exprs: SelectColumns = None + + @property + def type(self): + if self.returning_exprs: + return ITable + return None + + def returning(self, *exprs): + """Add a 'RETURNING' clause to the current node. + + Note: Not all databases support this feature! + """ + if self.returning_exprs: + raise ValueError("A returning clause has already been specified") + + exprs = args_as_tuple(exprs) + exprs = _drop_skips(exprs) + if not exprs: + return self + + resolve_names(self.path, exprs) + return self.replace(returning_exprs=exprs) + + +@dataclass +class DeleteFromTable(Statement_MaybeReturning): + """a DELETE FROM statement""" + + path: TablePath + where_exprs: Sequence[Expr] = None + returning_exprs: SelectColumns = None + + +@dataclass +class UpdateTable(Statement_MaybeReturning): + """an UPDATE statement""" + + path: TablePath + updates: Dict[str, Expr] + where_exprs: Sequence[Expr] = None + returning_exprs: SelectColumns = None + + +@dataclass +class InsertToTable(Statement_MaybeReturning): + """an INSERT INTO statement""" + + path: TablePath + expr: Expr + columns: List[str] = None + returning_exprs: SelectColumns = None + + +@dataclass +class Commit(Statement): + """Generate a COMMIT statement, if we're in the middle of a transaction, or in auto-commit. Otherwise SKIP.""" + + +@dataclass +class Param(ExprNode, ITable): + """A value placeholder, to be specified at compilation time using the `cv_params` context variable.""" + + name: str + + @property + def source_table(self): + return self diff --git a/reladiff/sqeleton/queries/base.py b/reladiff/sqeleton/queries/base.py new file mode 100644 index 00000000..a52133b0 --- /dev/null +++ b/reladiff/sqeleton/queries/base.py @@ -0,0 +1,27 @@ +from typing import Generator + +from ..abcs import DbPath, DbKey +from ..schema import Schema + + +class T_SKIP: + def __repr__(self): + return "SKIP" + + def __deepcopy__(self, memo): + return self + + +SKIP = T_SKIP() + + +class SqeletonError(Exception): + pass + + +def args_as_tuple(exprs): + if len(exprs) == 1: + (e,) = exprs + if isinstance(e, Generator): + return tuple(e) + return exprs diff --git a/reladiff/sqeleton/queries/compiler.py b/reladiff/sqeleton/queries/compiler.py new file mode 100644 index 00000000..cb66c44c --- /dev/null +++ b/reladiff/sqeleton/queries/compiler.py @@ -0,0 +1,535 @@ +import contextvars +import decimal +import json +import random +import re +import typing as t +from dataclasses import is_dataclass +from datetime import date, datetime +from enum import Enum +from typing import Any, Dict, List, Optional, Sequence, Tuple +from uuid import UUID + +from runtype import Dispatch, dataclass + +from ..abcs import AbstractCompiler, AbstractDatabase, AbstractDialect, Compilable, DbPath +from ..abcs.mixins import AbstractMixin_Regex, AbstractMixin_TimeTravel +from ..schema import _Field +from ..utils import ArithString, join_iter +from . import ast_classes as ast +from .base import SKIP + +cv_params = contextvars.ContextVar("params") + + +class CompileError(Exception): + pass + + +md = Dispatch() + + +@dataclass +class CompiledCode: + code: str + args: List[Any] + type: Optional[type] + + +def eval_template(query_template: str, data_dict: Dict[str, Any], arg_symbol) -> Tuple[str, list]: + args = [] + + def replace_match(match): + varname = match.group(1) + args.append(data_dict[varname]) + return arg_symbol + + return re.sub("\xff" + r"\[(\w+)\]", replace_match, query_template), args + + +@dataclass +class Compiler(AbstractCompiler): + database: AbstractDatabase + in_select: bool = False # Compilation runtime flag + in_join: bool = False # Compilation runtime flag + + _table_context: List = [] # List[ITable] + _subqueries: Dict[str, Any] = {} # XXX not thread-safe + _args: Dict[str, Any] = {} + _args_enabled: bool = False + _is_root: bool = True + + _counter: List = [0] + + @property + def dialect(self) -> AbstractDialect: + return self.database.dialect + + def compile(self, elem: Any, params: Optional[Dict[str, Any]] = None) -> str: + if params: + cv_params.set(params) + + if self._is_root and isinstance(elem, Compilable) and not isinstance(elem, ast.Root): + from .ast_classes import Select + + elem = Select(columns=[elem]) + + res = self._compile(elem) + if self._is_root and self._subqueries: + subq = ", ".join(f"\n {k} AS ({v})" for k, v in self._subqueries.items()) + self._subqueries.clear() + return f"WITH {subq}\n{res}" + + return res + + def compile_with_args(self, elem: Any, params: Optional[Dict[str, Any]] = None) -> CompiledCode: + assert self._is_root + + if self.dialect.ARG_SYMBOL is not None: + # Only enable if the database supports args. Otherwise compile normally + self = self.replace(_args_enabled=True) + + res = self.compile(elem, params) + if res is SKIP: + return SKIP + + if self._args: + res, args = eval_template(res, self._args, self.dialect.ARG_SYMBOL) + self._args.clear() + else: + args = [] + + return CompiledCode(res, args, elem.type) + + def _add_as_param(self, elem): + if self._args_enabled: + name = self.new_unique_name() + self._args[name] = elem + return f"\xff[{name}]" + + if isinstance(elem, bytes): + return f"b'{elem.decode()}'" + elif isinstance(elem, bytearray): + return f"'{elem.decode()}'" + elif isinstance(elem, (str, UUID)): + escaped = elem.replace("'", "''") + return f"'{escaped}'" + + raise NotImplementedError() + + def _compile(self, elem) -> str: + if elem is None: + return "NULL" + elif isinstance(elem, UUID): + return self.dialect.uuid_value(elem) + elif isinstance(elem, ArithString): + return f"'{elem}'" + elif isinstance(elem, str): + return self._add_as_param(elem) + elif isinstance(elem, (bytes, bytearray)): + return self._add_as_param(elem.decode()) + elif isinstance(elem, Compilable): + # return elem.compile(self.replace(_is_root=False)) + return self.replace(_is_root=False).compile_node(elem) + elif isinstance(elem, (int, float)): + return str(elem) + elif isinstance(elem, datetime): + return self.dialect.timestamp_value(elem) + elif isinstance(elem, date): + return self._compile(str(elem)) + elif isinstance(elem, decimal.Decimal): + return str(elem) + elif elem is ...: + return "*" + elif is_dataclass(elem): + return self._add_as_param(json.dumps(elem.json())) + elif isinstance(elem, (list, dict)): + return self._add_as_param(json.dumps(elem)) + elif isinstance(elem, Enum): + return self._add_as_param(elem.value) + + assert False, elem + + def new_unique_name(self, prefix="tmp"): + self._counter[0] += 1 + return f"{prefix}{self._counter[0]}" + + 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)}") + + def add_table_context(self, *tables: Sequence, **kw): + new_context = self._table_context + list(filter(None, tables)) + if len({t.name for t in new_context}) < len(new_context): + raise ValueError("Duplicate table alias", {t.name for t in new_context}) + return self.replace(_table_context=new_context, **kw) + + def quote(self, s: str): + return self.dialect.quote(s) + + @md + def compile_node(self, c: ast.Code) -> str: + if not c.args: + return c.code + + args = {k: self.compile(v) for k, v in c.args.items()} + return c.code.format(**args) + + @md + def compile_node(self, c: ast.Alias) -> str: + return f"{self.compile(c.expr)} AS {self.quote(c.name)}" + + @md + def compile_node(self, c: ast.Concat) -> str: + # We coalesce because on some DBs (e.g. MySQL) concat('a', NULL) is NULL + # TODO expression can be simpler? + items = [ + f"coalesce({self.compile(ast.Code(self.dialect.to_string(self.compile(expr))))}, '')" + for expr in c.exprs + ] + assert items + if len(items) == 1: + return items[0] + + if c.sep: + items = list(join_iter(f"'{c.sep}'", items)) + return self.dialect.concat(items) + + @md + def compile_node(self, c: ast.Count) -> str: + expr = self.compile(c.expr) if c.expr else "*" + if c.distinct: + return f"count(distinct {expr})" + + return f"count({expr})" + + @md + def compile_node(self, n: ast.TestRegex) -> str: + if not isinstance(self.dialect, AbstractMixin_Regex): + raise NotImplementedError(f"No regex implementation for database '{self.database}'") + regex = self.dialect.test_regex(n.string, n.pattern) + return self.compile(regex) + + @md + def compile_node(self, c: ast.Func) -> str: + args = ", ".join(self.compile(e) for e in c.args) + return f"{c.name}({args})" + + @md + def compile_node(self, n: ast.WhenThen) -> str: + return f"WHEN {self.compile(n.when)} THEN {self.compile(n.then)}" + + @md + def compile_node(self, c: ast.CaseWhen) -> str: + when_thens = " ".join(self.compile(case) for case in c.cases) + else_expr = (" ELSE " + self.compile(c.else_expr)) if c.else_expr is not None else "" + return f"CASE {when_thens}{else_expr} END" + + @md + def compile_node(self, n: ast.IsDistinctFrom) -> str: + return self.dialect.is_distinct_from(self.compile(n.a), self.compile(n.b)) + + @md + def compile_node(self, n: ast.BinOp) -> str: + expr = f" {n.op} ".join(self.compile(a) for a in n.args) + return f"({expr})" + + @md + def compile_node(self, n: ast.UnaryOp) -> str: + return f"({n.op}{self.compile(n.expr)})" + + @md + def compile_node(self, n: ast.Column) -> str: + if self._table_context: + if len(self._table_context) > 1: + possible_owners = [t for t in self._table_context if t.schema is None or n.name in t.schema] + if len(possible_owners) > 1: + owners = [t for t in possible_owners if t is n.source_table] + if owners: + (owner,) = owners + if isinstance(owner, ast.TablePath): + return f"{self.compile(owner)}.{self.quote(n.name)}" + elif isinstance(owner, ast.TableAlias): + return f"{self.quote(owner.name)}.{self.quote(n.name)}" + + aliases = [ + t + for t in self._table_context + if isinstance(t, ast.TableAlias) and (t.source_table is n.source_table or t is n.source_table) + ] + if not aliases: + return self.quote(n.name) + elif len(aliases) > 1: + names = [a.name for a in aliases] + raise CompileError(f"Too many aliases for column {n.name} between tables: {names}") + (alias,) = aliases + + return f"{self.quote(alias.name)}.{self.quote(n.name)}" + + return self.quote(n.name) + + @md + def compile_node(self, n: ast.TableAlias) -> str: + return f"{self.compile(n.source_table)} {self.quote(n.name)}" + + @md + def compile_node(self, n: ast.Exists) -> str: + # TODO use context to avoid replace + self = self.replace(in_select=False) + return f"EXISTS ({self.compile(n.expr)})" + + @md + def compile_node(self, n: ast.TablePath) -> str: + # TODO normalize path? + return n.to_string(self.dialect) + + @md + def compile_node(self, n: ast.TableOp) -> str: + # TODO contextvar for in_select + c = self.replace(in_select=False) + table_expr = f"{c.compile(n.table1)} {n.op} {c.compile(n.table2)}" + if self.in_select: + table_expr = f"({table_expr})" + return table_expr + + @md + def compile_node(self, n: ast.Desc) -> str: + e = self.compile(n.expr) + return f"{e} DESC" + + @md + def compile_node(self, n: ast.Cte) -> str: + # TODO contextvar for _table_context + c = self.replace(_table_context=[], in_select=False) + compiled = c.compile(n.source_table) + + name = n.name or self.new_unique_name() + name_params = f"{name}({', '.join(n.params)})" if n.params else name + self._subqueries[name_params] = compiled + + return name + + @md + def compile_node(self, n: ast._ResolveColumn) -> str: + return self.compile_node(n._get_resolved()) + + @md + def compile_node(self, n: ast.In): + elems = ", ".join(map(self.compile, n.list)) + return f"({self.compile(n.expr)} IN ({elems}))" + + @md + def compile_node(self, n: ast.InTable): + table = self.replace(in_select=False).compile_node(n.source_table) + return f"({self.compile(n.expr)} IN ({table}))" + + @md + def compile_node(self, n: ast.Cast) -> str: + return f"cast({self.compile(n.expr)} as {self.compile(n.target_type)})" + + @md + def compile_node(self, n: ast.Random) -> str: + return self.dialect.random() + + @md + def compile_node(self, n: ast.Explain): + return self.dialect.explain_as_text(self.compile(n.select)) + + @md + def compile_node(self, n: ast.CurrentTimestamp): + return self.dialect.current_timestamp() + + @md + def compile_node(self, n: ast.TimeTravel): + assert isinstance(self.dialect, AbstractMixin_TimeTravel) + return self.compile( + self.dialect.time_travel( + n.table, before=n.before, timestamp=n.timestamp, offset=n.offset, statement=n.statement + ) + ) + + @md + def compile_node(self, n: ast.Param) -> str: + params = cv_params.get() + return self._compile(params[n.name]) + + @md + def compile_node(self, n: ast.Commit) -> str: + return "COMMIT" if not self.database.is_autocommit else SKIP + + @md + def compile_node(self, n: ast.InsertToTable) -> str: + if isinstance(n.expr, ast.ConstantTable): + compiled_rows = [[self.compile(v) for v in r] for r in n.expr.rows] + expr = self.dialect.immediate_values(compiled_rows) + else: + expr = self.compile(n.expr) + + columns = "(%s)" % ", ".join(map(self.quote, n.columns)) if n.columns is not None else "" + + q = f"INSERT INTO {self.compile(n.path)}{columns} {expr}" + + q += self._compile_returning(n) + return q + + @md + def compile_node(self, n: ast.UpdateTable) -> str: + updates = [f"{k} = {self.compile(v)}" for k, v in n.updates.items()] + update = f"UPDATE {self.compile(n.path)} SET " + ", ".join(updates) + + if n.where_exprs: + update += " WHERE " + " AND ".join(map(self.compile, n.where_exprs)) + update += self._compile_returning(n) + return update + + @md + def compile_node(self, n: ast.DeleteFromTable) -> str: + delete = f"DELETE FROM {self.compile(n.path)}" + if n.where_exprs: + delete += " WHERE " + " AND ".join(map(self.compile, n.where_exprs)) + + delete += self._compile_returning(n) + return delete + + def _compile_returning(self, n: ast.Statement_MaybeReturning): + if not n.returning_exprs: + return "" + + columns = ", ".join(map(self.compile, n.returning_exprs)) + return " RETURNING " + columns + + @md + def compile_node(self, n: ast.TruncateTable) -> str: + return f"TRUNCATE TABLE {self.compile(n.path)}" + + @md + def compile_node(self, n: ast.DropTable) -> str: + ie = "IF EXISTS " if n.if_exists else "" + return f"DROP TABLE {ie}{self.compile(n.path)}" + + @md + def compile_node(self, n: ast.CreateTable) -> str: + ne = "IF NOT EXISTS " if n.if_not_exists else "" + if n.source_table: + return f"CREATE TABLE {ne}{self.compile(n.path)} AS {self.compile(n.source_table)}" + + primary_keys = [k for k, v in n.path.schema.items() if isinstance(v, _Field) and v.options.primary_key] + + if n.primary_keys is not None: + if primary_keys: + assert n.primary_keys == primary_keys + else: + primary_keys = n.primary_keys + + if primary_keys and self.dialect.SUPPORTS_PRIMARY_KEY: + pks = ", PRIMARY KEY (%s)" % ", ".join(primary_keys) + else: + pks = "" + + schema = ", ".join(f"{self.dialect.quote(k)} {self.dialect.type_repr(v)}" for k, v in n.path.schema.items()) + return f"CREATE TABLE {ne}{self.compile(n.path)}({schema}{pks})" + + @md + def compile_node(self, n: ast.Join) -> str: + tables = [ + t if isinstance(t, ast.TableAlias) else ast.TableAlias(t, self.new_unique_name()) for t in n.source_tables + ] + c = self.replace(in_select=True) + op = " JOIN " if n.op is None else f" {n.op} JOIN " + joined = op.join(c.compile(t) for t in tables) + c = self.add_table_context(*tables, in_select=True) + + if n.on_exprs: + on = " AND ".join(c.compile(e) for e in n.on_exprs) + res = f"{joined} ON {on}" + else: + res = joined + + columns = "*" if not n.columns else ", ".join(map(c.compile, n.columns)) + select = f"SELECT {columns} FROM {res}" + + if self.in_select: + select = f"({select})" + return select + + @md + def compile_node(self, n: ast.GroupBy) -> str: + if n.values_ is None: + raise CompileError(".group_by() must be followed by a call to .agg()") + keys = [str(i + 1) for i in range(len(n.keys_))] + columns = (n.keys_ or []) + (n.values_ or []) + if isinstance(n.table, ast.Select) and not n.table.columns and n.table.group_by_exprs is None: + return self.compile( + n.table.replace( + columns=columns, + group_by_exprs=[ast.Code(k) for k in keys], + having_exprs=n.having_exprs, + ) + ) + + keys_str = ", ".join(keys) + columns_str = ", ".join(self.compile(x) for x in columns) + having_str = " HAVING " + " AND ".join(map(self.compile, n.having_exprs)) if n.having_exprs is not None else "" + select = f"SELECT {columns_str} FROM {SelectCompiler(self).compile(n.table)} GROUP BY {keys_str}{having_str}" + + if self.in_select: + select = f"({select})" + return select + + @md + def compile_node(self, n: ast.Select) -> str: + c = SelectCompiler(self) + + if isinstance(n.table, (ast.TablePath,)): + c = c.add_table_context(n.table) + + columns = ", ".join(map(c.compile, n.columns)) if n.columns else "*" + distinct = "DISTINCT " if n.distinct else "" + optimizer_hints = self.dialect.optimizer_hints(n.optimizer_hints) if n.optimizer_hints else "" + select = f"SELECT {optimizer_hints}{distinct}{columns}" + + if n.table: + select += " FROM " + c.compile(n.table) + elif self.dialect.PLACEHOLDER_TABLE: + select += f" FROM {self.dialect.PLACEHOLDER_TABLE}" + + if n.where_exprs: + select += " WHERE " + " AND ".join(map(c.compile, n.where_exprs)) + + if n.group_by_exprs: + select += " GROUP BY " + ", ".join(map(c.compile, n.group_by_exprs)) + + if n.having_exprs: + assert n.group_by_exprs + select += " HAVING " + " AND ".join(map(c.compile, n.having_exprs)) + + if n.order_by_exprs: + select += " ORDER BY " + ", ".join(map(c.compile, n.order_by_exprs)) + + if n.limit_expr is not None: + select += " " + self.dialect.offset_limit(0, n.limit_expr) + + if n.postfix: + select += n.postfix + + if self.in_select: + select = f"({select})" + return select + + +@dataclass +class SelectCompiler(AbstractCompiler): + c: Compiler + + def compile(self, elem: t.Any, params: t.Dict[str, t.Any] = None) -> str: + if isinstance(elem, (ast.Select, ast.TableOp, ast.GroupBy, ast.Join)): + elem = ast.TableAlias(elem, self.c.new_unique_name()) + c = self.c.replace(in_select=True) + return c.compile(elem, params) + + @property + def dialect(self): + return self.c.dialect + + def add_table_context(self, *tables: t.Sequence, **kw): + return SelectCompiler(self.c.add_table_context(*tables, **kw)) diff --git a/reladiff/sqeleton/queries/extras.py b/reladiff/sqeleton/queries/extras.py new file mode 100644 index 00000000..382407c4 --- /dev/null +++ b/reladiff/sqeleton/queries/extras.py @@ -0,0 +1,71 @@ +"Useful AST classes that don't quite fall within the scope of regular SQL" + +from typing import Callable, Sequence +from runtype import dataclass + +from ..abcs.database_types import ColType, Native_UUID + +from .compiler import Compiler, md +from .ast_classes import Expr, ExprNode, Concat, Code + + + + +@dataclass +class NormalizeAsString(ExprNode): + expr: ExprNode + expr_type: ColType = None + type = str + + +@dataclass +class ApplyFuncAndNormalizeAsString(ExprNode): + expr: ExprNode + apply_func: Callable = None + + +@dataclass +class Checksum(ExprNode): + exprs: Sequence[Expr] + + +class Compiler(Compiler): + @md + def compile_node(c: Compiler, n: NormalizeAsString) -> str: + expr = c.compile(n.expr) + return c.dialect.normalize_value_by_type(expr, n.expr_type or n.expr.type) + + + @md + def compile_node(c: Compiler, n: ApplyFuncAndNormalizeAsString) -> str: + expr = n.expr + expr_type = expr.type + + if isinstance(expr_type, Native_UUID): + # Normalize first, apply template after (for uuids) + # Needed because min/max(uuid) fails in postgresql + expr = NormalizeAsString(expr, expr_type) + if n.apply_func is not None: + expr = n.apply_func(expr) # Apply template using Python's string formatting + + else: + # Apply template before normalizing (for ints) + if n.apply_func is not None: + expr = n.apply_func(expr) # Apply template using Python's string formatting + expr = NormalizeAsString(expr, expr_type) + + return c.compile(expr) + + + @md + def compile_node(c: Compiler, n: Checksum) -> str: + if len(n.exprs) > 1: + exprs = [Code(f"coalesce({c.compile(expr)}, '')") for expr in n.exprs] + # exprs = [c.compile(e) for e in exprs] + expr = Concat(exprs, "|") + else: + # No need to coalesce - safe to assume that key cannot be null + (expr,) = n.exprs + expr = c.compile(expr) + md5 = c.dialect.md5_as_int(expr) + return f"sum({md5})" diff --git a/reladiff/sqeleton/query_utils.py b/reladiff/sqeleton/query_utils.py new file mode 100644 index 00000000..4eb07445 --- /dev/null +++ b/reladiff/sqeleton/query_utils.py @@ -0,0 +1,54 @@ +"Module for query utilities that didn't make it into the query-builder (yet)" + +from contextlib import suppress + +from sqeleton.databases import DbPath, QueryError, Oracle +from sqeleton.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)) diff --git a/reladiff/sqeleton/repl.py b/reladiff/sqeleton/repl.py new file mode 100644 index 00000000..ad0812bd --- /dev/null +++ b/reladiff/sqeleton/repl.py @@ -0,0 +1,284 @@ +import rich.table +import logging + +from pathlib import Path +from time import time + +### XXX Fix for Python 3.8 bug (https://github.com/prompt-toolkit/python-prompt-toolkit/issues/1023) +import asyncio +import selectors + +selector = selectors.SelectSelector() +loop = asyncio.SelectorEventLoop(selector) +asyncio.set_event_loop(loop) +### XXX End of fix + +from pygments.lexers.sql import SqlLexer +from pygments.styles import get_style_by_name +from prompt_toolkit import PromptSession +from prompt_toolkit.lexers import PygmentsLexer +from prompt_toolkit.filters import Condition +from prompt_toolkit.application.current import get_app +from prompt_toolkit.history import FileHistory +from prompt_toolkit.auto_suggest import AutoSuggestFromHistory +from prompt_toolkit.output.color_depth import ColorDepth +from prompt_toolkit.completion import WordCompleter +from prompt_toolkit.styles.pygments import style_from_pygments_cls + +from . import __version__ + +STYLE = style_from_pygments_cls(get_style_by_name("dracula")) + + +sql_keywords = [ + "abort", + "action", + "add", + "after", + "all", + "alter", + "analyze", + "and", + "as", + "asc", + "attach", + "autoincrement", + "before", + "begin", + "between", + "by", + "cascade", + "case", + "cast", + "check", + "collate", + "column", + "commit", + "conflict", + "constraint", + "create", + "cross", + "current_date", + "current_time", + "current_timestamp", + "database", + "default", + "deferrable", + "deferred", + "delete", + "desc", + "detach", + "distinct", + "drop", + "each", + "else", + "end", + "escape", + "except", + "exclusive", + "exists", + "explain", + "fail", + "for", + "foreign", + "from", + "full", + "glob", + "group", + "having", + "if", + "ignore", + "immediate", + "in", + "index", + "indexed", + "initially", + "inner", + "insert", + "instead", + "intersect", + "into", + "is", + "isnull", + "join", + "key", + "left", + "like", + "limit", + "match", + "natural", + "no", + "not", + "notnull", + "null", + "of", + "offset", + "on", + "or", + "order", + "outer", + "plan", + "pragma", + "primary", + "query", + "raise", + "recursive", + "references", + "regexp", + "reindex", + "release", + "rename", + "replace", + "restrict", + "right", + "rollback", + "row", + "savepoint", + "select", + "set", + "table", + "temp", + "temporary", + "then", + "to", + "transaction", + "trigger", + "union", + "unique", + "update", + "using", + "vacuum", + "values", + "view", + "virtual", + "when", + "where", + "with", + "without", +] +sql_completer = WordCompleter(sql_keywords, ignore_case=True) + + +def add_keywords(new_keywords): + global sql_keywords + new = set(new_keywords) - set(sql_keywords) + sql_keywords += new + + +def _code_is_valid(code: str): + if code: + if code[0].isalnum(): + return code[-1] in ";\n" + return True + + +def repl(uri, prompt=" >> "): + rich.print(f"[purple]Sqeleton {__version__} interactive prompt. Enter '?' for help[/purple]") + + db = connect(uri) + db_name = db.name + + try: + session = PromptSession( + lexer=PygmentsLexer(SqlLexer), + completer=sql_completer, + # key_bindings=kb + history=FileHistory(str(Path.home() / ".sqeleton_repl_history")), + auto_suggest=AutoSuggestFromHistory(), + color_depth=ColorDepth.TRUE_COLOR, + style=STYLE, + ) + + @Condition + def multiline_filter(): + text = get_app().layout.get_buffer_by_name("DEFAULT_BUFFER").text + return not _code_is_valid(text) + + prompt = f"{db_name}> " + while True: + # Read + try: + q = session.prompt(prompt, multiline=multiline_filter) + if not q.strip(): + continue + + start_time = time() + run_command(db, q) + + duration = time() - start_time + if duration > 1: + rich.print("(Query took %.2f seconds)" % duration) + + except KeyboardInterrupt: + rich.print("Interrupted (Ctrl+C)") + + except (KeyboardInterrupt, EOFError): + rich.print("Exiting Sqeleton interaction") + + +from . import connect + +import sys + + +def print_table(rows, schema, table_name=""): + # Print rows in a rich table + t = rich.table.Table(title=table_name, caption=f"{len(rows)} rows") + for col in schema: + t.add_column(col) + for r in rows: + t.add_row(*map(str, r)) + rich.print(t) + + +def help(): + rich.print("Commands:") + rich.print(" ?mytable - shows schema of table 'mytable'") + rich.print(" * - shows list of all tables") + rich.print(" *pattern - shows list of all tables with name like pattern") + rich.print("Otherwise, runs regular SQL query") + + +def run_command(db, q): + if q.startswith("*"): + pattern = q[1:] + try: + names = db.query(db.dialect.list_tables(db.default_schema, like=f"%{pattern}%" if pattern else None)) + except Exception as e: + logging.exception(e) + else: + print_table(names, ["name"], "List of tables") + add_keywords([".".join(n) for n in names]) + elif q.startswith("?"): + table_name = q[1:] + if not table_name: + help() + return + try: + path = db.parse_table_name(table_name) + print("->", path) + schema = db.query_table_schema(path) + except Exception as e: + logging.error(e) + else: + print_table([(k, v[1]) for k, v in schema.items()], ["name", "type"], f"Table '{table_name}'") + add_keywords(schema.keys()) + else: + # Normal SQL query + try: + res = db.query(q) + except Exception as e: + logging.error(e) + else: + if res: + print_table(res.rows, res.columns, None) + add_keywords(res.columns) + + +def main(): + uri = sys.argv[1] + return repl(uri) + + +if __name__ == "__main__": + main() diff --git a/reladiff/sqeleton/schema.py b/reladiff/sqeleton/schema.py new file mode 100644 index 00000000..a557ba85 --- /dev/null +++ b/reladiff/sqeleton/schema.py @@ -0,0 +1,73 @@ +import logging +from typing import Any, Union, Type + +from runtype import dataclass + +from .utils import CaseAwareMapping, CaseInsensitiveDict, CaseSensitiveDict +from .abcs import AbstractDatabase, DbPath + +logger = logging.getLogger("schema") + +Schema = CaseAwareMapping + +class TableType: + pass + # TODO: This should replace the current Schema type + + @classmethod + def is_superclass(cls, t): + return isinstance(t, type) and issubclass(t, cls) + + +SchemaInput = Union[Type[TableType], Schema, dict] + +@dataclass +class Options: + default: Any = None + primary_key: bool = False + auto: bool = False + # TODO: foreign_key, unique + # TODO: index? + +@dataclass +class _Field: + type: type + options: Options + +class _Schema(CaseAwareMapping[Union[type, _Field]]): + pass + + @classmethod + def make(cls, schema: SchemaInput): + assert schema + if TableType.is_superclass(schema): + def _make_field(k: str, v: type): + field = getattr(schema, k) + if field: + if not isinstance(field, Options): + field = Options(default=v) + return _Field(v, field) + return v + + schema = CaseSensitiveDict({k:_make_field(k, v) for k,v in schema.__annotations__.items()}) + + elif isinstance(schema, CaseAwareMapping): + pass + else: + assert isinstance(schema, dict), schema + schema = CaseSensitiveDict(schema) + + return schema + +def options(**kw) -> Any: # Any, so that type-checking doesn't complain + return Options(**kw) + + +def create_schema(db: AbstractDatabase, table_path: DbPath, schema: dict, case_sensitive: bool) -> CaseAwareMapping: + if case_sensitive: + return CaseSensitiveDict(schema) + + if len({k.lower() for k in schema}) < len(schema): + logger.warning(f'Ambiguous schema for {db}:{".".join(table_path)} | Columns = {", ".join(list(schema))}') + # logger.warning("We recommend to disable case-insensitivity (set --case-sensitive).") + return CaseInsensitiveDict(schema) diff --git a/reladiff/sqeleton/utils.py b/reladiff/sqeleton/utils.py new file mode 100644 index 00000000..1ceee4a2 --- /dev/null +++ b/reladiff/sqeleton/utils.py @@ -0,0 +1,345 @@ +from typing import ( + Iterable, + Iterator, + MutableMapping, + Union, + Any, + Sequence, + Dict, + Hashable, + TypeVar, + Generator, + List, +) +from abc import ABC, abstractmethod +from weakref import ref +import math +import string +import re +from uuid import UUID +from urllib.parse import urlparse + +# -- Common -- + +try: + from typing import Self +except ImportError: + Self = Any + + +class WeakCache: + def __init__(self): + self._cache = {} + + def _hashable_key(self, k: Union[dict, Hashable]) -> Hashable: + if isinstance(k, dict): + return tuple(k.items()) + return k + + def add(self, key: Union[dict, Hashable], value: Any): + key = self._hashable_key(key) + self._cache[key] = ref(value) + + def get(self, key: Union[dict, Hashable]) -> Any: + key = self._hashable_key(key) + + value = self._cache[key]() + if value is None: + del self._cache[key] + raise KeyError(f"Key {key} not found, or no longer a valid reference") + + return value + + +def join_iter(joiner: Any, iterable: Iterable) -> Iterable: + it = iter(iterable) + try: + yield next(it) + except StopIteration: + return + for i in it: + yield joiner + yield i + + +def safezip(*args): + "zip but makes sure all sequences are the same length" + lens = list(map(len, args)) + if len(set(lens)) != 1: + raise ValueError(f"Mismatching lengths in arguments to safezip: {lens}") + return zip(*args) + + +def is_uuid(u): + try: + UUID(u) + except ValueError: + return False + return True + + +def match_regexps(regexps: Dict[str, Any], s: str) -> Generator[tuple, None, None]: + for regexp, v in regexps.items(): + m = re.match(regexp + "$", s) + if m: + yield m, v + + +# -- Schema -- + +V = TypeVar("V") + + +class CaseAwareMapping(MutableMapping[str, V]): + @abstractmethod + def get_key(self, key: str) -> str: + ... + + @abstractmethod + def __init__(self, initial): + ... + + def new(self, initial=()): + return type(self)(initial) + + +class CaseInsensitiveDict(CaseAwareMapping[V]): + def __init__(self, initial): + self._dict = {k.lower(): (k, v) for k, v in dict(initial).items()} + + def __getitem__(self, key: str) -> V: + return self._dict[key.lower()][1] + + def __iter__(self) -> Iterator[str]: + return iter(self._dict) + + def __len__(self) -> int: + return len(self._dict) + + def __setitem__(self, key: str, value): + k = key.lower() + if k in self._dict: + key = self._dict[k][0] + self._dict[k] = key, value + + def __delitem__(self, key: str): + del self._dict[key.lower()] + + def get_key(self, key: str) -> str: + return self._dict[key.lower()][0] + + def __repr__(self) -> str: + return repr(dict(self.items())) + + +class CaseSensitiveDict(dict, CaseAwareMapping): + def get_key(self, key) -> str: + self[key] # Throw KeyError if key doesn't exist + return key + + def as_insensitive(self): + return CaseInsensitiveDict(self) + + +# -- Alphanumerics -- + +alphanums = " -" + string.digits + string.ascii_uppercase + "_" + string.ascii_lowercase + + +class ArithString(ABC): + @classmethod + def new(cls, *args, **kw): + return cls(*args, **kw) + + def range(self, other: "ArithString", count: int): + assert isinstance(other, ArithString) + checkpoints = split_space(self.int, other.int, count) + return [self.new(int=i) for i in checkpoints] + + # @property + # @abstractmethod + # def int(self): + # ... + + +class ArithUUID(UUID, ArithString): + "A UUID that supports basic arithmetic (add, sub)" + + def __int__(self): + return self.int + + def __add__(self, other: int): + if isinstance(other, int): + return self.new(int=self.int + other) + return NotImplemented + + def __sub__(self, other: Union[UUID, int]): + if isinstance(other, int): + return self.new(int=self.int - other) + elif isinstance(other, UUID): + return self.int - other.int + return NotImplemented + + +def numberToAlphanum(num: int, base: str = alphanums) -> str: + digits = [] + while num > 0: + num, remainder = divmod(num, len(base)) + digits.append(remainder) + return "".join(base[i] for i in digits[::-1]) + + +def alphanumToNumber(alphanum: str, base: str = alphanums) -> int: + num = 0 + for c in alphanum: + num = num * len(base) + base.index(c) + return num + + +def justify_alphanums(s1: str, s2: str): + max_len = max(len(s1), len(s2)) + s1 = s1.ljust(max_len) + s2 = s2.ljust(max_len) + return s1, s2 + + +def alphanums_to_numbers(s1: str, s2: str): + s1, s2 = justify_alphanums(s1, s2) + n1 = alphanumToNumber(s1) + n2 = alphanumToNumber(s2) + return n1, n2 + + +class ArithAlphanumeric(ArithString): + def __init__(self, s: str, max_len=None): + if s is None: + raise ValueError("Alphanum string cannot be None") + if max_len and len(s) > max_len: + raise ValueError(f"Length of alphanum value '{str}' is longer than the expected {max_len}") + + for ch in s: + if ch not in alphanums: + raise ValueError(f"Unexpected character {ch} in alphanum string") + + self._str = s + self._max_len = max_len + + def __str__(self): + s = self._str + if self._max_len: + s = s.rjust(self._max_len, alphanums[0]) + return s + + def __len__(self): + return len(self._str) + + def __repr__(self): + return f'alphanum"{self._str}"' + + def __add__(self, other: "Union[ArithAlphanumeric, int]") -> "ArithAlphanumeric": + if isinstance(other, int): + if other != 1: + raise NotImplementedError("not implemented for arbitrary numbers") + num = alphanumToNumber(self._str) + return self.new(numberToAlphanum(num + 1)) + + return NotImplemented + + def range(self, other: "ArithAlphanumeric", count: int): + assert isinstance(other, ArithAlphanumeric) + n1, n2 = alphanums_to_numbers(self._str, other._str) + split = split_space(n1, n2, count) + return [self.new(numberToAlphanum(s)) for s in split] + + def __sub__(self, other: "Union[ArithAlphanumeric, int]") -> float: + if isinstance(other, ArithAlphanumeric): + n1, n2 = alphanums_to_numbers(self._str, other._str) + return n1 - n2 + + return NotImplemented + + def __ge__(self, other): + if not isinstance(other, type(self)): + return NotImplemented + return self._str >= other._str + + def __lt__(self, other): + if not isinstance(other, type(self)): + return NotImplemented + return self._str < other._str + + def __eq__(self, other): + if not isinstance(other, type(self)): + return NotImplemented + return self._str == other._str + + def new(self, *args, **kw): + return type(self)(*args, **kw, max_len=self._max_len) + + +def number_to_human(n): + millnames = ["", "k", "m", "b"] + n = float(n) + millidx = max( + 0, + min(len(millnames) - 1, int(math.floor(0 if n == 0 else math.log10(abs(n)) / 3))), + ) + + return "{:.0f}{}".format(n / 10 ** (3 * millidx), millnames[millidx]) + + +def split_space(start, end, count) -> List[int]: + size = end - start + assert count <= size, (count, size) + return list(range(start, end, (size + 1) // (count + 1)))[1 : count + 1] + + +def remove_passwords_in_dict(d: dict, replace_with: str = "***"): + for k, v in d.items(): + if k == "password": + d[k] = replace_with + elif isinstance(v, dict): + remove_passwords_in_dict(v, replace_with) + elif k.startswith("database"): + d[k] = remove_password_from_url(v, replace_with) + + +def _join_if_any(sym, args): + args = list(args) + if not args: + return "" + return sym.join(str(a) for a in args if a) + + +def remove_password_from_url(url: str, replace_with: str = "***") -> str: + parsed = urlparse(url) + account = parsed.username or "" + if parsed.password: + account += ":" + replace_with + host = _join_if_any(":", filter(None, [parsed.hostname, parsed.port])) + netloc = _join_if_any("@", filter(None, [account, host])) + replaced = parsed._replace(netloc=netloc) + return replaced.geturl() + + +def match_like(pattern: str, strs: Sequence[str]) -> Iterable[str]: + reo = re.compile(pattern.replace("%", ".*").replace("?", ".") + "$") + for s in strs: + if reo.match(s): + yield s + + +class UnknownMeta(type): + def __instancecheck__(self, instance): + return instance is Unknown + + def __repr__(self): + return "Unknown" + + +class Unknown(metaclass=UnknownMeta): + def __nonzero__(self): + raise TypeError() + + def __new__(cls, *args, **kwargs): + raise RuntimeError("Unknown is a singleton") diff --git a/sqeleton-project/pyproject.toml b/sqeleton-project/pyproject.toml new file mode 100644 index 00000000..544c081b --- /dev/null +++ b/sqeleton-project/pyproject.toml @@ -0,0 +1,114 @@ +[tool.poetry] +name = "sqeleton" +version = "0.1.8" +description = "Python library for querying SQL databases" +authors = ["Erez Shinan "] +license = "MIT" +readme = "README.md" +repository = "https://github.com/erezsh/sqeleton" +documentation = "https://sqeleton.readthedocs.io/en/latest/" +classifiers = [ + "Intended Audience :: Developers", + "Intended Audience :: Information Technology", + "Intended Audience :: System Administrators", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Development Status :: 3 - Alpha", + "Environment :: Console", + "Topic :: Database :: Database Engines/Servers", + "Typing :: Typed" +] +packages = [{ include = "sqeleton" }] + +[tool.poetry.dependencies] +python = "^3.8" +runtype = ">=0.5.0" +dsnparse = "*" +click = ">=8.1" +rich = "*" +toml = ">=0.10.2" +mysql-connector-python = {version=">=8.0.29", optional=true} +# psycopg2 = {version="*", optional=true} +psycopg2-binary = {version="*", optional=true} +snowflake-connector-python = {version=">=2.7.2", optional=true} +cryptography = {version="*", optional=true} +trino = {version=">=0.314.0", optional=true} +presto-python-client = {version="*", optional=true} +clickhouse-driver = {version="*", optional=true} +duckdb = {version=">=0.7.0", optional=true} +textual = {version=">=0.9.1", optional=true} +textual-select = {version="*", optional=true} +pygments = {version=">=2.13.0", optional=true} +prompt-toolkit = {version=">=3.0.36", optional=true} +pyarrow = [ + {version=">=18.0.0", markers="python_version >= '3.9'", optional=true}, + {version="<18.0.0", markers="python_version < '3.9'", optional=true}, +] +pandas = [ + {version=">=2.1", markers="python_version >= '3.9'", optional=true}, + {version="*", markers="python_version < '3.9'", optional=true}, +] +numpy = [ + {version=">=1.26.4", markers="python_version >= '3.9'", optional=true}, + {version=">=1.24.4", markers="python_version < '3.9'", optional=true}, +] +# sqlalchemy-dremio 3.0.4 is missing a type conversion for datetime[ms] and is not compatible with Python 3.8. Created a +# fork with a fix and opening a PR. +sqlalchemy-dremio = {git = "https://github.com/scottkwalker/sqlalchemy_dremio", rev="308e686b166787941180850ce5813c313d3a9985", optional=true} + +[tool.poetry.dev-dependencies] +parameterized = "*" +unittest-parallel = "*" + +duckdb = "<1.2" +mysql-connector-python = "*" +# psycopg2 = "*" +psycopg2-binary = "*" +snowflake-connector-python = ">=2.7.2" +cryptography = "*" +trino = ">=0.314.0" +presto-python-client = "*" +clickhouse-driver = "*" +vertica-python = "*" +pyarrow = [ + {version=">=18.0.0", markers="python_version >= '3.9'"}, + {version="<18.0.0", markers="python_version < '3.9'"}, +] +pandas = [ + {version=">=2.1", markers="python_version >= '3.9'"}, + {version="*", markers="python_version < '3.9'"}, +] +numpy = [ + {version=">=1.26.4", markers="python_version >= '3.9'"}, + {version=">=1.24.4", markers="python_version < '3.9'"}, +] +# sqlalchemy-dremio 3.0.4 is missing a type conversion for datetime[ms] and is not compatible with Python 3.8. Created a +# fork with a fix and opening a PR. +sqlalchemy-dremio = {git = "https://github.com/scottkwalker/sqlalchemy_dremio", rev="308e686b166787941180850ce5813c313d3a9985"} + + +[tool.poetry.extras] +mysql = ["mysql-connector-python"] +postgresql = ["psycopg2-binary"] +snowflake = ["snowflake-connector-python", "cryptography"] +presto = ["presto-python-client"] +oracle = ["cx_Oracle"] +databricks = ["databricks-sql-connector"] +trino = ["trino"] +clickhouse = ["clickhouse-driver"] +vertica = ["vertica-python"] +duckdb = ["duckdb"] +tui = ["textual", "textual-select", "pygments", "prompt-toolkit"] +dremio = ["sqlalchemy-dremio", "pandas", "pyarrow"] + +[build-system] +requires = ["poetry-core>=1.0.0"] +build-backend = "poetry.core.masonry.api" + +[tool.poetry.scripts] +sqeleton = 'sqeleton.__main__:main' + +[tool.ruff] +line-length = 120 +target-version = "py38" From a6cdb664d158293788f927539905fc93f71aebc0 Mon Sep 17 00:00:00 2001 From: vmatt Date: Fri, 13 Feb 2026 12:06:36 +0100 Subject: [PATCH 02/11] update pyproject --- pyproject.toml | 41 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 40fc72f0..0aeb4e78 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,26 @@ cryptography = {version="*", optional=true} trino = {version=">=0.314.0", optional=true} presto-python-client = {version="*", optional=true} clickhouse-driver = {version="*", optional=true} -duckdb = {version=">=0.6.0", optional=true} +duckdb = {version=">=0.7.0", optional=true} +textual = {version=">=0.9.1", optional=true} +textual-select = {version="*", optional=true} +pygments = {version=">=2.13.0", optional=true} +prompt-toolkit = {version=">=3.0.36", optional=true} +pyarrow = [ + {version=">=18.0.0", markers="python_version >= '3.9'", optional=true}, + {version="<18.0.0", markers="python_version < '3.9'", optional=true}, +] +pandas = [ + {version=">=2.1", markers="python_version >= '3.9'", optional=true}, + {version="*", markers="python_version < '3.9'", optional=true}, +] +numpy = [ + {version=">=1.26.4", markers="python_version >= '3.9'", optional=true}, + {version=">=1.24.4", markers="python_version < '3.9'", optional=true}, +] +# sqlalchemy-dremio 3.0.4 is missing a type conversion for datetime[ms] and is not compatible with Python 3.8. Created a +# fork with a fix and opening a PR. +sqlalchemy-dremio = {git = "https://github.com/scottkwalker/sqlalchemy_dremio", rev="308e686b166787941180850ce5813c313d3a9985", optional=true} [tool.poetry.dev-dependencies] parameterized = "*" @@ -51,9 +70,24 @@ trino = ">=0.314.0" presto-python-client = "*" clickhouse-driver = "*" vertica-python = "*" -duckdb = ">=0.6.0" +duckdb = "<1.2" # google-cloud-bigquery = "*" # databricks-sql-connector = "*" +pyarrow = [ + {version=">=18.0.0", markers="python_version >= '3.9'"}, + {version="<18.0.0", markers="python_version < '3.9'"}, +] +pandas = [ + {version=">=2.1", markers="python_version >= '3.9'"}, + {version="*", markers="python_version < '3.9'"}, +] +numpy = [ + {version=">=1.26.4", markers="python_version >= '3.9'"}, + {version=">=1.24.4", markers="python_version < '3.9'"}, +] +# sqlalchemy-dremio 3.0.4 is missing a type conversion for datetime[ms] and is not compatible with Python 3.8. Created a +# fork with a fix and opening a PR. +sqlalchemy-dremio = {git = "https://github.com/scottkwalker/sqlalchemy_dremio", rev="308e686b166787941180850ce5813c313d3a9985"} [tool.poetry.extras] # When adding, update also: README + dev deps just above @@ -68,6 +102,8 @@ trino = ["trino"] clickhouse = ["clickhouse-driver"] vertica = ["vertica-python"] duckdb = ["duckdb"] +tui = ["textual", "textual-select", "pygments", "prompt-toolkit"] +dremio = ["sqlalchemy-dremio", "pandas", "pyarrow"] all = ["mysql-connector-python", "psycopg2-binary", "snowflake-connector-python", "cryptography", "presto-python-client", "cx_Oracle", "trino", "clickhouse-driver", "vertica-python", "duckdb"] @@ -83,6 +119,7 @@ no_implicit_optional=false [tool.ruff] line-length = 120 +target-version = "py38" [tool.black] line-length = 120 From 87bc6af2ce2a51dbfe501f956fe7b6a27423d930 Mon Sep 17 00:00:00 2001 From: vmatt Date: Fri, 13 Feb 2026 12:06:42 +0100 Subject: [PATCH 03/11] moved --- {reladiff/sqeleton => sqeleton}/__init__.py | 0 {reladiff/sqeleton => sqeleton}/__main__.py | 0 {reladiff/sqeleton => sqeleton}/abcs/__init__.py | 0 {reladiff/sqeleton => sqeleton}/abcs/compiler.py | 0 {reladiff/sqeleton => sqeleton}/abcs/database_types.py | 0 {reladiff/sqeleton => sqeleton}/abcs/mixins.py | 0 {reladiff/sqeleton => sqeleton}/bound_exprs.py | 0 {reladiff/sqeleton => sqeleton}/conn_editor.py | 0 {reladiff/sqeleton => sqeleton}/databases/__init__.py | 0 {reladiff/sqeleton => sqeleton}/databases/_connect.py | 0 {reladiff/sqeleton => sqeleton}/databases/base.py | 0 {reladiff/sqeleton => sqeleton}/databases/bigquery.py | 0 {reladiff/sqeleton => sqeleton}/databases/clickhouse.py | 0 {reladiff/sqeleton => sqeleton}/databases/databricks.py | 0 {reladiff/sqeleton => sqeleton}/databases/dremio.py | 0 {reladiff/sqeleton => sqeleton}/databases/duckdb.py | 0 {reladiff/sqeleton => sqeleton}/databases/mssql.py | 0 {reladiff/sqeleton => sqeleton}/databases/mysql.py | 0 {reladiff/sqeleton => sqeleton}/databases/oracle.py | 0 {reladiff/sqeleton => sqeleton}/databases/postgresql.py | 0 {reladiff/sqeleton => sqeleton}/databases/presto.py | 0 {reladiff/sqeleton => sqeleton}/databases/redshift.py | 0 {reladiff/sqeleton => sqeleton}/databases/snowflake.py | 0 {reladiff/sqeleton => sqeleton}/databases/trino.py | 0 {reladiff/sqeleton => sqeleton}/databases/vertica.py | 0 {reladiff/sqeleton => sqeleton}/queries/__init__.py | 0 {reladiff/sqeleton => sqeleton}/queries/api.py | 0 {reladiff/sqeleton => sqeleton}/queries/ast_classes.py | 0 {reladiff/sqeleton => sqeleton}/queries/base.py | 0 {reladiff/sqeleton => sqeleton}/queries/compiler.py | 0 {reladiff/sqeleton => sqeleton}/queries/extras.py | 0 {reladiff/sqeleton => sqeleton}/query_utils.py | 0 {reladiff/sqeleton => sqeleton}/repl.py | 0 {reladiff/sqeleton => sqeleton}/schema.py | 0 {reladiff/sqeleton => sqeleton}/utils.py | 0 35 files changed, 0 insertions(+), 0 deletions(-) rename {reladiff/sqeleton => sqeleton}/__init__.py (100%) rename {reladiff/sqeleton => sqeleton}/__main__.py (100%) rename {reladiff/sqeleton => sqeleton}/abcs/__init__.py (100%) rename {reladiff/sqeleton => sqeleton}/abcs/compiler.py (100%) rename {reladiff/sqeleton => sqeleton}/abcs/database_types.py (100%) rename {reladiff/sqeleton => sqeleton}/abcs/mixins.py (100%) rename {reladiff/sqeleton => sqeleton}/bound_exprs.py (100%) rename {reladiff/sqeleton => sqeleton}/conn_editor.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/__init__.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/_connect.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/base.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/bigquery.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/clickhouse.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/databricks.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/dremio.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/duckdb.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/mssql.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/mysql.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/oracle.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/postgresql.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/presto.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/redshift.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/snowflake.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/trino.py (100%) rename {reladiff/sqeleton => sqeleton}/databases/vertica.py (100%) rename {reladiff/sqeleton => sqeleton}/queries/__init__.py (100%) rename {reladiff/sqeleton => sqeleton}/queries/api.py (100%) rename {reladiff/sqeleton => sqeleton}/queries/ast_classes.py (100%) rename {reladiff/sqeleton => sqeleton}/queries/base.py (100%) rename {reladiff/sqeleton => sqeleton}/queries/compiler.py (100%) rename {reladiff/sqeleton => sqeleton}/queries/extras.py (100%) rename {reladiff/sqeleton => sqeleton}/query_utils.py (100%) rename {reladiff/sqeleton => sqeleton}/repl.py (100%) rename {reladiff/sqeleton => sqeleton}/schema.py (100%) rename {reladiff/sqeleton => sqeleton}/utils.py (100%) diff --git a/reladiff/sqeleton/__init__.py b/sqeleton/__init__.py similarity index 100% rename from reladiff/sqeleton/__init__.py rename to sqeleton/__init__.py diff --git a/reladiff/sqeleton/__main__.py b/sqeleton/__main__.py similarity index 100% rename from reladiff/sqeleton/__main__.py rename to sqeleton/__main__.py diff --git a/reladiff/sqeleton/abcs/__init__.py b/sqeleton/abcs/__init__.py similarity index 100% rename from reladiff/sqeleton/abcs/__init__.py rename to sqeleton/abcs/__init__.py diff --git a/reladiff/sqeleton/abcs/compiler.py b/sqeleton/abcs/compiler.py similarity index 100% rename from reladiff/sqeleton/abcs/compiler.py rename to sqeleton/abcs/compiler.py diff --git a/reladiff/sqeleton/abcs/database_types.py b/sqeleton/abcs/database_types.py similarity index 100% rename from reladiff/sqeleton/abcs/database_types.py rename to sqeleton/abcs/database_types.py diff --git a/reladiff/sqeleton/abcs/mixins.py b/sqeleton/abcs/mixins.py similarity index 100% rename from reladiff/sqeleton/abcs/mixins.py rename to sqeleton/abcs/mixins.py diff --git a/reladiff/sqeleton/bound_exprs.py b/sqeleton/bound_exprs.py similarity index 100% rename from reladiff/sqeleton/bound_exprs.py rename to sqeleton/bound_exprs.py diff --git a/reladiff/sqeleton/conn_editor.py b/sqeleton/conn_editor.py similarity index 100% rename from reladiff/sqeleton/conn_editor.py rename to sqeleton/conn_editor.py diff --git a/reladiff/sqeleton/databases/__init__.py b/sqeleton/databases/__init__.py similarity index 100% rename from reladiff/sqeleton/databases/__init__.py rename to sqeleton/databases/__init__.py diff --git a/reladiff/sqeleton/databases/_connect.py b/sqeleton/databases/_connect.py similarity index 100% rename from reladiff/sqeleton/databases/_connect.py rename to sqeleton/databases/_connect.py diff --git a/reladiff/sqeleton/databases/base.py b/sqeleton/databases/base.py similarity index 100% rename from reladiff/sqeleton/databases/base.py rename to sqeleton/databases/base.py diff --git a/reladiff/sqeleton/databases/bigquery.py b/sqeleton/databases/bigquery.py similarity index 100% rename from reladiff/sqeleton/databases/bigquery.py rename to sqeleton/databases/bigquery.py diff --git a/reladiff/sqeleton/databases/clickhouse.py b/sqeleton/databases/clickhouse.py similarity index 100% rename from reladiff/sqeleton/databases/clickhouse.py rename to sqeleton/databases/clickhouse.py diff --git a/reladiff/sqeleton/databases/databricks.py b/sqeleton/databases/databricks.py similarity index 100% rename from reladiff/sqeleton/databases/databricks.py rename to sqeleton/databases/databricks.py diff --git a/reladiff/sqeleton/databases/dremio.py b/sqeleton/databases/dremio.py similarity index 100% rename from reladiff/sqeleton/databases/dremio.py rename to sqeleton/databases/dremio.py diff --git a/reladiff/sqeleton/databases/duckdb.py b/sqeleton/databases/duckdb.py similarity index 100% rename from reladiff/sqeleton/databases/duckdb.py rename to sqeleton/databases/duckdb.py diff --git a/reladiff/sqeleton/databases/mssql.py b/sqeleton/databases/mssql.py similarity index 100% rename from reladiff/sqeleton/databases/mssql.py rename to sqeleton/databases/mssql.py diff --git a/reladiff/sqeleton/databases/mysql.py b/sqeleton/databases/mysql.py similarity index 100% rename from reladiff/sqeleton/databases/mysql.py rename to sqeleton/databases/mysql.py diff --git a/reladiff/sqeleton/databases/oracle.py b/sqeleton/databases/oracle.py similarity index 100% rename from reladiff/sqeleton/databases/oracle.py rename to sqeleton/databases/oracle.py diff --git a/reladiff/sqeleton/databases/postgresql.py b/sqeleton/databases/postgresql.py similarity index 100% rename from reladiff/sqeleton/databases/postgresql.py rename to sqeleton/databases/postgresql.py diff --git a/reladiff/sqeleton/databases/presto.py b/sqeleton/databases/presto.py similarity index 100% rename from reladiff/sqeleton/databases/presto.py rename to sqeleton/databases/presto.py diff --git a/reladiff/sqeleton/databases/redshift.py b/sqeleton/databases/redshift.py similarity index 100% rename from reladiff/sqeleton/databases/redshift.py rename to sqeleton/databases/redshift.py diff --git a/reladiff/sqeleton/databases/snowflake.py b/sqeleton/databases/snowflake.py similarity index 100% rename from reladiff/sqeleton/databases/snowflake.py rename to sqeleton/databases/snowflake.py diff --git a/reladiff/sqeleton/databases/trino.py b/sqeleton/databases/trino.py similarity index 100% rename from reladiff/sqeleton/databases/trino.py rename to sqeleton/databases/trino.py diff --git a/reladiff/sqeleton/databases/vertica.py b/sqeleton/databases/vertica.py similarity index 100% rename from reladiff/sqeleton/databases/vertica.py rename to sqeleton/databases/vertica.py diff --git a/reladiff/sqeleton/queries/__init__.py b/sqeleton/queries/__init__.py similarity index 100% rename from reladiff/sqeleton/queries/__init__.py rename to sqeleton/queries/__init__.py diff --git a/reladiff/sqeleton/queries/api.py b/sqeleton/queries/api.py similarity index 100% rename from reladiff/sqeleton/queries/api.py rename to sqeleton/queries/api.py diff --git a/reladiff/sqeleton/queries/ast_classes.py b/sqeleton/queries/ast_classes.py similarity index 100% rename from reladiff/sqeleton/queries/ast_classes.py rename to sqeleton/queries/ast_classes.py diff --git a/reladiff/sqeleton/queries/base.py b/sqeleton/queries/base.py similarity index 100% rename from reladiff/sqeleton/queries/base.py rename to sqeleton/queries/base.py diff --git a/reladiff/sqeleton/queries/compiler.py b/sqeleton/queries/compiler.py similarity index 100% rename from reladiff/sqeleton/queries/compiler.py rename to sqeleton/queries/compiler.py diff --git a/reladiff/sqeleton/queries/extras.py b/sqeleton/queries/extras.py similarity index 100% rename from reladiff/sqeleton/queries/extras.py rename to sqeleton/queries/extras.py diff --git a/reladiff/sqeleton/query_utils.py b/sqeleton/query_utils.py similarity index 100% rename from reladiff/sqeleton/query_utils.py rename to sqeleton/query_utils.py diff --git a/reladiff/sqeleton/repl.py b/sqeleton/repl.py similarity index 100% rename from reladiff/sqeleton/repl.py rename to sqeleton/repl.py diff --git a/reladiff/sqeleton/schema.py b/sqeleton/schema.py similarity index 100% rename from reladiff/sqeleton/schema.py rename to sqeleton/schema.py diff --git a/reladiff/sqeleton/utils.py b/sqeleton/utils.py similarity index 100% rename from reladiff/sqeleton/utils.py rename to sqeleton/utils.py From 4ba742a6b306aefb58593a32d9ee3f482d610c39 Mon Sep 17 00:00:00 2001 From: vmatt Date: Fri, 13 Feb 2026 17:44:45 +0100 Subject: [PATCH 04/11] feat: add MD5 hashing for Text columns to avoid concatenation-too-long errors --- reladiff/hashdiff_tables.py | 34 ++++++++++++++++++++++++++++++++-- reladiff/table_segment.py | 25 +++++++++++++++++++++---- sqeleton/queries/extras.py | 27 ++++++++++++++++++++++++--- 3 files changed, 77 insertions(+), 9 deletions(-) diff --git a/reladiff/hashdiff_tables.py b/reladiff/hashdiff_tables.py index 7a6b4f4e..76653f68 100644 --- a/reladiff/hashdiff_tables.py +++ b/reladiff/hashdiff_tables.py @@ -1,4 +1,5 @@ import os +import re from functools import cmp_to_key from numbers import Number import logging @@ -27,6 +28,18 @@ logger = logging.getLogger("hashdiff_tables") +# Patterns for concatenation-too-long errors from various databases +_CONCAT_TOO_LONG_PATTERNS = [ + re.compile(r"ORA-01489", re.IGNORECASE), # Oracle: result of string concatenation is too long + re.compile(r"too long and would be truncated", re.IGNORECASE), # Snowflake +] + + +def _is_concat_too_long_error(e: Exception) -> bool: + """Check if the exception is a concatenation-too-long error from the database.""" + msg = str(e) + return any(p.search(msg) for p in _CONCAT_TOO_LONG_PATTERNS) + def compare_element(a, b): """Compare a and b, treat None as the smallest value. @@ -189,7 +202,20 @@ def _diff_segments( count1, count2 = self._threaded_call("count", [table1, table2]) checksum1 = checksum2 = None else: - (count1, checksum1), (count2, checksum2) = self._threaded_call("count_and_checksum", [table1, table2]) + try: + (count1, checksum1), (count2, checksum2) = self._threaded_call("count_and_checksum", [table1, table2]) + except Exception as e: + if _is_concat_too_long_error(e): + logger.warning( + "Concatenation too long error detected. Retrying with MD5-hashed text columns." + ) + table1 = table1.with_md5_text_columns() + table2 = table2.with_md5_text_columns() + (count1, checksum1), (count2, checksum2) = self._threaded_call( + "count_and_checksum", [table1, table2] + ) + else: + raise assert not info_tree.info.rowcounts info_tree.info.rowcounts = {1: count1, 2: count2} @@ -232,7 +258,11 @@ def _bisect_and_diff_segments( # If count is below the threshold, just download and compare the columns locally # This saves time, as bisection speed is limited by ping and query performance. if max_rows < self.bisection_threshold or max_space_size < self.bisection_factor * 2: - rows1, rows2 = self._threaded_call("get_values", [table1, table2]) + # get_values() downloads each column individually (no concatenation), + # so always use non-MD5 mode to get actual values for comparison. + plain_table1 = table1.new(_md5_text_columns=False) if table1._md5_text_columns else table1 + plain_table2 = table2.new(_md5_text_columns=False) if table2._md5_text_columns else table2 + rows1, rows2 = self._threaded_call("get_values", [plain_table1, plain_table2]) diff = list(diff_sets(rows1, rows2, self.skip_sort_results, self.duplicate_rows_support)) info_tree.info.set_diff(diff) diff --git a/reladiff/table_segment.py b/reladiff/table_segment.py index f64247c8..f5fdbbb5 100644 --- a/reladiff/table_segment.py +++ b/reladiff/table_segment.py @@ -8,10 +8,10 @@ from .utils import safezip, Vector from sqeleton.utils import ArithString, split_space from sqeleton.databases import Database, DbPath, DbKey, DbTime -from sqeleton.abcs.database_types import String_UUID +from sqeleton.abcs.database_types import String_UUID, Text from sqeleton.schema import Schema, create_schema from sqeleton.queries import Count, Checksum, SKIP, table, this, Expr, min_, max_, Code -from sqeleton.queries.extras import ApplyFuncAndNormalizeAsString, NormalizeAsString +from sqeleton.queries.extras import ApplyFuncAndNormalizeAsString, NormalizeAsString, Md5AsString logger = logging.getLogger("table_segment") @@ -126,6 +126,7 @@ class TableSegment: case_sensitive: bool = True _schema: Schema = None + _md5_text_columns: bool = False def __post_init__(self): if not self.update_column and (self.min_update or self.max_update): @@ -206,7 +207,10 @@ def make_select(self): def get_values(self) -> list: "Download all the relevant values of the segment from the database" - select = self.make_select().select(*self._relevant_columns_repr) + # get_values() downloads each column individually (no concatenation), + # so always use plain NormalizeAsString to get actual values. + plain_repr = [NormalizeAsString(this[c]) for c in self.relevant_columns] + select = self.make_select().select(*plain_repr) return self.database.query(select, List[Tuple]) def choose_checkpoints(self, count: int) -> List[List[DbKey]]: @@ -250,8 +254,18 @@ def relevant_columns(self) -> List[str]: @property def _relevant_columns_repr(self) -> List[Expr]: + if self._md5_text_columns and self._schema: + return [ + Md5AsString(this[c]) if isinstance(self._schema.get(c), Text) + else NormalizeAsString(this[c]) + for c in self.relevant_columns + ] return [NormalizeAsString(this[c]) for c in self.relevant_columns] + def with_md5_text_columns(self) -> "TableSegment": + """Return a copy of this segment that MD5-hashes Text columns to avoid concatenation length errors.""" + return self.new(_md5_text_columns=True) + def count(self) -> int: """Count how many rows are in the segment, in one pass.""" return self.database.query(self.make_select().select(Count()), int) @@ -330,7 +344,7 @@ def count_and_checksum(self) -> Tuple[int, int]: return (0, None) def __getattr__(self, attr): - assert attr in ("database", "key_columns", "key_types", "relevant_columns", "_schema") + assert attr in ("database", "key_columns", "key_types", "relevant_columns", "_schema", "_md5_text_columns") return getattr(self._table_segment, attr) @property @@ -357,5 +371,8 @@ def make_select(self): # XXX shouldn't be called return self._table_segment.make_select() + def with_md5_text_columns(self) -> "EmptyTableSegment": + return self + def get_values(self) -> list: return [] diff --git a/sqeleton/queries/extras.py b/sqeleton/queries/extras.py index 382407c4..fdc51dda 100644 --- a/sqeleton/queries/extras.py +++ b/sqeleton/queries/extras.py @@ -18,6 +18,13 @@ class NormalizeAsString(ExprNode): type = str +@dataclass +class Md5AsString(ExprNode): + """MD5-hash the expression and return as a normalized string. Used to avoid concatenation length errors.""" + expr: ExprNode + type = str + + @dataclass class ApplyFuncAndNormalizeAsString(ExprNode): expr: ExprNode @@ -57,12 +64,26 @@ def compile_node(c: Compiler, n: ApplyFuncAndNormalizeAsString) -> str: return c.compile(expr) + @md + def compile_node(c: Compiler, n: Md5AsString) -> str: + expr = c.compile(n.expr) + str_expr = c.dialect.to_string(expr) + coalesced = f"coalesce({str_expr}, '')" + return c.dialect.to_string(c.dialect.md5_as_int(coalesced)) + @md def compile_node(c: Compiler, n: Checksum) -> str: if len(n.exprs) > 1: - exprs = [Code(f"coalesce({c.compile(expr)}, '')") for expr in n.exprs] - # exprs = [c.compile(e) for e in exprs] - expr = Concat(exprs, "|") + compiled = [] + for expr in n.exprs: + compiled_expr = c.compile(expr) + # NormalizeAsString and Md5AsString already guarantee non-NULL string output, + # so skip the redundant coalesce wrapper for them. + if isinstance(expr, (NormalizeAsString, Md5AsString)): + compiled.append(Code(compiled_expr)) + else: + compiled.append(Code(f"coalesce({compiled_expr}, '')")) + expr = Concat(compiled, "|") else: # No need to coalesce - safe to assume that key cannot be null (expr,) = n.exprs From 29b312ed81e698c38ba4c5b5ddd017b0701d02ff Mon Sep 17 00:00:00 2001 From: vmatt Date: Sat, 14 Feb 2026 23:31:03 +0100 Subject: [PATCH 05/11] include sqeleton package --- pyproject.toml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 0aeb4e78..a4747095 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,10 @@ classifiers = [ "Topic :: Database :: Database Engines/Servers", "Typing :: Typed" ] -packages = [{ include = "reladiff" }] +packages = [ + { include = "reladiff" }, + { include = "sqeleton" } + ] [tool.poetry.dependencies] python = "^3.8" From 3f9e47d958198634125b3515e0d79bc000565826 Mon Sep 17 00:00:00 2001 From: vmatt Date: Mon, 16 Feb 2026 00:24:11 +0100 Subject: [PATCH 06/11] include dot in alphanums --- sqeleton/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sqeleton/utils.py b/sqeleton/utils.py index 1ceee4a2..84ff1909 100644 --- a/sqeleton/utils.py +++ b/sqeleton/utils.py @@ -143,7 +143,7 @@ def as_insensitive(self): # -- Alphanumerics -- -alphanums = " -" + string.digits + string.ascii_uppercase + "_" + string.ascii_lowercase +alphanums = " -." + string.digits + string.ascii_uppercase + "_" + string.ascii_lowercase class ArithString(ABC): From 7e2cde5a5426dfe16a6df7b51cb4d912dcd5b390 Mon Sep 17 00:00:00 2001 From: vmatt Date: Wed, 25 Feb 2026 17:57:55 +0100 Subject: [PATCH 07/11] reduce ts precision to 3 for relevant dbs, version bump --- pyproject.toml | 2 +- sqeleton/databases/duckdb.py | 2 +- sqeleton/databases/mysql.py | 4 ++-- sqeleton/databases/oracle.py | 2 +- sqeleton/databases/postgresql.py | 2 +- sqeleton/databases/snowflake.py | 2 +- 6 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a4747095..ebd72fa2 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "reladiff" -version = "0.6.0" +version = "0.6.3" description = "Command-line tool and Python library to efficiently diff rows across two different databases." authors = ["Erez Shinan "] license = "MIT" diff --git a/sqeleton/databases/duckdb.py b/sqeleton/databases/duckdb.py index 154997db..bf4423a9 100644 --- a/sqeleton/databases/duckdb.py +++ b/sqeleton/databases/duckdb.py @@ -53,7 +53,7 @@ def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: if coltype.rounds and coltype.precision > 0: return f"CONCAT(SUBSTRING(STRFTIME({value}::TIMESTAMP, '%Y-%m-%d %H:%M:%S.'),1,23), LPAD(((ROUND(strftime({value}::timestamp, '%f')::DECIMAL(15,7)/100000,{coltype.precision-1})*100000)::INT)::VARCHAR,6,'0'))" - return f"rpad(substring(strftime({value}::timestamp, '%Y-%m-%d %H:%M:%S.%f'),1,{TIMESTAMP_PRECISION_POS+coltype.precision}),26,'0')" + return f"rpad(substring(strftime({value}::timestamp, '%Y-%m-%d %H:%M:%S.%f'),1,{TIMESTAMP_PRECISION_POS+3}),26,'0')" def normalize_number(self, value: str, coltype: FractionalType) -> str: return self.to_string(f"{value}::DECIMAL(38, {coltype.precision})") diff --git a/sqeleton/databases/mysql.py b/sqeleton/databases/mysql.py index 75522fec..e9b6dc29 100644 --- a/sqeleton/databases/mysql.py +++ b/sqeleton/databases/mysql.py @@ -37,10 +37,10 @@ def md5_as_int(self, s: str) -> str: class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: if coltype.rounds: - return self.to_string(f"cast( cast({value} as datetime({coltype.precision})) as datetime(6))") + return self.to_string(f"cast( cast({value} as datetime({coltype.precision})) as datetime(3))") s = self.to_string(f"cast({value} as datetime(6))") - return f"RPAD(RPAD({s}, {TIMESTAMP_PRECISION_POS+coltype.precision}, '.'), {TIMESTAMP_PRECISION_POS+6}, '0')" + return f"RPAD(RPAD({s}, {TIMESTAMP_PRECISION_POS+coltype.precision}, '.'), {TIMESTAMP_PRECISION_POS+3}, '0')" def normalize_number(self, value: str, coltype: FractionalType) -> str: return self.to_string(f"cast({value} as decimal(38, {coltype.precision}))") diff --git a/sqeleton/databases/oracle.py b/sqeleton/databases/oracle.py index 913b393a..c2d4d12b 100644 --- a/sqeleton/databases/oracle.py +++ b/sqeleton/databases/oracle.py @@ -60,7 +60,7 @@ def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: truncated = f"to_char({value}, 'YYYY-MM-DD HH24:MI:SS.FF{coltype.precision}')" else: truncated = f"to_char({value}, 'YYYY-MM-DD HH24:MI:SS.')" - return f"RPAD({truncated}, {TIMESTAMP_PRECISION_POS+6}, '0')" + return f"RPAD({truncated}, {TIMESTAMP_PRECISION_POS+3}, '0')" def normalize_number(self, value: str, coltype: FractionalType) -> str: # FM999.9990 diff --git a/sqeleton/databases/postgresql.py b/sqeleton/databases/postgresql.py index 154cf651..ab905a80 100644 --- a/sqeleton/databases/postgresql.py +++ b/sqeleton/databases/postgresql.py @@ -42,7 +42,7 @@ def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: timestamp6 = f"to_char({value}::timestamp(6), 'YYYY-mm-dd HH24:MI:SS.US')" return ( - f"RPAD(LEFT({timestamp6}, {TIMESTAMP_PRECISION_POS+coltype.precision}), {TIMESTAMP_PRECISION_POS+6}, '0')" + f"RPAD(LEFT({timestamp6}, {TIMESTAMP_PRECISION_POS+min(coltype.precision, 3)}), {TIMESTAMP_PRECISION_POS+3}, '0')" ) def normalize_number(self, value: str, coltype: FractionalType) -> str: diff --git a/sqeleton/databases/snowflake.py b/sqeleton/databases/snowflake.py index c78a26e7..8ca40451 100644 --- a/sqeleton/databases/snowflake.py +++ b/sqeleton/databases/snowflake.py @@ -53,7 +53,7 @@ def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: else: timestamp = f"cast(convert_timezone('UTC', {value}) as timestamp({coltype.precision}))" - return f"to_char({timestamp}, 'YYYY-MM-DD HH24:MI:SS.FF6')" + return f"to_char({timestamp}, 'YYYY-MM-DD HH24:MI:SS.FF3')" def normalize_number(self, value: str, coltype: FractionalType) -> str: return self.to_string(f"cast({value} as decimal(38, {coltype.precision}))") From c593fa9e5ef73b5dbdcc3e605e0949702bdeb4d4 Mon Sep 17 00:00:00 2001 From: vmatt Date: Wed, 25 Feb 2026 18:06:01 +0100 Subject: [PATCH 08/11] trim for mysql instead of round, vbump --- pyproject.toml | 2 +- sqeleton/databases/mysql.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ebd72fa2..d1393ee0 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "reladiff" -version = "0.6.3" +version = "0.6.4" description = "Command-line tool and Python library to efficiently diff rows across two different databases." authors = ["Erez Shinan "] license = "MIT" diff --git a/sqeleton/databases/mysql.py b/sqeleton/databases/mysql.py index e9b6dc29..cce79917 100644 --- a/sqeleton/databases/mysql.py +++ b/sqeleton/databases/mysql.py @@ -37,7 +37,8 @@ def md5_as_int(self, s: str) -> str: class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: if coltype.rounds: - return self.to_string(f"cast( cast({value} as datetime({coltype.precision})) as datetime(3))") + s = self.to_string(f"cast({value} as datetime(6))") + return f"LEFT({s}, {TIMESTAMP_PRECISION_POS + 3})" s = self.to_string(f"cast({value} as datetime(6))") return f"RPAD(RPAD({s}, {TIMESTAMP_PRECISION_POS+coltype.precision}, '.'), {TIMESTAMP_PRECISION_POS+3}, '0')" From 3e31b2dbf3a84d2090d0600f79940048913b9193 Mon Sep 17 00:00:00 2001 From: vmatt Date: Mon, 2 Mar 2026 19:24:07 +0100 Subject: [PATCH 09/11] fix oracle rounding return, vbump --- pyproject.toml | 2 +- sqeleton/databases/oracle.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d1393ee0..e265d4e3 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "reladiff" -version = "0.6.4" +version = "0.6.5" description = "Command-line tool and Python library to efficiently diff rows across two different databases." authors = ["Erez Shinan "] license = "MIT" diff --git a/sqeleton/databases/oracle.py b/sqeleton/databases/oracle.py index c2d4d12b..9d16d658 100644 --- a/sqeleton/databases/oracle.py +++ b/sqeleton/databases/oracle.py @@ -54,7 +54,7 @@ def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str: def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: if coltype.rounds: - return f"to_char(cast({value} as timestamp({coltype.precision})), 'YYYY-MM-DD HH24:MI:SS.FF6')" + return f"to_char(cast({value} as timestamp({coltype.precision})), 'YYYY-MM-DD HH24:MI:SS.FF3')" if coltype.precision > 0: truncated = f"to_char({value}, 'YYYY-MM-DD HH24:MI:SS.FF{coltype.precision}')" From 008002cb7e2141816314c5db219e34e16677d214 Mon Sep 17 00:00:00 2001 From: vmatt Date: Fri, 6 Mar 2026 12:49:01 +0100 Subject: [PATCH 10/11] convert timestamp to no timezone explicitly for snowflake --- sqeleton/databases/snowflake.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/sqeleton/databases/snowflake.py b/sqeleton/databases/snowflake.py index 8ca40451..4ddf3fee 100644 --- a/sqeleton/databases/snowflake.py +++ b/sqeleton/databases/snowflake.py @@ -49,11 +49,14 @@ def md5_as_int(self, s: str) -> str: class Mixin_NormalizeValue(AbstractMixin_NormalizeValue): def normalize_timestamp(self, value: str, coltype: TemporalType) -> str: if coltype.rounds: - timestamp = f"to_timestamp(round(date_part(epoch_nanosecond, convert_timezone('UTC', {value})::timestamp(9))/1000000000, {coltype.precision}))" + # Round the epoch to the desired precision, then reconstruct timestamp + timestamp = f"to_timestamp_ntz(round(date_part(epoch_nanosecond, {value})::number / 1000000000, {coltype.precision}))" else: - timestamp = f"cast(convert_timezone('UTC', {value}) as timestamp({coltype.precision}))" + # Truncate by casting to lower precision (truncation, no rounding) + timestamp = f"cast({value} as timestamp_ntz({coltype.precision}))" - return f"to_char({timestamp}, 'YYYY-MM-DD HH24:MI:SS.FF3')" + # Always format with exactly 3 fractional digits to match MySQL output + return f"to_char({timestamp}::timestamp_ntz(3), 'YYYY-MM-DD HH24:MI:SS.FF3')" def normalize_number(self, value: str, coltype: FractionalType) -> str: return self.to_string(f"cast({value} as decimal(38, {coltype.precision}))") From 8596a1c9672ef585c410fa84c563738e2693d338 Mon Sep 17 00:00:00 2001 From: vmatt Date: Fri, 6 Mar 2026 12:55:33 +0100 Subject: [PATCH 11/11] bump version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index e265d4e3..4f7ded7b 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "reladiff" -version = "0.6.5" +version = "0.6.6" description = "Command-line tool and Python library to efficiently diff rows across two different databases." authors = ["Erez Shinan "] license = "MIT"