sqlscope 1.0.0__py3-none-any.whl

sqlscope/query/set_operations/set_operation.py

@@ -0,0 +1,59 @@
+import sqlglot
+from sqlglot import exp
+from ...catalog import Table, Column
+
+from abc import ABC, abstractmethod
+
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from .select import Select
+
+
+class SetOperation(ABC):
+    '''
+    Abstract base class for SQL set operations (i.e., SELECT, UNION, INTERSECT, EXCEPT).
+    '''
+
+    def __init__(self, sql: str, parent_query: 'Select | None' = None) -> None:
+        self.sql = sql
+        '''The SQL string representing the operation.'''
+        
+        self.parent_query = parent_query
+        '''The parent Select if this is a subquery.'''
+
+    @property
+    @abstractmethod
+    def output(self) -> Table:
+        '''Returns the output table schema of the set operation.'''
+        pass
+
+    @property
+    @abstractmethod
+    def referenced_tables(self) -> list[Table]:
+        '''Returns a list of tables that are referenced in the SQL query.'''
+        pass
+    
+    def __repr__(self, pre: str = '') -> str:
+        return f'{pre}{self.__class__.__name__}'
+
+    
+    @abstractmethod
+    def print_tree(self, pre: str = '') -> None:
+        pass
+    
+    @property
+    @abstractmethod
+    def main_selects(self) -> list['Select']:
+        '''Returns a list of selects that are part of a set operation.'''
+        return []
+
+    @property
+    @abstractmethod
+    def selects(self) -> list['Select']:
+        '''Returns a list of all Select nodes in the tree.'''
+        return []
+    
+    
+
+    
+

sqlscope/query/smt.py

@@ -0,0 +1,334 @@
+'''Convert SQL expressions to Z3 expressions for logical reasoning.'''
+
+from typing import Any, Callable
+from sqlglot import exp
+from z3 import (
+    Int, IntVal,
+    Real, RealVal,
+    Bool, BoolVal,
+    String, StringVal,
+    And, Or, Not,
+    Solver,
+    unsat,
+    is_expr,
+    BoolSort,
+    ExprRef,
+    Re,
+    AllChar,
+    Concat,
+    InRe,
+    PrefixOf,
+    SuffixOf,
+    Contains,
+)
+
+from ..catalog import Table
+
+
+# ----------------------------------------------------------------------
+#  Z3 variable creation
+# ----------------------------------------------------------------------
+
+def create_z3_var(variables: dict[str, Any], table_name: str | None,
+                  col_name: str, col_type: Callable[[str], ExprRef] | None = None) -> None:
+    '''
+    Create a Z3 variable for the given column name and type, and add it to the
+    variables dictionary. If col_type is None, default to Int.
+    '''
+    if col_type is None:
+        col_type = Int  # default type
+
+    # unqualified
+    variables[col_name] = col_type(col_name)
+    variables[f'{col_name}_isnull'] = Bool(f'{col_name}_isnull')
+
+    # qualified
+    if table_name:
+        variables[f'{table_name}.{col_name}'] = col_type(f'{table_name}.{col_name}')
+        variables[f'{table_name}.{col_name}_isnull'] = Bool(f'{table_name}.{col_name}_isnull')
+
+
+def fresh_symbol(prefix: str, sort: str):
+    '''Generate a fresh Z3 symbol with the given prefix and sort.'''
+    if sort == 'int':
+        return Int(f'{prefix}_{id(prefix)}')
+    if sort == 'real':
+        return Real(f'{prefix}_{id(prefix)}')
+    if sort == 'bool':
+        return Bool(f'{prefix}_{id(prefix)}')
+    return String(f'{prefix}_{id(prefix)}')
+
+
+# ----------------------------------------------------------------------
+#  Infer expected type of a subquery based on parent expression
+# ----------------------------------------------------------------------
+
+def infer_subquery_sort_from_parent(expr) -> str:
+    '''
+    Infer the expected Z3 sort of a subquery based on its parent expression.
+    '''
+    parent = expr.parent
+
+    # Arithmetic context → numeric
+    if isinstance(parent, (exp.Add, exp.Sub, exp.Mul, exp.Div, exp.Mod, exp.Pow)):
+        return 'real'
+
+    # Comparison context → numeric
+    if isinstance(parent, (exp.GT, exp.GTE, exp.LT, exp.LTE)):
+        return 'real'
+
+    # BETWEEN → numeric
+    if isinstance(parent, exp.Between):
+        return 'real'
+
+    # LIKE → string
+    if isinstance(parent, exp.Like):
+        return 'string'
+
+    # String concatenation (|| operator)
+    if isinstance(parent, exp.Concat):
+        return 'string'
+
+    # Default: boolean (EXISTS, WHERE (...))
+    return 'bool'
+
+
+# ----------------------------------------------------------------------
+#  Catalog → Z3 vars
+# ----------------------------------------------------------------------
+
+def catalog_table_to_z3_vars(table: Table) -> dict[str, ExprRef]:
+    '''Convert catalog table columns to Z3 variables.'''
+    variables = {}
+    for column in table.columns:
+        col_name = column.name
+        col_type = column.column_type.upper()
+
+        if col_type in ('INT', 'INTEGER', 'BIGINT', 'SMALLINT'):
+            create_z3_var(variables, table.name, col_name, Int)
+        elif col_type in ('FLOAT', 'REAL', 'DOUBLE'):
+            create_z3_var(variables, table.name, col_name, Real)
+        elif col_type in ('BOOLEAN', 'BOOL'):
+            create_z3_var(variables, table.name, col_name, Bool)
+        elif col_type in ('VARCHAR', 'CHAR', 'TEXT', 'CHARACTER VARYING'):
+            create_z3_var(variables, table.name, col_name, String)
+        else:
+            create_z3_var(variables, table.name, col_name)
+    return variables
+
+
+# ----------------------------------------------------------------------
+#  SQL → Z3 conversion
+# ----------------------------------------------------------------------
+
+def sql_to_z3(expr, variables: dict[str, ExprRef] = {}) -> Any:
+    '''Convert a SQLGlot expression to a Z3 expression.'''
+
+    # --- Columns ---
+    if isinstance(expr, exp.Column):
+        name = expr.name.lower()
+        if name not in variables:
+            create_z3_var(variables, None, name)
+        return variables[name]
+
+    # --- Literals ---
+    elif isinstance(expr, exp.Literal):
+        val = expr.this
+        if expr.is_int:
+            return IntVal(int(val))
+        elif expr.is_number:
+            return RealVal(float(val))
+        elif expr.is_string:
+            return StringVal(val.strip("'"))
+        elif val.upper() in ('TRUE', 'FALSE'):
+            return BoolVal(val.upper() == 'TRUE')
+        elif val.upper() == 'NULL':
+            return None
+        else:
+            raise NotImplementedError(f"Unsupported literal: {val}")
+
+    elif isinstance(expr, exp.Null):
+        return None
+
+    # --- Boolean comparisons ---
+    elif isinstance(expr, exp.EQ):
+        return sql_to_z3(expr.left, variables) == sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.NEQ):
+        return sql_to_z3(expr.left, variables) != sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.GT):
+        return sql_to_z3(expr.left, variables) > sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.GTE):
+        return sql_to_z3(expr.left, variables) >= sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.LT):
+        return sql_to_z3(expr.left, variables) < sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.LTE):
+        return sql_to_z3(expr.left, variables) <= sql_to_z3(expr.right, variables)
+
+    # --- Logical connectives ---
+    elif isinstance(expr, exp.And):
+        return And(sql_to_z3(expr.left, variables), sql_to_z3(expr.right, variables))
+    elif isinstance(expr, exp.Or):
+        return Or(sql_to_z3(expr.left, variables), sql_to_z3(expr.right, variables))
+    elif isinstance(expr, exp.Not):
+        return Not(sql_to_z3(expr.this, variables))
+    elif isinstance(expr, exp.Paren):
+        return sql_to_z3(expr.this, variables)
+
+    # --- Arithmetic ---
+    elif isinstance(expr, exp.Add):
+        return sql_to_z3(expr.left, variables) + sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.Sub):
+        return sql_to_z3(expr.left, variables) - sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.Mul):
+        return sql_to_z3(expr.left, variables) * sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.Div):
+        return sql_to_z3(expr.left, variables) / sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.Mod):
+        return sql_to_z3(expr.left, variables) % sql_to_z3(expr.right, variables)
+    elif isinstance(expr, exp.Pow):
+        return sql_to_z3(expr.left, variables) ** sql_to_z3(expr.right, variables)
+
+    # --- BETWEEN ---
+    elif isinstance(expr, exp.Between):
+        target = sql_to_z3(expr.this, variables)
+        low = sql_to_z3(expr.args['low'], variables)
+        high = sql_to_z3(expr.args['high'], variables)
+        return And(target >= low, target <= high)
+
+    # --- IN ---
+    elif isinstance(expr, exp.In):
+        target = sql_to_z3(expr.this, variables)
+
+        if isinstance(expr.args.get('query'), exp.Subquery):
+            # subquery → symbolic value
+            sym = fresh_symbol('subq_in', 'string')
+            return target == sym
+
+        options = [sql_to_z3(e, variables) for e in expr.expressions]
+
+        return Or(*[target == o for o in options])
+
+    # --- IS / IS NOT ---
+    elif isinstance(expr, exp.Is):
+        target_expr = expr.this
+        right_expr = expr.args.get('expression')
+
+        if isinstance(right_expr, exp.Null):
+            if isinstance(target_expr, exp.Column):
+                name = target_expr.name.lower()
+                flag = variables.setdefault(f'{name}_isnull', Bool(f'{name}_isnull'))
+                return flag
+            return BoolVal(False)
+
+        if isinstance(right_expr, exp.Not) and isinstance(right_expr.this, exp.Null):
+            if isinstance(target_expr, exp.Column):
+                name = target_expr.name.lower()
+                flag = variables.setdefault(f'{name}_isnull', Bool(f'{name}_isnull'))
+                return Not(flag)
+            return BoolVal(True)
+
+        return sql_to_z3(target_expr, variables) == sql_to_z3(right_expr, variables)
+
+    # --- LIKE ---
+    elif isinstance(expr, exp.Like):
+        target = sql_to_z3(expr.this, variables)
+        pattern_expr = sql_to_z3(expr.expression, variables)
+
+        # If pattern is a variable → fallback
+        if not isinstance(expr.expression, exp.Literal):
+            return target == pattern_expr
+
+        pattern = expr.expression.this.strip("'")
+        wildcard_count = pattern.count('%') + pattern.count('_')
+
+        if wildcard_count > 2:
+            return target == StringVal(pattern)
+
+        if '%' in pattern and '_' not in pattern:
+            # PREFIX pattern: abc%
+            if pattern.endswith('%') and pattern.count('%') == 1:
+                prefix = pattern[:-1]
+                return PrefixOf(StringVal(prefix), target)
+
+            # CONTAINS: %abc%
+            if pattern.startswith('%') and pattern.endswith('%') and pattern.count('%') == 2:
+                mid = pattern[1:-1]
+                return Contains(target, StringVal(mid))
+
+            # SUFFIX: %abc
+            if pattern.startswith('%') and pattern.count('%') == 1:
+                suffix = pattern[1:]
+                return SuffixOf(StringVal(suffix), target)
+
+        # EXACTLY ONE '_' wildcard
+        if '_' in pattern and '%' not in pattern and wildcard_count == 1:
+            parts = pattern.split('_')
+            regex = None
+            for i, p in enumerate(parts):
+                r = Re(StringVal(p))
+                regex = r if regex is None else Concat(regex, r)
+                if i < len(parts) - 1:
+                    regex = Concat(regex, AllChar(r.sort()))
+            return InRe(target, regex)
+
+        return target == StringVal(pattern)
+
+    # --- EXISTS ---
+    elif isinstance(expr, exp.Exists):
+        return fresh_symbol('subq_exists', 'bool')
+
+    # --- SUBQUERY ---
+    elif isinstance(expr, exp.Subquery):
+        sort = infer_subquery_sort_from_parent(expr)
+        if sort == 'int':
+            return fresh_symbol('subq_val', 'int')
+        elif sort == 'real':
+            return fresh_symbol('subq_val', 'real')
+        elif sort == 'string':
+            return fresh_symbol('subq_val', 'string')
+        else:
+            return fresh_symbol('subq_bool', 'bool')
+
+    # --- Fallback ---
+    return BoolVal(True)
+
+
+# ----------------------------------------------------------------------
+#  Formula checking
+# ----------------------------------------------------------------------
+
+def check_formula(expr) -> str:
+    '''Check if the given SQLGlot expression is a tautology, contradiction, or contingent.'''
+    
+    formula = sql_to_z3(expr, {})
+    
+    if formula is None:
+        return 'unknown'
+
+    solver = Solver()
+
+    solver.push()
+    solver.add(formula)
+
+    if solver.check() == unsat:
+        return 'contradiction'
+
+    solver.pop()
+    solver.push()
+    solver.add(Not(formula))
+
+    if solver.check() == unsat:
+        return 'tautology'
+
+    return 'contingent'
+
+def is_satisfiable(expr_z3) -> bool:
+
+    solver = Solver()
+    solver.add(expr_z3)
+    result = solver.check() != unsat
+
+    return result
+
+def is_bool_expr(e) -> bool:
+    return is_expr(e) and e.sort().kind() == BoolSort().kind()

sqlscope/query/tokenized_sql.py

@@ -0,0 +1,70 @@
+'''Query representation as a list of tokens. Works even for invalid SQL.'''
+
+import sqlparse
+import sqlparse.tokens
+from sqlparse.tokens import Whitespace, Newline
+
+from . import extractors
+
+class TokenizedSQL:
+    '''Base class for tokenizing SQL queries.'''
+
+    def __init__(self, sql: str) -> None:
+        self.sql = sql
+        '''The full SQL query string.'''
+
+        # Lazy properties
+        self._tokens = None
+        self._functions = None
+        self._comparisons = None
+        # End of lazy properties
+
+        parsed_statements = sqlparse.parse(self.sql)
+        if not parsed_statements:
+            self.all_statements: list[sqlparse.sql.Statement] = []
+            self.parsed = sqlparse.sql.Statement()
+        else:
+            self.all_statements = list(parsed_statements)
+            self.parsed = parsed_statements[0]
+
+    # region Properties
+    @property
+    def tokens(self) -> list[tuple[sqlparse.tokens._TokenType, str]]:
+        '''Returns a flattened list of tokens as (ttype, value) tuples, excluding whitespace and newlines.'''
+        if not self._tokens:
+            self._tokens = self._flatten()
+        return self._tokens
+
+    @property
+    def functions(self) -> list[tuple[sqlparse.sql.Function, str]]:
+        '''Returns a list of (function, clause) tuples found in the SQL query.'''
+
+        if self._functions is None:
+            self._functions = extractors.extract_functions(self.parsed.tokens)
+        return self._functions
+    
+    @property
+    def comparisons(self) -> list[tuple[sqlparse.sql.Comparison, str]]:
+        '''Returns a list of (comparison, clause) tuples found in the SQL query.'''
+        if self._comparisons is None:
+            self._comparisons = extractors.extract_comparisons(self.parsed.tokens)
+        return self._comparisons    
+
+    # endregion
+
+    def _flatten(self) -> list[tuple[sqlparse.tokens._TokenType, str]]:
+        '''Flattens the parsed SQL statement into a list of (ttype, value) tuples. Ignores whitespace and newlines.'''
+
+        if not self.parsed:
+            return []
+
+        # Flatten tokens into (ttype, value)
+        return [
+            (tok.ttype, tok.value) for tok in self.parsed.flatten()
+            if tok.ttype not in (Whitespace, Newline)
+        ]
+
+    def print_tree(self) -> None:
+        for stmt in self.all_statements:
+            stmt._pprint_tree()
+

sqlscope/query/typechecking/__init__.py

@@ -0,0 +1,21 @@
+from .base import get_type
+from . import primitives, functions, queries, unary_ops, binary_ops, predicates
+from ...catalog import Catalog
+from sqlglot import exp
+from sqlglot.optimizer.annotate_types import annotate_types
+from sqlglot.optimizer.qualify import qualify
+
+__all__ = ["get_type"]
+
+def rewrite_expression(expression: exp.Expression, catalog: Catalog, search_path: str = 'public') -> exp.Expression:
+    '''
+    Rewrites the expression by annotating types to its nodes based on the catalog.
+    '''
+
+    schema = catalog.to_sqlglot_schema()
+
+    return annotate_types(qualify(expression, schema=schema, db=search_path, validate_qualify_columns=False), schema)
+
+# This function needs to be called on a typed expression
+def collect_errors(expression: exp.Expression, catalog: Catalog, search_path: str = 'public') -> list[tuple[str, str, str | None]]:
+    return get_type(expression, catalog, search_path).messages

sqlscope/query/typechecking/base.py

@@ -0,0 +1,9 @@
+import sqlglot.expressions as exp
+from .types import AtomicType, ResultType
+from ...catalog import Catalog
+from functools import singledispatch
+
+@singledispatch
+def get_type(expression: exp.Expression, catalog: Catalog, search_path: str) -> ResultType:
+    '''Returns the type of the given SQL expression.'''
+    return AtomicType() # Default to unhandled expression

sqlscope/query/typechecking/binary_ops.py

@@ -0,0 +1,57 @@
+from .base import get_type
+from ...catalog import Catalog
+from sqlglot import exp
+from .types import ResultType, AtomicType
+from sqlglot.expressions import DataType
+from .util import is_number, to_number, to_date, error_message
+
+@get_type.register
+def _(expression: exp.Binary, catalog: Catalog, search_path: str) -> ResultType:
+    left_type = get_type(expression.this, catalog, search_path)
+    right_type = get_type(expression.expression, catalog, search_path)
+
+    old_messages = left_type.messages + right_type.messages
+    
+    # handle comparison operators
+    if isinstance(expression, exp.Predicate):
+        return typecheck_comparisons(left_type, right_type, expression, old_messages)
+    
+
+    if left_type != right_type:
+
+        if left_type.data_type != DataType.Type.UNKNOWN and not to_number(left_type) and left_type.data_type != DataType.Type.NULL:
+            old_messages.append(error_message(expression, left_type, "numeric"))
+
+        if right_type.data_type != DataType.Type.UNKNOWN and not to_number(right_type) and right_type.data_type != DataType.Type.NULL:
+            old_messages.append(error_message(expression, right_type, "numeric"))
+
+    elif DataType.Type.UNKNOWN != left_type.data_type and not is_number(left_type.data_type) and not is_number(right_type.data_type):
+        if left_type.data_type != DataType.Type.NULL or right_type.data_type != DataType.Type.NULL:
+            old_messages.append(error_message(expression, left_type, "numeric"))
+
+    return AtomicType(data_type=expression.type.this, nullable=left_type.nullable or right_type.nullable, constant=left_type.constant and right_type.constant, messages=old_messages)
+
+# handle comparison typechecking (e.g =, <, >, etc.)
+def typecheck_comparisons(left_type: ResultType, right_type: ResultType, expression: exp.Binary, old_messages: list) -> ResultType:
+
+    if DataType.Type.UNKNOWN in (left_type.data_type, right_type.data_type):
+        return AtomicType(data_type=expression.type.this,messages=old_messages)
+
+    # for boolean comparisons we can have only equality/inequality
+    if DataType.Type.BOOLEAN == left_type.data_type == right_type.data_type:
+        if not isinstance(expression, (exp.EQ, exp.NEQ)):
+            old_messages.append(error_message(expression, left_type, "boolean"))
+
+    if left_type != right_type and left_type.data_type != DataType.Type.NULL and right_type.data_type != DataType.Type.NULL:
+        
+        # handle implicit casts
+        if to_number(left_type) and to_number(right_type):
+            return AtomicType(data_type=expression.type.this, nullable=False, constant=True, messages=old_messages)
+
+        if to_date(left_type) and to_date(right_type):
+            return AtomicType(data_type=expression.type.this, nullable=False, constant=True, messages=old_messages)
+
+        old_messages.append(error_message(expression, left_type.data_type_str + " & " + right_type.data_type_str))
+
+    # Always returns boolean
+    return AtomicType(data_type=expression.type.this, nullable=False, constant=True, messages=old_messages)

sqlscope/query/typechecking/functions.py

@@ -0,0 +1,81 @@
+from .base import get_type
+from ...catalog import Catalog
+from sqlglot import exp
+from .types import ResultType, AtomicType
+from sqlglot.expressions import DataType
+from .util import is_number, error_message
+
+@get_type.register
+def _(expression: exp.Count, catalog: Catalog, search_path: str) -> ResultType:
+    old_messages = get_type(expression.this, catalog, search_path).messages
+
+    return AtomicType(data_type=expression.type.this, nullable=False, constant=True, messages=old_messages)
+
+@get_type.register
+def _(expression: exp.Avg, catalog: Catalog, search_path: str) -> ResultType:
+    inner_type = get_type(expression.this, catalog, search_path)
+
+    old_messages = inner_type.messages
+
+    if inner_type.data_type != DataType.Type.UNKNOWN and not is_number(inner_type.data_type):
+        old_messages.append(error_message(expression, inner_type, "NUMERIC"))
+
+    return AtomicType(data_type=expression.type.this, nullable=True, constant=True, messages=old_messages)
+
+@get_type.register
+def _(expression: exp.Sum, catalog: Catalog, search_path: str) -> ResultType:
+    inner_type = get_type(expression.this, catalog, search_path)
+
+    old_messages = inner_type.messages
+
+    if inner_type.data_type != DataType.Type.UNKNOWN and not is_number(inner_type.data_type):
+        old_messages.append(error_message(expression, inner_type, "NUMERIC"))
+
+    return AtomicType(data_type=expression.type.this, nullable=True, constant=True, messages=old_messages)
+
+@get_type.register
+def _(expression: exp.Min, catalog: Catalog, search_path: str) -> ResultType:
+    inner_type = get_type(expression.this, catalog, search_path)
+
+    old_messages = inner_type.messages
+
+    if inner_type.data_type != DataType.Type.UNKNOWN and inner_type.data_type == DataType.Type.BOOLEAN:
+        old_messages.append(error_message(expression, inner_type))
+
+    return AtomicType(data_type=inner_type.data_type, nullable=inner_type.nullable, constant=True, messages=old_messages)
+
+@get_type.register
+def _(expression: exp.Max, catalog: Catalog, search_path: str) -> ResultType:
+    inner_type = get_type(expression.this, catalog, search_path)
+
+    old_messages = inner_type.messages
+
+    if inner_type.data_type != DataType.Type.UNKNOWN and inner_type.data_type == DataType.Type.BOOLEAN:
+        old_messages.append(error_message(expression, inner_type))
+
+    return AtomicType(data_type=inner_type.data_type, nullable=inner_type.nullable, constant=True, messages=old_messages)
+
+@get_type.register
+def _(expression: exp.Concat, catalog: Catalog, search_path: str) -> ResultType:
+    old_messages = []
+    args_type = []
+
+    for arg in expression.expressions:
+        arg_type = get_type(arg, catalog, search_path)
+        if arg_type.messages:
+            old_messages.extend(arg_type.messages)
+        args_type.append(arg_type)
+        
+
+    if not args_type:
+        old_messages.append(error_message(expression, "Empty arguments"))
+    
+    # if all args are NULL, result is NULL
+    if all(target_type.data_type == DataType.Type.NULL for target_type in args_type):
+        return AtomicType(data_type=DataType.Type.NULL, constant=True, messages=old_messages)
+
+    constant = all(target_type.constant for target_type in args_type)
+    nullable = any(target_type.nullable for target_type in args_type)
+
+    # Always returns VARCHAR
+    return AtomicType(data_type=expression.type.this, nullable=nullable, constant=constant, messages=old_messages)