# SQL code generation

import copy

import itertools

import operator

from collections import OrderedDict

import re

import UserDict

import warnings

import dicts

import exc

import iters

import lists

import objects

import regexp

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 unesc_name(name, quote='"'):

    removed_ref = [False]

    name = strings.remove_prefix(quote, name, removed_ref)

    if removed_ref[0]:

        name = strings.remove_suffix(quote, name, removed_ref)

        assert removed_ref[0]

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

    return name

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

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

def lstrip(str_):

    '''Also removes comments.'''

    if str_.startswith('/*'): comment, sep, str_ = str_.partition('*/')

    return str_.lstrip()

##### 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 __str__(self): return clean_name(strings.repr_no_u(self))

##### Unparameterized code objects

class Code(BasicObject):

    def __init__(self, lang='sql'):

        self.lang = lang

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

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

class CustomCode(Code):

    def __init__(self, str_):

        Code.__init__(self)

        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 isinstance(value, Code): return 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):

        Code.__init__(self)

        self.expr = expr

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

##### Names

class Name(Code):

    def __init__(self, name):

        Code.__init__(self)

        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

#### Primitives

class Literal(Code):

    def __init__(self, value):

        Code.__init__(self)

        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 get_value(value):

    '''Unwraps a Literal's value'''

    value = remove_col_rename(value)

    if isinstance(value, Literal): return value.value

    else:

        assert not isinstance(value, Code)

        return value

def is_literal(value): return isinstance(value, Literal)

def is_null(value): return is_literal(value) and value.value == None

#### Composites

class List(Code):

    def __init__(self, values):

        Code.__init__(self)

        self.values = values

    def to_str(self, db): return ', '.join((v.to_str(db) for v in self.values))

class Tuple(List):

    def __init__(self, *values):

        List.__init__(self, values)

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

class Row(Tuple):

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

### Arrays

class Array(List):

    def __init__(self, values):

        values = map(remove_col_rename, values)

        List.__init__(self, values)

    def to_str(self, db): return 'ARRAY['+List.to_str(self, db)+']'

def to_Array(value):

    if isinstance(value, Array): return value

    return Array(lists.mk_seq(value))

##### 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()

        '''

        Code.__init__(self)

        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.order_by = None

        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['order_by'] # ignore

        del compare_on['index_cols'] # ignore

        return compare_on

def is_underlying_table(table):

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

def table2regclass_text(db, table):

    assert isinstance(table, Table)

    return db.esc_value(table.to_str(db))

class NoUnderlyingTableException(Exception):

    def __init__(self, ref):

        Exception.__init__(self, 'for: '+strings.as_tt(strings.urepr(ref)))

        self.ref = ref

def underlying_table(table):

    table = remove_table_rename(table)

    if table != None and table.srcs:

        table, = table.srcs # for derived tables or row vars

    if not is_underlying_table(table): raise NoUnderlyingTableException(table)

    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(strings.ustr(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_col(col): return isinstance(col, Col)

def is_table_col(col): return is_col(col) and col.table != None

def col2col_ref(db, col):

    assert isinstance(col, Col)

    return db.esc_value((col.table.to_str(db), col.name))

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

    elif util.is_str(col): return Col(col, table)

    else: return Literal(col)

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

    elif isinstance(col, Col):

        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)

def has_srcs(col): return is_col(col) and col.srcs

def cross_join_srcs(cols):

    cols = filter(has_srcs, cols) # empty srcs will mess up the cross join

    srcs = [[s.name for s in c.srcs] for c in cols]

    if not srcs: return [] # itertools.product() returns [()] for empty input

    return [Col(','.join(s)) for s in itertools.product(*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(col)

    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.

    Anything that isn't a column is wrapped in a NamedCol with the key's column

    name by `as_Col(value, name=key.name)`.

    '''

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

        dicts.DictProxy.__init__(self, OrderedDict())

        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:

            try: value = self.db.col_info(key).default

            except NoUnderlyingTableException: pass # not a table column

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

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

##### Definitions

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)+' '+as_Code(self.type).to_str(db)

        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)

class SetOf(Code):

    def __init__(self, type_):

        Code.__init__(self)

        self.type = type_

    def to_str(self, db):

        return 'SETOF '+self.type.to_str(db)

class RowType(Code):

    def __init__(self, table):

        Code.__init__(self)

        self.table = table

    def to_str(self, db):

        return self.table.to_str(db)+'%ROWTYPE'

class ColType(Code):

    def __init__(self, col):

        Code.__init__(self)

        self.col = col

    def to_str(self, db):

        return self.col.to_str(db)+'%TYPE'

class ArrayType(Code):

    def __init__(self, elem_type):

        Code.__init__(self)

        elem_type = as_Code(elem_type)

        self.elem_type = elem_type

    def to_str(self, db):

        return self.elem_type.to_str(db)+'[]'

##### Functions

Function = Table

as_Function = as_Table

class InternalFunction(CustomCode): pass

#### Calls

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

        '''

        Code.__init__(self)

        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]

#### Definitions

class FunctionDef(Code):

    def __init__(self, function, return_type, body, params=[], modifiers=None):

        Code.__init__(self)

        return_type = as_Code(return_type)

        body = as_Code(body)

        self.function = function

        self.return_type = return_type

        self.body = body

        self.params = params

        self.modifiers = modifiers

    def to_str(self, db):

        params_str = (', '.join((p.to_str(db) for p in self.params)))

        str_ = '''\

CREATE FUNCTION '''+self.function.to_str(db)+'''('''+params_str+''')

RETURNS '''+self.return_type.to_str(db)+'''

LANGUAGE '''+self.body.lang+'''

'''

        if self.modifiers != None: str_ += self.modifiers+'\n'

        str_ += '''\

AS $$

'''+self.body.to_str(db)+'''

$$;

'''

        return str_

class FunctionParam(TypedCol):

    def __init__(self, name, type_, default=None, out=False):

        TypedCol.__init__(self, name, type_, default)

        self.out = out

    def to_str(self, db):

        str_ = TypedCol.to_str(self, db)

        if self.out: str_ = 'OUT '+str_

        return str_

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

### PL/pgSQL

class ReturnQuery(Code):

    def __init__(self, query):

        Code.__init__(self)

        query = as_Code(query)

        self.query = query

    def to_str(self, db):

        return 'RETURN QUERY\n'+strings.indent(self.query.to_str(db))+';\n'

## Exceptions

class BaseExcHandler(BasicObject):

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

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

suppress_exc = 'NULL;\n';

reraise_exc = 'RAISE USING ERRCODE = SQLSTATE, MESSAGE = SQLERRM;\n';

class ExcHandler(BaseExcHandler):

    def __init__(self, exc, handler=None):

        if handler != None: handler = as_Code(handler)

        self.exc = exc

        self.handler = handler

    def to_str(self, db, body):

        body = as_Code(body)

        if self.handler != None:

            handler_str = '\n'+strings.indent(self.handler.to_str(db), 2)

        else: handler_str = ' '+suppress_exc

        str_ = '''\

BEGIN

'''+strings.indent(body.to_str(db))+'''\

EXCEPTION

    WHEN '''+self.exc+''' THEN'''+handler_str+'''\

END;\

'''

        return str_

class NestedExcHandler(BaseExcHandler):

    def __init__(self, *handlers):

        '''

        @param handlers Sorted from outermost to innermost

        '''

        self.handlers = handlers

    def to_str(self, db, body):

        for handler in reversed(self.handlers): body = handler.to_str(db, body)

        return body

class ExcToWarning(Code):

    def __init__(self, return_):

        '''

        @param return_ Statement to return a default value in case of error

        '''

        Code.__init__(self)

        return_ = as_Code(return_)

        self.return_ = return_

    def to_str(self, db):

        return "RAISE WARNING '%', SQLERRM;\n"+self.return_.to_str(db)

unique_violation_handler = ExcHandler('unique_violation')

# Note doubled "\"s because inside Python string

plpythonu_error_handler = ExcHandler('internal_error', '''\

-- Handle PL/Python exceptions

DECLARE

    matches text[] := regexp_matches(SQLERRM,

        E'^(?:PL/Python: )?(\\\\w+): (.*)$'); -- .* also matches \\n

    exc_name text := matches[1];

    msg text := matches[2];

BEGIN

    /* Re-raise PL/Python exceptions with the PL/Python prefix removed.

    This allows the exception to be parsed like a native exception.

    Always raise as data_exception so it goes in the errors table. */

    IF exc_name IS NOT NULL THEN

        RAISE data_exception USING MESSAGE = msg;

    -- Re-raise non-PL/Python exceptions

    ELSE

        '''+reraise_exc+'''\

    END IF;

END;

''')

def data_exception_handler(handler):

    return ExcHandler('data_exception', handler)

row_var = Table('row')

class RowExcIgnore(Code):

    def __init__(self, row_type, select_query, with_row, cols=None,

        exc_handler=unique_violation_handler, row_var=row_var):

        '''

        @param row_type Ignored if a custom row_var is used.

        @pre If a custom row_var is used, it must already be defined.

        '''

        Code.__init__(self, lang='plpgsql')

        row_type = as_Code(row_type)

        select_query = as_Code(select_query)

        with_row = as_Code(with_row)

        row_var = as_Table(row_var)

        self.row_type = row_type

        self.select_query = select_query

        self.with_row = with_row

        self.cols = cols

        self.exc_handler = exc_handler

        self.row_var = row_var

    def to_str(self, db):

        if self.cols == None: row_vars = [self.row_var]

        else: row_vars = [Col(c.name, self.row_var) for c in self.cols]

        # Need an EXCEPTION block for each individual row because "When an error

        # is caught by an EXCEPTION clause, [...] all changes to persistent

        # database state within the block are rolled back."

        # This is unfortunate because "A block containing an EXCEPTION clause is

        # significantly more expensive to enter and exit than a block without

        # one."

        # (http://www.postgresql.org/docs/8.3/static/\

        # plpgsql-control-structures.html#PLPGSQL-ERROR-TRAPPING)

        str_ = '''\

FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN

'''+strings.indent(self.select_query.to_str(db))+'''\

LOOP

'''+strings.indent(self.exc_handler.to_str(db, self.with_row))+'''\

END LOOP;

'''

        if self.row_var == row_var:

            str_ = '''\

DECLARE

    '''+self.row_var.to_str(db)+''' '''+self.row_type.to_str(db)+''';

BEGIN

'''+strings.indent(str_)+'''\