Project

General

Profile

1 2211 aaronmk
# SQL code generation
2
3 2276 aaronmk
import operator
4
5 2211 aaronmk
import sql
6 2222 aaronmk
import strings
7 2227 aaronmk
import util
8 2211 aaronmk
9 2219 aaronmk
##### SQL code objects
10
11 2211 aaronmk
class Code:
12
    def to_str(self, db): raise NotImplemented()
13 2228 aaronmk
14
    def __str__(self): return str(self.__dict__)
15 2211 aaronmk
16 2269 aaronmk
class CustomCode(Code):
17 2256 aaronmk
    def __init__(self, str_): self.str_ = str_
18
19
    def to_str(self, db): return self.str_
20
21 2216 aaronmk
class Literal(Code):
22 2211 aaronmk
    def __init__(self, value): self.value = value
23 2213 aaronmk
24
    def to_str(self, db): return db.esc_value(self.value)
25 2211 aaronmk
26 2216 aaronmk
def is_null(value): return isinstance(value, Literal) and value.value == None
27
28 2211 aaronmk
class Table(Code):
29
    def __init__(self, name, schema=None):
30
        '''
31
        @param schema str|None (for no schema)
32
        '''
33
        self.name = name
34
        self.schema = schema
35
36
    def to_str(self, db): return sql.qual_name(db, self.schema, self.name)
37
38 2219 aaronmk
def as_Table(table):
39 2270 aaronmk
    if table == None or isinstance(table, Code): return table
40 2219 aaronmk
    elif isinstance(table, tuple):
41
        schema, table = table
42
        return Table(table, schema)
43
    else: return Table(table)
44
45 2211 aaronmk
class Col(Code):
46
    def __init__(self, name, table=None):
47
        '''
48
        @param table Table|None (for no table)
49
        '''
50 2241 aaronmk
        if util.is_str(table): table = Table(table)
51 2211 aaronmk
        assert table == None or isinstance(table, Table)
52
53
        self.name = name
54
        self.table = table
55
56
    def to_str(self, db):
57
        str_ = ''
58
        if self.table != None: str_ += self.table.to_str(db)+'.'
59
        str_ += sql.esc_name(db, self.name)
60
        return str_
61
62 2260 aaronmk
def as_Col(col, table=None):
63
    if col == None or isinstance(col, Code): return col
64
    else: return Col(col, table)
65
66 2229 aaronmk
class NamedCode(Code):
67
    def __init__(self, name, code):
68
        if not isinstance(code, Code): code = Literal(code)
69
70
        self.name = name
71
        self.code = code
72
73
    def to_str(self, db):
74
        return self.code.to_str(db)+' AS '+sql.esc_name(db, self.name)
75
76 2259 aaronmk
##### Parameterized SQL code objects
77
78 2214 aaronmk
class ValueCond:
79 2213 aaronmk
    def __init__(self, value):
80 2225 aaronmk
        if not isinstance(value, Code): value = Literal(value)
81 2213 aaronmk
82
        self.value = value
83 2214 aaronmk
84 2216 aaronmk
    def to_str(self, db, left_value):
85 2214 aaronmk
        '''
86 2216 aaronmk
        @param left_value The Code object that the condition is being applied on
87 2214 aaronmk
        '''
88
        raise NotImplemented()
89 2228 aaronmk
90
    def __str__(self): return str(self.__dict__)
91 2211 aaronmk
92
class CompareCond(ValueCond):
93
    def __init__(self, value, operator='='):
94 2222 aaronmk
        '''
95
        @param operator By default, compares NULL values literally. Use '~=' or
96
            '~!=' to pass NULLs through.
97
        '''
98 2211 aaronmk
        ValueCond.__init__(self, value)
99
        self.operator = operator
100
101 2216 aaronmk
    def to_str(self, db, left_value):
102
        if not isinstance(left_value, Code): left_value = Col(left_value)
103
104 2222 aaronmk
        right_value = self.value
105
        left = left_value.to_str(db)
106
        right = right_value.to_str(db)
107
108
        # Parse operator
109 2216 aaronmk
        operator = self.operator
110 2222 aaronmk
        passthru_null_ref = [False]
111
        operator = strings.remove_prefix('~', operator, passthru_null_ref)
112
        neg_ref = [False]
113
        operator = strings.remove_prefix('!', operator, neg_ref)
114
        equals = operator.endswith('=')
115
        if equals and is_null(self.value): operator = 'IS'
116
117
        # Create str
118
        str_ = left+' '+operator+' '+right
119
        if equals and not passthru_null_ref[0] and isinstance(right_value, Col):
120
            str_ += ' OR ('+left+' IS NULL AND '+right+' IS NULL)'
121
        if neg_ref[0]: str_ = 'NOT ('+str_+')'
122
        return str_
123 2216 aaronmk
124 2260 aaronmk
# Tells as_ValueCond() to assume a non-ValueCond is a literal value
125
assume_literal = object()
126
127
def as_ValueCond(value, default_table=assume_literal):
128
    if not isinstance(value, ValueCond):
129
        if default_table is not assume_literal:
130
            value = as_Col(value, default_table)
131
        return CompareCond(value)
132 2216 aaronmk
    else: return value
133 2219 aaronmk
134 2260 aaronmk
join_using = object() # tells Join to join the column with USING
135
136
filter_out = object() # tells Join to filter out rows that match the join
137
138
class Join(Code):
139
    def __init__(self, table, mapping, type_=None):
140
        '''
141
        @param mapping dict(right_table_col=left_table_col, ...)
142
            * if left_table_col is join_using: left_table_col = right_table_col
143
        @param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out
144
            * filter_out: equivalent to 'LEFT' with the query filtered by
145
              `table_pkey IS NULL` (indicating no match)
146
        '''
147
        if util.is_str(table): table = Table(table)
148
        assert type_ == None or util.is_str(type_) or type_ is filter_out
149
150
        self.table = table
151
        self.mapping = mapping
152
        self.type_ = type_
153
154
    def to_str(self, db, left_table):
155
        def join(entry):
156
            '''Parses non-USING joins'''
157
            right_table_col, left_table_col = entry
158
159
            # Parse special values
160
            if left_table_col is join_using: left_table_col = right_table_col
161
162
            cond = as_ValueCond(right_table_col, self.table)
163
            return cond.to_str(db, as_Col(left_table_col, left_table))
164
165 2265 aaronmk
        # Create join condition
166
        type_ = self.type_
167 2276 aaronmk
        joins = self.mapping
168 2265 aaronmk
        if type_ is not filter_out and reduce(operator.and_,
169
            (v is join_using for v in joins.itervalues())):
170 2260 aaronmk
            # all cols w/ USING, so can use simpler USING syntax
171
            join_cond = 'USING ('+(', '.join(joins.iterkeys()))+')'
172
        else: join_cond = 'ON '+(' AND '.join(map(join, joins.iteritems())))
173
174
        # Create join
175
        if type_ is filter_out: type_ = 'LEFT'
176 2266 aaronmk
        str_ = ''
177
        if type_ != None: str_ += type_+' '
178 2276 aaronmk
        str_ += 'JOIN '+self.table.to_str(db)+' '+join_cond
179 2266 aaronmk
        return str_
180 2260 aaronmk
181 2219 aaronmk
##### Old-style format support
182
183
def unescape_table(table):
184
    '''Currently only works with PostgreSQL.'''
185
    if table == None: return table
186
187
    assert table.count('.') <= 1
188
    parts = tuple((v.replace('"', '') for v in table.split('"."', 2)))
189
    if len(parts) == 1: parts, = parts
190
    return parts
191
192 2262 aaronmk
def table2sql_gen(table, table_is_esc=False):
193
    '''Converts old-style (tuple-based) tables to sql_gen-compatible values.
194
    @param table_is_esc If False, assumes any table name is not escaped or that
195
        re-escaping it will produce the same value.
196
    '''
197
    if util.is_str(table) and table_is_esc: table = unescape_table(table)
198
    return as_Table(table)