# SQL code generation

import copy

import operator

import re

import UserDict

import warnings

import dicts

import exc

import iters

import lists

import objects

import strings

import util

##### Names

identifier_max_len = 63 # works for both PostgreSQL and MySQL

def concat(str_, suffix):

    '''Preserves version so that it won't be truncated off the string, leading

    to collisions.'''

    # Preserve version

    match = re.match(r'^(.*?)((?:(?:#\d+)?\)?)*(?:\.\w+)?(?:::[\w ]+)*)$', str_)

    if match:

        str_, old_suffix = match.groups()

        suffix = old_suffix+suffix

    return strings.concat(str_, suffix, identifier_max_len)

def truncate(str_): return concat(str_, '')

def is_safe_name(name):

    '''A name is safe *and unambiguous* if it:

    * contains only *lowercase* word (\w) characters

    * doesn't start with a digit

    * contains "_", so that it's not a keyword

    '''

    return re.match(r'^(?=.*_)(?!\d)[^\WA-Z]+$', name)

def esc_name(name, quote='"'):

    return quote + name.replace(quote, quote+quote) + quote

        # doubling an embedded quote escapes it in both PostgreSQL and MySQL

def clean_name(name): return name.replace('"', '').replace('`', '')

def esc_comment(comment): return '/*'+comment.replace('*/', '* /')+'*/'

##### General SQL code objects

class MockDb:

    def esc_value(self, value): return strings.repr_no_u(value)

    def esc_name(self, name): return esc_name(name)

    def col_info(self, col):

        return TypedCol(col.name, '<type>', CustomCode('<default>'), True)

mockDb = MockDb()

class BasicObject(objects.BasicObject):

    def __init__(self, value): self.value = value

    def __str__(self): return clean_name(strings.repr_no_u(self))

##### Unparameterized code objects

class Code(BasicObject):

    def to_str(self, db): raise NotImplementedError()

    def __repr__(self): return self.to_str(mockDb)

class CustomCode(Code):

    def __init__(self, str_): self.str_ = str_

    def to_str(self, db): return self.str_

def as_Code(value, db=None):

    '''

    @param db If set, runs db.std_code() on the value.

    '''

    if util.is_str(value):

        if db != None: value = db.std_code(value)

        return CustomCode(value)

    else: return Literal(value)

class Expr(Code):

    def __init__(self, expr): self.expr = expr

    def to_str(self, db): return '('+self.expr.to_str(db)+')'

##### Names

class Name(Code):

    def __init__(self, name):

        name = truncate(name)

        self.name = name

    def to_str(self, db): return db.esc_name(self.name)

def as_Name(value):

    if isinstance(value, Code): return value

    else: return Name(value)

##### Literal values

class Literal(Code):

    def __init__(self, value): self.value = value

    def to_str(self, db): return db.esc_value(self.value)

def as_Value(value):

    if isinstance(value, Code): return value

    else: return Literal(value)

def is_null(value): return isinstance(value, Literal) and value.value == None

##### Derived elements

src_self = object() # tells Col that it is its own source column

class Derived(Code):

    def __init__(self, srcs):

        '''An element which was derived from some other element(s).

        @param srcs See self.set_srcs()

        '''

        self.set_srcs(srcs)

    def set_srcs(self, srcs, overwrite=True):

        '''

        @param srcs (self_type...)|src_self The element(s) this is derived from

        '''

        if not overwrite and self.srcs != (): return # already set

        if srcs == src_self: srcs = (self,)

        srcs = tuple(srcs) # make Col hashable

        self.srcs = srcs

    def _compare_on(self):

        compare_on = self.__dict__.copy()

        del compare_on['srcs'] # ignore

        return compare_on

def cols_srcs(cols): return lists.uniqify(iters.flatten((v.srcs for v in cols)))

##### Tables

class Table(Derived):

    def __init__(self, name, schema=None, srcs=(), is_temp=False):

        '''

        @param schema str|None (for no schema)

        @param srcs (Table...)|src_self See Derived.set_srcs()

        '''

        Derived.__init__(self, srcs)

        if util.is_str(name): name = truncate(name)

        self.name = name

        self.schema = schema

        self.is_temp = is_temp

        self.index_cols = {}

    def to_str(self, db):

        str_ = ''

        if self.schema != None: str_ += as_Name(self.schema).to_str(db)+'.'

        str_ += as_Name(self.name).to_str(db)

        return str_

    def to_Table(self): return self

    def _compare_on(self):

        compare_on = Derived._compare_on(self)

        del compare_on['index_cols'] # ignore

        return compare_on

def is_underlying_table(table):

    return isinstance(table, Table) and table.to_Table() is table

class NoUnderlyingTableException(Exception): pass

def underlying_table(table):

    table = remove_table_rename(table)

    if not is_underlying_table(table): raise NoUnderlyingTableException

    return table

def as_Table(table, schema=None):

    if table == None or isinstance(table, Code): return table

    else: return Table(table, schema)

def suffixed_table(table, suffix):

    table = copy.copy(table) # don't modify input!

    table.name = concat(table.name, suffix)

    return table

class NamedTable(Table):

    def __init__(self, name, code, cols=None):

        Table.__init__(self, name)

        code = as_Table(code)

        if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)

        if cols != None: cols = [to_name_only_col(c).to_Col() for c in cols]

        self.code = code

        self.cols = cols

    def to_str(self, db):

        str_ = self.code.to_str(db)

        if str_.find('\n') >= 0: whitespace = '\n'

        else: whitespace = ' '

        str_ += whitespace+'AS '+Table.to_str(self, db)

        if self.cols != None:

            str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'

        return str_

    def to_Table(self): return Table(self.name)

def remove_table_rename(table):

    if isinstance(table, NamedTable): table = table.code

    return table

##### Columns

class Col(Derived):

    def __init__(self, name, table=None, srcs=()):

        '''

        @param table Table|None (for no table)

        @param srcs (Col...)|src_self See Derived.set_srcs()

        '''

        Derived.__init__(self, srcs)

        if util.is_str(name): name = truncate(name)

        if util.is_str(table): table = Table(table)

        assert table == None or isinstance(table, Table)

        self.name = name

        self.table = table

    def to_str(self, db, for_str=False):

        str_ = as_Name(self.name).to_str(db)

        if for_str: str_ = clean_name(str_)

        if self.table != None:

            table = self.table.to_Table()

            if for_str: str_ = concat(str(table), '.'+str_)

            else: str_ = table.to_str(db)+'.'+str_

        return str_

    def __str__(self): return self.to_str(mockDb, for_str=True)

    def to_Col(self): return self

def is_table_col(col): return isinstance(col, Col) and col.table != None

def index_col(col):

    if not is_table_col(col): return None

    table = col.table

    try: name = table.index_cols[col.name]

    except KeyError: return None

    else: return Col(name, table, col.srcs)

def is_temp_col(col): return is_table_col(col) and col.table.is_temp

def as_Col(col, table=None, name=None):

    '''

    @param name If not None, any non-Col input will be renamed using NamedCol.

    '''

    if name != None:

        col = as_Value(col)

        if not isinstance(col, Col): col = NamedCol(name, col)

    if isinstance(col, Code): return col

    else: return Col(col, table)

def with_table(col, table):

    if isinstance(col, NamedCol): pass # doesn't take a table

    elif isinstance(col, FunctionCall):

        col = copy.deepcopy(col) # don't modify input!

        col.args[0].table = table

    else:

        col = copy.copy(col) # don't modify input!

        col.table = table

    return col

def with_default_table(col, table):

    col = as_Col(col)

    if col.table == None: col = with_table(col, table)

    return col

def set_cols_table(table, cols):

    table = as_Table(table)

    for i, col in enumerate(cols):

        col = cols[i] = as_Col(col)

        col.table = table

def to_name_only_col(col, check_table=None):

    col = as_Col(col)

    if not is_table_col(col): return col

    if check_table != None:

        table = col.table

        assert table == None or table == check_table

    return Col(col.name)

def suffixed_col(col, suffix):

    return Col(concat(col.name, suffix), col.table, col.srcs)

class NamedCol(Col):

    def __init__(self, name, code):

        Col.__init__(self, name)

        code = as_Value(code)

        self.code = code

    def to_str(self, db):

        return self.code.to_str(db)+' AS '+Col.to_str(self, db)

    def to_Col(self): return Col(self.name)

def remove_col_rename(col):

    if isinstance(col, NamedCol): col = col.code

    return col

def underlying_col(col):

    col = remove_col_rename(col)

    if not isinstance(col, Col): raise NoUnderlyingTableException

    return Col(col.name, underlying_table(col.table), col.srcs)

def wrap(wrap_func, value):

    '''Wraps a value, propagating any column renaming to the returned value.'''

    if isinstance(value, NamedCol):

        return NamedCol(value.name, wrap_func(value.code))

    else: return wrap_func(value)

class ColDict(dicts.DictProxy):

    '''A dict that automatically makes inserted entries Col objects'''

    def __init__(self, db, keys_table, dict_={}):

        dicts.DictProxy.__init__(self, {})

        keys_table = as_Table(keys_table)

        self.db = db

        self.table = keys_table

        self.update(dict_) # after setting vars because __setitem__() needs them

    def copy(self): return ColDict(self.db, self.table, self.inner.copy())

    def __getitem__(self, key):

        return dicts.DictProxy.__getitem__(self, self._key(key))

    def __setitem__(self, key, value):

        key = self._key(key)

        if value == None: value = self.db.col_info(key).default

        dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))

    def _key(self, key): return as_Col(key, self.table)

##### Functions

Function = Table

as_Function = as_Table

class InternalFunction(CustomCode): pass

class NamedArg(NamedCol):

    def __init__(self, name, value):

        NamedCol.__init__(self, name, value)

    def to_str(self, db):

        return Col.to_str(self, db)+' := '+self.code.to_str(db)

class FunctionCall(Code):

    def __init__(self, function, *args, **kw_args):

        '''

        @param args [Code|literal-value...] The function's arguments

        '''

        function = as_Function(function)

        def filter_(arg): return remove_col_rename(as_Value(arg))

        args = map(filter_, args)

        args += [NamedArg(k, filter_(v)) for k, v in kw_args.iteritems()]

        self.function = function

        self.args = args

    def to_str(self, db):

        args_str = ', '.join((v.to_str(db) for v in self.args))

        return self.function.to_str(db)+'('+args_str+')'

def wrap_in_func(function, value):

    '''Wraps a value inside a function call.

    Propagates any column renaming to the returned value.

    '''

    return wrap(lambda v: FunctionCall(function, v), value)

def unwrap_func_call(func_call, check_name=None):

    '''Unwraps any function call to its first argument.

    Also removes any column renaming.

    '''

    func_call = remove_col_rename(func_call)

    if not isinstance(func_call, FunctionCall): return func_call

    if check_name != None:

        name = func_call.function.name

        assert name == None or name == check_name

    return func_call.args[0]

##### Casts

class Cast(FunctionCall):

    def __init__(self, type_, value):

        value = as_Value(value)

        self.type_ = type_

        self.value = value

    def to_str(self, db):

        return 'CAST('+self.value.to_str(db)+' AS '+self.type_+')'

##### Conditions

class ColValueCond(Code):

    def __init__(self, col, value):

        value = as_ValueCond(value)

        self.col = col

        self.value = value

    def to_str(self, db): return self.value.to_str(db, self.col)

def combine_conds(conds, keyword=None):

    '''

    @param keyword The keyword to add before the conditions, if any

    '''

    str_ = ''

    if keyword != None:

        if conds == []: whitespace = ''

        elif len(conds) == 1: whitespace = ' '

        else: whitespace = '\n'

        str_ += keyword+whitespace

    str_ += '\nAND '.join(conds)

    return str_

##### Condition column comparisons

class ValueCond(BasicObject):

    def __init__(self, value):

        value = remove_col_rename(as_Value(value))

        self.value = value

    def to_str(self, db, left_value):

        '''

        @param left_value The Code object that the condition is being applied on

        '''

        raise NotImplemented()

    def __repr__(self): return self.to_str(mockDb, '<left_value>')

class CompareCond(ValueCond):

    def __init__(self, value, operator='='):

        '''

        @param operator By default, compares NULL values literally. Use '~=' or

            '~!=' to pass NULLs through.

        '''

        ValueCond.__init__(self, value)

        self.operator = operator

    def to_str(self, db, left_value):

        left_value = remove_col_rename(as_Col(left_value))

        right_value = self.value

        # Parse operator

        operator = self.operator

        passthru_null_ref = [False]

        operator = strings.remove_prefix('~', operator, passthru_null_ref)

        neg_ref = [False]

        operator = strings.remove_prefix('!', operator, neg_ref)

        equals = operator.endswith('=') # also includes <=, >=

        # Handle nullable columns

        check_null = False

        if not passthru_null_ref[0]: # NULLs compare equal

            try: left_value = ensure_not_null(db, left_value)

            except ensure_not_null_excs: # fall back to alternate method

                check_null = equals and isinstance(right_value, Col)

            else:

                if isinstance(left_value, EnsureNotNull):

                    right_value = ensure_not_null(db, right_value,

                        left_value.type) # apply same function to both sides

        if equals and is_null(right_value): operator = 'IS'

        left = left_value.to_str(db)

        right = right_value.to_str(db)

        # Create str

        str_ = left+' '+operator+' '+right

        if check_null:

            str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'

        if neg_ref[0]: str_ = 'NOT '+str_

        return str_

# Tells as_ValueCond() to assume a non-ValueCond is a literal value

assume_literal = object()

def as_ValueCond(value, default_table=assume_literal):

    if not isinstance(value, ValueCond):

        if default_table is not assume_literal:

            value = with_default_table(value, default_table)

        return CompareCond(value)

    else: return value

##### Joins

join_same = object() # tells Join the left and right columns have the same name

# Tells Join the left and right columns have the same name and are never NULL

join_same_not_null = object()

filter_out = object() # tells Join to filter out rows that match the join

class Join(BasicObject):

    def __init__(self, table, mapping={}, type_=None):

        '''

        @param mapping dict(right_table_col=left_table_col, ...)

            * if left_table_col is join_same: left_table_col = right_table_col

              * Note that right_table_col must be a string

            * if left_table_col is join_same_not_null:

              left_table_col = right_table_col and both have NOT NULL constraint

              * Note that right_table_col must be a string

        @param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out

            * filter_out: equivalent to 'LEFT' with the query filtered by

              `table_pkey IS NULL` (indicating no match)

        '''

        if util.is_str(table): table = Table(table)

        assert type_ == None or util.is_str(type_) or type_ is filter_out

        self.table = table

        self.mapping = mapping

        self.type_ = type_

    def to_str(self, db, left_table_):

        def join(entry):

            '''Parses non-USING joins'''

            right_table_col, left_table_col = entry

            # Switch order (right_table_col is on the left in the comparison)

            left = right_table_col

            right = left_table_col

            left_table = self.table

            right_table = left_table_

            # Parse left side

            left = with_default_table(left, left_table)

            # Parse special values

            left_on_right = Col(left.name, right_table)

            if right is join_same: right = left_on_right

            elif right is join_same_not_null:

                right = CompareCond(left_on_right, '~=')

            # Parse right side

            right = as_ValueCond(right, right_table)

            return right.to_str(db, left)

        # Create join condition

        type_ = self.type_

        joins = self.mapping

        if joins == {}: join_cond = None

        elif type_ is not filter_out and reduce(operator.and_,

            (v is join_same_not_null for v in joins.itervalues())):

            # all cols w/ USING, so can use simpler USING syntax

            cols = map(to_name_only_col, joins.iterkeys())

            join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'

        else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')

        if isinstance(self.table, NamedTable): whitespace = '\n'

        else: whitespace = ' '

        # Create join

        if type_ is filter_out: type_ = 'LEFT'

        str_ = ''

        if type_ != None: str_ += type_+' '

        str_ += 'JOIN'+whitespace+self.table.to_str(db)

        if join_cond != None: str_ += whitespace+join_cond

        return str_

    def __repr__(self): return self.to_str(mockDb, '<left_table>')

##### Value exprs

all_cols = CustomCode('*')

default = CustomCode('DEFAULT')

row_count = FunctionCall(InternalFunction('COUNT'), all_cols)

class Coalesce(FunctionCall):

    def __init__(self, *args):

        FunctionCall.__init__(self, InternalFunction('COALESCE'), *args)

class Nullif(FunctionCall):

    def __init__(self, *args):

        FunctionCall.__init__(self, InternalFunction('NULLIF'), *args)

# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>

null_sentinels = {

    'character varying': r'\N',

    'double precision': 'NaN',

    'integer': 2147483647,

    'text': r'\N',

    'timestamp with time zone': 'infinity'

}

class EnsureNotNull(Coalesce):

    def __init__(self, value, type_):

        Coalesce.__init__(self, as_Col(value),

            Cast(type_, null_sentinels[type_]))

        self.type = type_

    def to_str(self, db):

        col = self.args[0]

        index_col_ = index_col(col)

        if index_col_ != None: return index_col_.to_str(db)

        return Coalesce.to_str(self, db)

##### Table exprs

class Values(Code):

    def __init__(self, values):

        '''

        @param values [...]|[[...], ...] Can be one or multiple rows.

        '''

        rows = values

        if len(values) >= 1 and not lists.is_seq(values[0]): # only one row

            rows = [values]

        for i, row in enumerate(rows):

            rows[i] = map(remove_col_rename, map(as_Value, row))

        self.rows = rows

    def to_str(self, db):

        def row_str(row):

            return '('+(', '.join((v.to_str(db) for v in row)))+')'

        return 'VALUES '+(', '.join(map(row_str, self.rows)))

def NamedValues(name, cols, values):

    '''

    @param cols None|[...]

    @post `cols` will be changed to Col objects with the table set to `name`.

    '''

    table = NamedTable(name, Values(values), cols)

    if cols != None: set_cols_table(table, cols)

    return table

##### Database structure

class TypedCol(Col):

    def __init__(self, name, type_, default=None, nullable=True,

        constraints=None):

        assert default == None or isinstance(default, Code)

        Col.__init__(self, name)

        self.type = type_

        self.default = default

        self.nullable = nullable

        self.constraints = constraints

    def to_str(self, db):

        str_ = Col.to_str(self, db)+' '+self.type

        if not self.nullable: str_ += ' NOT NULL'

        if self.default != None: str_ += ' DEFAULT '+self.default.to_str(db)

        if self.constraints != None: str_ += ' '+self.constraints

        return str_

    def to_Col(self): return Col(self.name)

ensure_not_null_excs = (NoUnderlyingTableException, KeyError)

def ensure_not_null(db, col, type_=None):

    '''

    @param col If type_ is not set, must have an underlying column.

    @param type_ If set, overrides the underlying column's type.

    @return EnsureNotNull|Col

    @throws ensure_not_null_excs

    '''

    nullable = True

    try: typed_col = db.col_info(underlying_col(col))

    except NoUnderlyingTableException:

        if type_ == None: raise

    else:

        if type_ == None: type_ = typed_col.type

        nullable = typed_col.nullable

    if nullable:

        try: col = EnsureNotNull(col, type_)

        except KeyError, e:

            # Warn of no null sentinel for type, even if caller catches error

            warnings.warn(UserWarning(exc.str_(e)))

            raise

    return col