1from autotest_lib.frontend.shared import exceptions 2 3class ConstraintError(Exception): 4 """Raised when an error occurs applying a Constraint.""" 5 6 7class QueryProcessor(object): 8 def __init__(self): 9 # maps selector name to (selector, constraint) 10 self._selectors = {} 11 self._alias_counter = 0 12 13 14 def add_field_selector(self, name, field=None, value_transform=None, 15 doc=None): 16 if not field: 17 field = name 18 self.add_selector(Selector(name, doc), 19 _FieldConstraint(field, value_transform)) 20 21 22 def add_related_existence_selector(self, name, model, field, doc=None): 23 self.add_selector( 24 Selector(name, doc), 25 _RelatedExistenceConstraint(model, field, 26 make_alias_fn=self.make_alias)) 27 28 29 def add_keyval_selector(self, name, model, key_field, value_field, 30 doc=None): 31 self.add_selector( 32 Selector(name, doc), 33 _KeyvalConstraint(model, key_field, value_field, 34 make_alias_fn=self.make_alias)) 35 36 37 def add_selector(self, selector, constraint): 38 if self._selectors is None: 39 self._selectors = {} 40 self._selectors[selector.name] = (selector, constraint) 41 42 43 def make_alias(self): 44 self._alias_counter += 1 45 return 'alias%s' % self._alias_counter 46 47 48 def selectors(self): 49 return tuple(selector for selector, constraint 50 in self._selectors.itervalues()) 51 52 53 def has_selector(self, selector_name): 54 return selector_name in self._selectors 55 56 57 def apply_selector(self, queryset, selector_name, value, 58 comparison_type=None, is_inverse=False): 59 if comparison_type is None: 60 comparison_type = 'equals' 61 _, constraint = self._selectors[selector_name] 62 try: 63 return constraint.apply_constraint(queryset, value, comparison_type, 64 is_inverse) 65 except ConstraintError, exc: 66 raise exceptions.BadRequest('Selector %s: %s' 67 % (selector_name, exc)) 68 69 70 # common value conversions 71 72 def read_boolean(self, boolean_input): 73 if boolean_input.lower() == 'true': 74 return True 75 if boolean_input.lower() == 'false': 76 return False 77 raise exceptions.BadRequest('Invalid input for boolean: %r' 78 % boolean_input) 79 80 81class Selector(object): 82 def __init__(self, name, doc): 83 self.name = name 84 self.doc = doc 85 86 87class Constraint(object): 88 def apply_constraint(self, queryset, value, comparison_type, is_inverse): 89 raise NotImplementedError 90 91 92class _FieldConstraint(Constraint): 93 def __init__(self, field, value_transform=None): 94 self._field = field 95 self._value_transform = value_transform 96 97 98 _COMPARISON_MAP = { 99 'equals': 'exact', 100 'lt': 'lt', 101 'le': 'lte', 102 'gt': 'gt', 103 'ge': 'gte', 104 'contains': 'contains', 105 'startswith': 'startswith', 106 'endswith': 'endswith', 107 'in': 'in', 108 } 109 110 111 def apply_constraint(self, queryset, value, comparison_type, is_inverse): 112 if self._value_transform: 113 value = self._value_transform(value) 114 115 kwarg_name = str(self._field + '__' + 116 self._COMPARISON_MAP[comparison_type]) 117 if comparison_type == 'in': 118 value = value.split(',') 119 120 if is_inverse: 121 return queryset.exclude(**{kwarg_name: value}) 122 else: 123 return queryset.filter(**{kwarg_name: value}) 124 125 126class _RelatedExistenceConstraint(Constraint): 127 def __init__(self, model, field, make_alias_fn): 128 self._model = model 129 self._field = field 130 self._make_alias_fn = make_alias_fn 131 132 133 def apply_constraint(self, queryset, value, comparison_type, is_inverse): 134 if comparison_type not in (None, 'equals'): 135 raise ConstraintError('Can only use equals or not equals with ' 136 'this selector') 137 related_query = self._model.objects.filter(**{self._field: value}) 138 if not related_query: 139 raise ConstraintError('%s %s not found' % (self._model_name, value)) 140 alias = self._make_alias_fn() 141 queryset = queryset.model.objects.join_custom_field(queryset, 142 related_query, 143 alias) 144 if is_inverse: 145 condition = '%s.%s IS NULL' 146 else: 147 condition = '%s.%s IS NOT NULL' 148 condition %= (alias, 149 queryset.model.objects.key_on_joined_table(related_query)) 150 151 queryset = queryset.model.objects.add_where(queryset, condition) 152 153 return queryset 154 155 156class _KeyvalConstraint(Constraint): 157 def __init__(self, model, key_field, value_field, make_alias_fn): 158 self._model = model 159 self._key_field = key_field 160 self._value_field = value_field 161 self._make_alias_fn = make_alias_fn 162 163 164 def apply_constraint(self, queryset, value, comparison_type, is_inverse): 165 if comparison_type not in (None, 'equals'): 166 raise ConstraintError('Can only use equals or not equals with ' 167 'this selector') 168 if '=' not in value: 169 raise ConstraintError('You must specify a key=value pair for this ' 170 'selector') 171 172 key, actual_value = value.split('=', 1) 173 related_query = self._model.objects.filter( 174 **{self._key_field: key, self._value_field: actual_value}) 175 alias = self._make_alias_fn() 176 queryset = queryset.model.objects.join_custom_field(queryset, 177 related_query, 178 alias) 179 if is_inverse: 180 condition = '%s.%s IS NULL' 181 else: 182 condition = '%s.%s IS NOT NULL' 183 condition %= (alias, 184 queryset.model.objects.key_on_joined_table(related_query)) 185 186 queryset = queryset.model.objects.add_where(queryset, condition) 187 188 return queryset 189