/lib/sql_gen.py - Annotate - BIEN 3 - NCEAS Projects

root/lib/sql_gen.py @ 3518

-            aaronmk
+# SQL code generation
-            aaronmk
+import copy
-            aaronmk
+import itertools
-            aaronmk
+import operator
-            aaronmk
+from ordereddict import OrderedDict
-            aaronmk
+import re
-            aaronmk
+import UserDict
-            aaronmk
+import warnings
             aaronmk
-            aaronmk
+import dicts
-            aaronmk
+import exc
-            aaronmk
+import iters
 import lists
-            aaronmk
+import objects
-            aaronmk
+import strings
-            aaronmk
+import util
             aaronmk
-            aaronmk
+##### Names
             aaronmk
-            aaronmk
+identifier_max_len = 63 # works for both PostgreSQL and MySQL
             aaronmk
-            aaronmk
+def concat(str_, suffix):
-            aaronmk
+    '''Preserves version so that it won't be truncated off the string, leading
     to collisions.'''
-            aaronmk
+    # Preserve version
-            aaronmk
+    match = re.match(r'^(.*?)((?:(?:#\d+)?\)?)*(?:\.\w+)?(?:::[\w ]+)*)$', str_)
-            aaronmk
+    if match:
         str_, old_suffix = match.groups()
         suffix = old_suffix+suffix
             aaronmk
-            aaronmk
+    return strings.concat(str_, suffix, identifier_max_len)
             aaronmk
-            aaronmk
+def truncate(str_): return concat(str_, '')
             aaronmk
-            aaronmk
+def is_safe_name(name):
-            aaronmk
+    '''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
-            aaronmk
+    '''
     return re.match(r'^(?=.*_)(?!\d)[^\WA-Z]+$', name)
             aaronmk
-            aaronmk
+def esc_name(name, quote='"'):
     return quote + name.replace(quote, quote+quote) + quote
         # doubling an embedded quote escapes it in both PostgreSQL and MySQL
-            aaronmk
+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
-            aaronmk
+def clean_name(name): return name.replace('"', '').replace('`', '')
-            aaronmk
+def esc_comment(comment): return '/*'+comment.replace('*/', '* /')+'*/'
-            aaronmk
+def lstrip(str_):
     '''Also removes comments.'''
     if str_.startswith('/*'): comment, sep, str_ = str_.partition('*/')
     return str_.lstrip()
-            aaronmk
+##### General SQL code objects
             aaronmk
-            aaronmk
+class MockDb:
-            aaronmk
+    def esc_value(self, value): return strings.repr_no_u(value)
             aaronmk
-            aaronmk
+    def esc_name(self, name): return esc_name(name)
             aaronmk
     def col_info(self, col):
         return TypedCol(col.name, '<type>', CustomCode('<default>'), True)
-            aaronmk
+mockDb = MockDb()
-            aaronmk
+class BasicObject(objects.BasicObject):
     def __str__(self): return clean_name(strings.repr_no_u(self))
-            aaronmk
+##### Unparameterized code objects
-            aaronmk
+class Code(BasicObject):
-            aaronmk
+    def __init__(self, lang='sql'):
         self.lang = lang
             aaronmk
-            aaronmk
+    def to_str(self, db): raise NotImplementedError()
             aaronmk
-            aaronmk
+    def __repr__(self): return self.to_str(mockDb)
             aaronmk
-            aaronmk
+class CustomCode(Code):
-            aaronmk
+    def __init__(self, str_):
         Code.__init__(self)
         self.str_ = str_
             aaronmk
     def to_str(self, db): return self.str_
-            aaronmk
+def as_Code(value, db=None):
     '''
     @param db If set, runs db.std_code() on the value.
     '''
-            aaronmk
+    if isinstance(value, Code): return value
-            aaronmk
+    if util.is_str(value):
         if db != None: value = db.std_code(value)
         return CustomCode(value)
-            aaronmk
+    else: return Literal(value)
-            aaronmk
+class Expr(Code):
-            aaronmk
+    def __init__(self, expr):
         Code.__init__(self)
         self.expr = expr
             aaronmk
     def to_str(self, db): return '('+self.expr.to_str(db)+')'
-            aaronmk
+##### Names
 class Name(Code):
-            aaronmk
+    def __init__(self, name):
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        name = truncate(name)
         self.name = name
             aaronmk
     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)
-            aaronmk
+##### Literal values
-            aaronmk
+class Literal(Code):
-            aaronmk
+    def __init__(self, value):
         Code.__init__(self)
         self.value = value
             aaronmk
     def to_str(self, db): return db.esc_value(self.value)
             aaronmk
-            aaronmk
+def as_Value(value):
     if isinstance(value, Code): return value
     else: return Literal(value)
-            aaronmk
+def is_literal(value): return isinstance(value, Literal)
             aaronmk
-            aaronmk
+def is_null(value): return is_literal(value) and value.value == None
-            aaronmk
+##### Derived elements
 src_self = object() # tells Col that it is its own source column
 class Derived(Code):
     def __init__(self, srcs):
-            aaronmk
+        '''An element which was derived from some other element(s).
-            aaronmk
+        @param srcs See self.set_srcs()
         '''
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        self.set_srcs(srcs)
-            aaronmk
+    def set_srcs(self, srcs, overwrite=True):
-            aaronmk
+        '''
         @param srcs (self_type...)|src_self The element(s) this is derived from
         '''
-            aaronmk
+        if not overwrite and self.srcs != (): return # already set
-            aaronmk
+        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)))
-            aaronmk
+##### Tables
-            aaronmk
+class Table(Derived):
-            aaronmk
+    def __init__(self, name, schema=None, srcs=(), is_temp=False):
-            aaronmk
+        '''
         @param schema str|None (for no schema)
-            aaronmk
+        @param srcs (Table...)|src_self See Derived.set_srcs()
-            aaronmk
+        '''
-            aaronmk
+        Derived.__init__(self, srcs)
-            aaronmk
+        if util.is_str(name): name = truncate(name)
             aaronmk
-            aaronmk
+        self.name = name
         self.schema = schema
-            aaronmk
+        self.is_temp = is_temp
-            aaronmk
+        self.index_cols = {}
             aaronmk
-            aaronmk
+    def to_str(self, db):
         str_ = ''
-            aaronmk
+        if self.schema != None: str_ += as_Name(self.schema).to_str(db)+'.'
         str_ += as_Name(self.name).to_str(db)
-            aaronmk
+        return str_
             aaronmk
     def to_Table(self): return self
             aaronmk
     def _compare_on(self):
         compare_on = Derived._compare_on(self)
         del compare_on['index_cols'] # ignore
         return compare_on
             aaronmk
-            aaronmk
+def is_underlying_table(table):
     return isinstance(table, Table) and table.to_Table() is table
             aaronmk
-            aaronmk
+class NoUnderlyingTableException(Exception): pass
 def underlying_table(table):
     table = remove_table_rename(table)
     if not is_underlying_table(table): raise NoUnderlyingTableException
     return table
-            aaronmk
+def as_Table(table, schema=None):
-            aaronmk
+    if table == None or isinstance(table, Code): return table
-            aaronmk
+    else: return Table(table, schema)
             aaronmk
-            aaronmk
+def suffixed_table(table, suffix):
-            aaronmk
+    table = copy.copy(table) # don't modify input!
     table.name = concat(table.name, suffix)
     return table
             aaronmk
-            aaronmk
+class NamedTable(Table):
     def __init__(self, name, code, cols=None):
         Table.__init__(self, name)
-            aaronmk
+        code = as_Table(code)
-            aaronmk
+        if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)
-            aaronmk
+        if cols != None: cols = [to_name_only_col(c).to_Col() for c in cols]
             aaronmk
         self.code = code
         self.cols = cols
     def to_str(self, db):
-            aaronmk
+        str_ = self.code.to_str(db)
         if str_.find('\n') >= 0: whitespace = '\n'
         else: whitespace = ' '
         str_ += whitespace+'AS '+Table.to_str(self, db)
-            aaronmk
+        if self.cols != None:
             str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'
-            aaronmk
+        return str_
     def to_Table(self): return Table(self.name)
-            aaronmk
+def remove_table_rename(table):
     if isinstance(table, NamedTable): table = table.code
     return table
-            aaronmk
+##### Columns
-            aaronmk
+class Col(Derived):
-            aaronmk
+    def __init__(self, name, table=None, srcs=()):
-            aaronmk
+        '''
         @param table Table|None (for no table)
-            aaronmk
+        @param srcs (Col...)|src_self See Derived.set_srcs()
-            aaronmk
+        '''
-            aaronmk
+        Derived.__init__(self, srcs)
-            aaronmk
+        if util.is_str(name): name = truncate(name)
-            aaronmk
+        if util.is_str(table): table = Table(table)
-            aaronmk
+        assert table == None or isinstance(table, Table)
         self.name = name
         self.table = table
-            aaronmk
+    def to_str(self, db, for_str=False):
-            aaronmk
+        str_ = as_Name(self.name).to_str(db)
-            aaronmk
+        if for_str: str_ = clean_name(str_)
-            aaronmk
+        if self.table != None:
-            aaronmk
+            table = self.table.to_Table()
             if for_str: str_ = concat(str(table), '.'+str_)
             else: str_ = table.to_str(db)+'.'+str_
-            aaronmk
+        return str_
             aaronmk
-            aaronmk
+    def __str__(self): return self.to_str(mockDb, for_str=True)
             aaronmk
-            aaronmk
+    def to_Col(self): return self
             aaronmk
-            aaronmk
+def is_col(col): return isinstance(col, Col)
             aaronmk
-            aaronmk
+def is_table_col(col): return is_col(col) and col.table != None
-            aaronmk
+def index_col(col):
     if not is_table_col(col): return None
             aaronmk
     table = col.table
     try: name = table.index_cols[col.name]
     except KeyError: return None
     else: return Col(name, table, col.srcs)
             aaronmk
-            aaronmk
+def is_temp_col(col): return is_table_col(col) and col.table.is_temp
             aaronmk
-            aaronmk
+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)
             aaronmk
     if isinstance(col, Code): return col
-            aaronmk
+    elif util.is_str(col): return Col(col, table)
     else: return Literal(col)
             aaronmk
-            aaronmk
+def with_table(col, table):
-            aaronmk
+    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
-            aaronmk
+    elif isinstance(col, Col):
-            aaronmk
+        col = copy.copy(col) # don't modify input!
         col.table = table
-            aaronmk
+    return col
-            aaronmk
+def with_default_table(col, table):
-            aaronmk
+    col = as_Col(col)
-            aaronmk
+    if col.table == None: col = with_table(col, table)
-            aaronmk
+    return col
-            aaronmk
+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
-            aaronmk
+def to_name_only_col(col, check_table=None):
     col = as_Col(col)
-            aaronmk
+    if not is_table_col(col): return col
             aaronmk
     if check_table != None:
         table = col.table
         assert table == None or table == check_table
     return Col(col.name)
-            aaronmk
+def suffixed_col(col, suffix):
     return Col(concat(col.name, suffix), col.table, col.srcs)
-            aaronmk
+def has_srcs(col): return is_col(col) and col.srcs
-            aaronmk
+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]
     return [Col(','.join(s)) for s in itertools.product(*srcs)]
             aaronmk
-            aaronmk
+class NamedCol(Col):
-            aaronmk
+    def __init__(self, name, code):
-            aaronmk
+        Col.__init__(self, name)
-            aaronmk
+        code = as_Value(code)
             aaronmk
         self.code = code
     def to_str(self, db):
-            aaronmk
+        return self.code.to_str(db)+' AS '+Col.to_str(self, db)
             aaronmk
     def to_Col(self): return Col(self.name)
             aaronmk
-            aaronmk
+def remove_col_rename(col):
     if isinstance(col, NamedCol): col = col.code
     return col
-            aaronmk
+def underlying_col(col):
     col = remove_col_rename(col)
-            aaronmk
+    if not isinstance(col, Col): raise NoUnderlyingTableException
-            aaronmk
+    return Col(col.name, underlying_table(col.table), col.srcs)
             aaronmk
-            aaronmk
+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)
-            aaronmk
+class ColDict(dicts.DictProxy):
-            aaronmk
+    '''A dict that automatically makes inserted entries Col objects'''
-            aaronmk
+    def __init__(self, db, keys_table, dict_={}):
-            aaronmk
+        dicts.DictProxy.__init__(self, OrderedDict())
             aaronmk
-            aaronmk
+        keys_table = as_Table(keys_table)
-            aaronmk
+        self.db = db
-            aaronmk
+        self.table = keys_table
-            aaronmk
+        self.update(dict_) # after setting vars because __setitem__() needs them
             aaronmk
-            aaronmk
+    def copy(self): return ColDict(self.db, self.table, self.inner.copy())
             aaronmk
-            aaronmk
+    def __getitem__(self, key):
         return dicts.DictProxy.__getitem__(self, self._key(key))
             aaronmk
-            aaronmk
+    def __setitem__(self, key, value):
-            aaronmk
+        key = self._key(key)
-            aaronmk
+        if value == None: value = self.db.col_info(key).default
-            aaronmk
+        dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))
             aaronmk
-            aaronmk
+    def _key(self, key): return as_Col(key, self.table)
             aaronmk
-            aaronmk
+##### Composite types
 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)))+')'
-            aaronmk
+class Tuple(List):
     def to_str(self, db): return 'ROW'+List.to_str(self, db)
-            aaronmk
+#### 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):
-            aaronmk
+        str_ = Col.to_str(self, db)+' '+as_Code(self.type).to_str(db)
-            aaronmk
+        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)
-            aaronmk
+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)
-            aaronmk
+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'
-            aaronmk
+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'
-            aaronmk
+##### Functions
-            aaronmk
+Function = Table
-            aaronmk
+as_Function = as_Table
-            aaronmk
+class InternalFunction(CustomCode): pass
-            aaronmk
+#### Calls
-            aaronmk
+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)
-            aaronmk
+class FunctionCall(Code):
-            aaronmk
+    def __init__(self, function, *args, **kw_args):
-            aaronmk
+        '''
-            aaronmk
+        @param args [Code|literal-value...] The function's arguments
-            aaronmk
+        '''
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        function = as_Function(function)
-            aaronmk
+        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()]
             aaronmk
         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+')'
-            aaronmk
+def wrap_in_func(function, value):
     '''Wraps a value inside a function call.
     Propagates any column renaming to the returned value.
     '''
-            aaronmk
+    return wrap(lambda v: FunctionCall(function, v), value)
             aaronmk
-            aaronmk
+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]
-            aaronmk
+#### Definitions
 class FunctionDef(Code):
-            aaronmk
+    def __init__(self, function, return_type, body, params=[], modifiers=None):
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        return_type = as_Code(return_type)
-            aaronmk
+        body = as_Code(body)
-            aaronmk
+        self.function = function
         self.return_type = return_type
         self.body = body
-            aaronmk
+        self.params = params
-            aaronmk
+        self.modifiers = modifiers
             aaronmk
     def to_str(self, db):
-            aaronmk
+        params_str = (', '.join((p.to_str(db) for p in self.params)))
-            aaronmk
+        str_ = '''\
-            aaronmk
+CREATE FUNCTION '''+self.function.to_str(db)+'''('''+params_str+''')
-            aaronmk
+RETURNS '''+self.return_type.to_str(db)+'''
-            aaronmk
+LANGUAGE '''+self.body.lang+'''
-            aaronmk
+'''
         if self.modifiers != None: str_ += self.modifiers+'\n'
         str_ += '''\
-            aaronmk
+AS $$
-            aaronmk
+'''+self.body.to_str(db)+'''
-            aaronmk
+$$;
 '''
         return str_
-            aaronmk
+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)
-            aaronmk
+### PL/pgSQL
-            aaronmk
+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'
-            aaronmk
+## Exceptions
-            aaronmk
+class BaseExcHandler(BasicObject):
     def to_str(self, db, body): raise NotImplementedError()
             aaronmk
     def __repr__(self): return self.to_str(mockDb, '<body>')
             aaronmk
 class ExcHandler(BaseExcHandler):
-            aaronmk
+    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)
-            aaronmk
+        if self.handler != None:
             handler_str = '\n'+strings.indent(self.handler.to_str(db), 2)
-            aaronmk
+        else: handler_str = ' NULL;\n'
             aaronmk
         str_ = '''\
 BEGIN
-            aaronmk
+'''+strings.indent(body.to_str(db))+'''\
-            aaronmk
+EXCEPTION
-            aaronmk
+    WHEN '''+self.exc+''' THEN'''+handler_str+'''\
-            aaronmk
+END;\
 '''
         return str_
-            aaronmk
+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
-            aaronmk
+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)
-            aaronmk
+unique_violation_handler = ExcHandler('unique_violation')
-            aaronmk
+plpythonu_error_handler = ExcHandler('internal_error', '''\
 RAISE data_exception USING MESSAGE =
     regexp_replace(SQLERRM, E'^PL/Python: \\w+: ', '');
 ''')
-            aaronmk
+def data_exception_handler(handler):
     return ExcHandler('data_exception', handler)
-            aaronmk
+class RowExcIgnore(Code):
     def __init__(self, row_type, select_query, with_row, cols=None,
-            aaronmk
+        exc_handler=unique_violation_handler, row_var='row'):
-            aaronmk
+        Code.__init__(self, lang='plpgsql')
-            aaronmk
+        row_type = as_Code(row_type)
-            aaronmk
+        select_query = as_Code(select_query)
         with_row = as_Code(with_row)
-            aaronmk
+        row_var = as_Table(row_var)
             aaronmk
         self.row_type = row_type
         self.select_query = select_query
         self.with_row = with_row
         self.cols = cols
-            aaronmk
+        self.exc_handler = exc_handler
-            aaronmk
+        self.row_var = row_var
             aaronmk
     def to_str(self, db):
-            aaronmk
+        if self.cols == None: row_vars = [self.row_var]
         else: row_vars = [Col(c.name, self.row_var) for c in self.cols]
             aaronmk
         str_ = '''\
 DECLARE
-            aaronmk
+    '''+self.row_var.to_str(db)+''' '''+self.row_type.to_str(db)+''';
-            aaronmk
+BEGIN
-            aaronmk
+    /* 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."
-            aaronmk
+    (http://www.postgresql.org/docs/8.3/static/plpgsql-control-structures.html\
 #PLPGSQL-ERROR-TRAPPING)
     */
     FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN
-            aaronmk
+'''+strings.indent(self.select_query.to_str(db), 2)+'''\
-            aaronmk
+    LOOP
-            aaronmk
+'''+strings.indent(self.exc_handler.to_str(db, self.with_row), 2)+'''\
-            aaronmk
+    END LOOP;
 END;\
 '''
         return str_
-            aaronmk
+##### 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_+')'
-            aaronmk
+def cast_literal(value):
-            aaronmk
+    if not is_literal(value): return value
             aaronmk
     if util.is_str(value.value): value = Cast('text', value)
     return value
-            aaronmk
+##### Conditions
             aaronmk
-            aaronmk
+class NotCond(Code):
     def __init__(self, cond):
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        self.cond = cond
     def to_str(self, db): return 'NOT '+self.cond.to_str(db)
-            aaronmk
+class ColValueCond(Code):
     def __init__(self, col, value):
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        value = as_ValueCond(value)
         self.col = col
         self.value = value
     def to_str(self, db): return self.value.to_str(db, self.col)
-            aaronmk
+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_
-            aaronmk
+##### Condition column comparisons
-            aaronmk
+class ValueCond(BasicObject):
-            aaronmk
+    def __init__(self, value):
-            aaronmk
+        value = remove_col_rename(as_Value(value))
             aaronmk
         self.value = value
             aaronmk
-            aaronmk
+    def to_str(self, db, left_value):
-            aaronmk
+        '''
-            aaronmk
+        @param left_value The Code object that the condition is being applied on
-            aaronmk
+        '''
         raise NotImplemented()
             aaronmk
-            aaronmk
+    def __repr__(self): return self.to_str(mockDb, '<left_value>')
             aaronmk
 class CompareCond(ValueCond):
     def __init__(self, value, operator='='):
-            aaronmk
+        '''
         @param operator By default, compares NULL values literally. Use '~=' or
             '~!=' to pass NULLs through.
         '''
-            aaronmk
+        ValueCond.__init__(self, value)
         self.operator = operator
-            aaronmk
+    def to_str(self, db, left_value):
-            aaronmk
+        left_value = remove_col_rename(as_Col(left_value))
             aaronmk
-            aaronmk
+        right_value = self.value
         # Parse operator
-            aaronmk
+        operator = self.operator
-            aaronmk
+        passthru_null_ref = [False]
         operator = strings.remove_prefix('~', operator, passthru_null_ref)
         neg_ref = [False]
         operator = strings.remove_prefix('!', operator, neg_ref)
-            aaronmk
+        equals = operator.endswith('=') # also includes <=, >=
             aaronmk
-            aaronmk
+        # Handle nullable columns
         check_null = False
-            aaronmk
+        if not passthru_null_ref[0]: # NULLs compare equal
-            aaronmk
+            try: left_value = ensure_not_null(db, left_value)
-            aaronmk
+            except ensure_not_null_excs: # fall back to alternate method
                 check_null = equals and isinstance(right_value, Col)
-            aaronmk
+            else:
-            aaronmk
+                if isinstance(left_value, EnsureNotNull):
                     right_value = ensure_not_null(db, right_value,
                         left_value.type) # apply same function to both sides
             aaronmk
-            aaronmk
+        if equals and is_null(right_value): operator = 'IS'
-            aaronmk
+        left = left_value.to_str(db)
         right = right_value.to_str(db)
-            aaronmk
+        # Create str
         str_ = left+' '+operator+' '+right
-            aaronmk
+        if check_null:
-            aaronmk
+            str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'
         if neg_ref[0]: str_ = 'NOT '+str_
-            aaronmk
+        return str_
             aaronmk
-            aaronmk
+# 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:
-            aaronmk
+            value = with_default_table(value, default_table)
-            aaronmk
+        return CompareCond(value)
-            aaronmk
+    else: return value
             aaronmk
-            aaronmk
+##### Joins
-            aaronmk
+join_same = object() # tells Join the left and right columns have the same name
             aaronmk
-            aaronmk
+# Tells Join the left and right columns have the same name and are never NULL
 join_same_not_null = object()
-            aaronmk
+filter_out = object() # tells Join to filter out rows that match the join
-            aaronmk
+class Join(BasicObject):
-            aaronmk
+    def __init__(self, table, mapping={}, type_=None):
-            aaronmk
+        '''
         @param mapping dict(right_table_col=left_table_col, ...)
-            aaronmk
+            * if left_table_col is join_same: left_table_col = right_table_col
-            aaronmk
+              * 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
-            aaronmk
+        @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_
-            aaronmk
+    def to_str(self, db, left_table_):
-            aaronmk
+        def join(entry):
             '''Parses non-USING joins'''
             right_table_col, left_table_col = entry
-            aaronmk
+            # Switch order (right_table_col is on the left in the comparison)
             left = right_table_col
             right = left_table_col
-            aaronmk
+            left_table = self.table
             right_table = left_table_
             aaronmk
-            aaronmk
+            # Parse left side
-            aaronmk
+            left = with_default_table(left, left_table)
             aaronmk
-            aaronmk
+            # Parse special values
-            aaronmk
+            left_on_right = Col(left.name, right_table)
             if right is join_same: right = left_on_right
-            aaronmk
+            elif right is join_same_not_null:
-            aaronmk
+                right = CompareCond(left_on_right, '~=')
             aaronmk
-            aaronmk
+            # Parse right side
-            aaronmk
+            right = as_ValueCond(right, right_table)
             aaronmk
             return right.to_str(db, left)
             aaronmk
-            aaronmk
+        # Create join condition
         type_ = self.type_
-            aaronmk
+        joins = self.mapping
-            aaronmk
+        if joins == {}: join_cond = None
         elif type_ is not filter_out and reduce(operator.and_,
-            aaronmk
+            (v is join_same_not_null for v in joins.itervalues())):
-            aaronmk
+            # all cols w/ USING, so can use simpler USING syntax
-            aaronmk
+            cols = map(to_name_only_col, joins.iterkeys())
             join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'
-            aaronmk
+        else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')
             aaronmk
-            aaronmk
+        if isinstance(self.table, NamedTable): whitespace = '\n'
         else: whitespace = ' '
-            aaronmk
+        # Create join
         if type_ is filter_out: type_ = 'LEFT'
-            aaronmk
+        str_ = ''
         if type_ != None: str_ += type_+' '
-            aaronmk
+        str_ += 'JOIN'+whitespace+self.table.to_str(db)
         if join_cond != None: str_ += whitespace+join_cond
-            aaronmk
+        return str_
             aaronmk
-            aaronmk
+    def __repr__(self): return self.to_str(mockDb, '<left_table>')
             aaronmk
 ##### Value exprs
-            aaronmk
+all_cols = CustomCode('*')
-            aaronmk
+default = CustomCode('DEFAULT')
-            aaronmk
+row_count = FunctionCall(InternalFunction('COUNT'), all_cols)
             aaronmk
-            aaronmk
+class Coalesce(FunctionCall):
     def __init__(self, *args):
         FunctionCall.__init__(self, InternalFunction('COALESCE'), *args)
             aaronmk
-            aaronmk
+class Nullif(FunctionCall):
     def __init__(self, *args):
         FunctionCall.__init__(self, InternalFunction('NULLIF'), *args)
-            aaronmk
+# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>
-            aaronmk
+null_sentinels = {
     'character varying': r'\N',
     'double precision': 'NaN',
     'integer': 2147483647,
     'text': r'\N',
     'timestamp with time zone': 'infinity'
+}
             aaronmk
-            aaronmk
+class EnsureNotNull(Coalesce):
-            aaronmk
+    def __init__(self, value, type_):
-            aaronmk
+        Coalesce.__init__(self, as_Col(value),
-            aaronmk
+            Cast(type_, null_sentinels[type_]))
             aaronmk
         self.type = type_
             aaronmk
     def to_str(self, db):
         col = self.args[0]
         index_col_ = index_col(col)
         if index_col_ != None: return index_col_.to_str(db)
-            aaronmk
+        return Coalesce.to_str(self, db)
             aaronmk
-            aaronmk
+##### Table exprs
 class Values(Code):
     def __init__(self, values):
-            aaronmk
+        '''
         @param values [...]|[[...], ...] Can be one or multiple rows.
         '''
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        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))
             aaronmk
-            aaronmk
+        self.rows = rows
             aaronmk
     def to_str(self, db):
-            aaronmk
+        return 'VALUES '+(', '.join((List(*r).to_str(db) for r in self.rows)))
             aaronmk
-            aaronmk
+def NamedValues(name, cols, values):
-            aaronmk
+    '''
-            aaronmk
+    @param cols None|[...]
-            aaronmk
+    @post `cols` will be changed to Col objects with the table set to `name`.
     '''
-            aaronmk
+    table = NamedTable(name, Values(values), cols)
-            aaronmk
+    if cols != None: set_cols_table(table, cols)
-            aaronmk
+    return table
             aaronmk
-            aaronmk
+##### Database structure
-            aaronmk
+ensure_not_null_excs = (NoUnderlyingTableException, KeyError)
-            aaronmk
+def ensure_not_null(db, col, type_=None):
-            aaronmk
+    '''
-            aaronmk
+    @param col If type_ is not set, must have an underlying column.
-            aaronmk
+    @param type_ If set, overrides the underlying column's type.
-            aaronmk
+    @return EnsureNotNull|Col
     @throws ensure_not_null_excs
     '''
-            aaronmk
+    nullable = True
     try: typed_col = db.col_info(underlying_col(col))
     except NoUnderlyingTableException:
-            aaronmk
+        col = remove_col_rename(col)
-            aaronmk
+        if is_literal(col) and not is_null(col): nullable = False
-            aaronmk
+        elif type_ == None: raise
-            aaronmk
+    else:
         if type_ == None: type_ = typed_col.type
         nullable = typed_col.nullable
-            aaronmk
+    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
-            aaronmk
+    return col

Project

General

Profile

root/lib/sql_gen.py @ 3518