def test_row_stream_query_set(self):
"""Build a RowStream from a Django query set."""
row_builder = DjangoRowBuilder(
[DjangoField('a'), DjangoField('a', rename='b')])
row_stream = from_django(row_builder, MockQuerySet())
self.assertEqual(['a', 'b'], row_stream.fields)
self.assertEqual([{'a': 'aa', 'b': 'aa'}], list(row_stream))
def test_django_field_tranform(self):
"""Transform the value in the Django field."""
field = DjangoField('tags', transform=extract_tags)
self.assertEqual([{
'id': 1,
'name': 'somename',
'hex_color': 'bbccdd'
}], [dict(r) for r in field.get_value(MockRecord())])
def test_row_stream_limit_fields(self):
row_builder = DjangoRowBuilder(
[DjangoField('a'),
DjangoField('name'),
DjangoField('n')])
row_stream = from_django(row_builder,
MockQuerySet(),
values=['a', 'name'])
self.assertEqual(['a', 'name'], row_stream.fields)
self.assertEqual([{'a': 'aa', 'name': 'somename'}], list(row_stream))
def test_row_builder(self):
"""Build a row from several DjangoFields."""
row_builder = DjangoRowBuilder([
DjangoField('a'),
DjangoField('a', rename='b'),
DjangoField('tags', transform=extract_tags)
])
self.assertEqual(['a', 'b', 'tags'], row_builder.get_fields())
expected_tag_list = [{
'id': 1,
'name': 'somename',
'hex_color': 'bbccdd'
}]
self.assertEqual({
'a': 'aa',
'b': 'aa',
'tags': expected_tag_list
}, row_builder.get_values(MockRecord()))
class PartnersDataSource(DataSource):
partner_row_builder = DjangoRowBuilder([
DjangoField('id', rename='partner_id'),
DjangoField('data_source'),
DjangoField('last_action_time', rename='date'),
DjangoField('name'),
DjangoField('primary_contact.name', rename='primary_contact'),
DjangoField('tags', transform=extract_tags),
DjangoField('uri')
])
def run(self, data_type, company, filter_spec, values):
return dispatch_run_by_data_type(self, data_type, company, filter_spec,
values)
def filtered_partners(self, company, filter_spec):
qs_filtered = filter_spec.filter_partners(company)
qs_distinct = qs_filtered.distinct()
return qs_distinct
def run_unaggregated(self, company, filter_spec, values):
qs_filtered = filter_spec.filter_partners(company)
qs_optimized = qs_filtered.prefetch_related('tags', 'primary_contact')
partners = from_django(self.partner_row_builder, qs_optimized, values)
return list(partners)
def run_count_comm_rec_per_month(self, company, filter_spec, values):
# Get a list of valid partners (according to our filter).
qs_filtered = filter_spec.filter_partners(company)
partner_ids = list(qs_filtered.values_list('id', flat=True))
# If there are no partners, we're done.
if not partner_ids:
return []
sql = count_comm_rec_per_month_per_partner_sql
cursor = connection.cursor()
cursor.execute(sql, {'partner_list': partner_ids})
count_stream = from_cursor(cursor)
# Iterate over counts results.
# Find year range and organize counts for later.
min_year = None
max_year = None
db_counts = {}
for record in count_stream:
partner_id = record['partner_id']
year = record['year']
month = record['month']
comm_rec_count = record['comm_rec_count']
if year < min_year or min_year is None:
min_year = year
if year > max_year or max_year is None:
max_year = year
key = (partner_id, year, month)
db_counts[key] = comm_rec_count
# If there are no communication records, assume current year.
if not db_counts:
min_year = datetime.now().year
max_year = min_year
# Create full range of time slots with counts.
# DB won't fill in zero's without a complex query.
# Do this instead.
full_counts = {}
for partner_id in partner_ids:
count_list = []
full_counts[partner_id] = count_list
for year in xrange(min_year, max_year + 1):
for month in xrange(1, 12 + 1):
key = (partner_id, year, month)
if key in db_counts:
count = db_counts[key]
count_list.append((year, month, count))
else:
count_list.append((year, month, 0))
qs_optimized = qs_filtered.prefetch_related('tags', 'primary_contact')
partners = from_django(self.partner_row_builder, qs_optimized)
# Join the two streams
joined_fields = partners.fields + ['year', 'month', 'comm_rec_count']
joined_data = []
for partner_record in partners:
count_list = full_counts.get(partner_record['partner_id'], [])
for count_tuple in count_list:
year, month, count = count_tuple
joined_record = dict(partner_record)
joined_record.update({
'year': year,
'month': month,
'comm_rec_count': count,
})
joined_data.append(joined_record)
joined = from_list(joined_fields, joined_data)
return list(joined)
def filter_type(self):
return PartnersFilter
def help(self, company, filter_spec, field, partial):
return dispatch_help_by_field_name(self, company, filter_spec, field,
partial)
def help_city(self, company, filter_spec, partial):
"""Get help for the city field."""
modified_filter_spec = filter_spec.clone_without_city()
partners_qs = modified_filter_spec.filter_partners(company)
locations_qs = (Location.objects.filter(
contacts__partner__in=partners_qs).filter(city__icontains=partial))
city_qs = locations_qs.values('city').distinct()
return [{'value': c['city'], 'display': c['city']} for c in city_qs]
def help_state(self, company, filter_spec, partial):
"""Get help for the state field."""
modified_filter_spec = filter_spec.clone_without_state()
partners_qs = modified_filter_spec.filter_partners(company)
locations_qs = (Location.objects.filter(
contacts__partner__in=partners_qs).filter(
state__icontains=partial))
state_qs = locations_qs.values('state').order_by('state').distinct()
return [{
'value': c['state'],
'display': states[c['state']]
} for c in state_qs if c['state']]
def help_tags(self, company, filter_spec, partial):
"""Get help for the tags field."""
partners_qs = PartnersFilter().filter_partners(company)
tags_qs = (Tag.objects.filter(partner__in=partners_qs).filter(
name__icontains=partial).values('name', 'hex_color').distinct())
return [{
'value': t['name'],
'display': t['name'],
'hexColor': t['hex_color'],
} for t in tags_qs]
def help_uri(self, company, filter_spec, partial):
"""Get help for the uri field."""
partners_qs = filter_spec.filter_partners(company)
uris_qs = (partners_qs.filter(
uri__icontains=partial).values('uri').distinct())
return [{'value': c['uri'], 'display': c['uri']} for c in uris_qs]
def help_data_source(self, company, filter_spec, partial):
"""Get help for the data_source field."""
partners_qs = filter_spec.filter_partners(company)
data_sources_qs = (partners_qs.filter(
data_source__icontains=partial).values('data_source').distinct())
return [{
'value': c['data_source'],
'display': c['data_source']
} for c in data_sources_qs]
def adorn_filter_items(self, company, found_filter_items):
adorned = {}
empty = PartnersFilter()
if 'tags' in found_filter_items:
adorned[u'tags'] = {
found['value']: found
for found in
[result for result in self.help_tags(company, empty, '')]
}
return adorned
def get_default_filter(self, data_type, company):
filter_spec = PartnersFilter(date=DateRangeFilter(
[datetime(2014, 1, 1), datetime.now()]))
adorned = adorn_filter(company, self, filter_spec)
return adorned
def test_django_field_tranform(self):
"""Transform the value in the Django field."""
field = DjangoField('tags', transform=extract_tags)
self.assertEqual(
[{'id': 1, 'name': 'somename', 'hex_color': 'bbccdd'}],
[dict(r) for r in field.get_value(MockRecord())])
def test_django_field_rename(self):
"""Rename the Django field in the Row."""
field = DjangoField('a', rename='b')
self.assertEqual('aa', field.get_value(MockRecord()))
self.assertEqual('b', field.get_name())
def test_django_dotted_name_field_none(self):
"""Recover from None in a dot relation."""
field = DjangoField('contact.n', rename='n')
self.assertEqual(None, field.get_value(MockRecord()))
self.assertEqual('n', field.get_name())
def test_django_dotted_name_field(self):
"""Follow a django dot relation."""
field = DjangoField('contact.a', rename='c')
self.assertEqual('aa', field.get_value(MockRecord()))
self.assertEqual('c', field.get_name())
def test_django_field(self):
"""Simple field extraction works."""
field = DjangoField('a')
self.assertEqual('aa', field.get_value(MockRecord()))
self.assertEqual('a', field.get_name())
def test_django_field_rename(self):
"""Rename the Django field in the Row."""
field = DjangoField('a', rename='b')
self.assertEqual('aa', field.get_value(MockRecord()))
self.assertEqual('b', field.get_name())
def test_django_dotted_name_field_none(self):
"""Recover from None in a dot relation."""
field = DjangoField('contact.n', rename='n')
self.assertEqual(None, field.get_value(MockRecord()))
self.assertEqual('n', field.get_name())
def test_django_dotted_name_field(self):
"""Follow a django dot relation."""
field = DjangoField('contact.a', rename='c')
self.assertEqual('aa', field.get_value(MockRecord()))
self.assertEqual('c', field.get_name())
def test_django_field(self):
"""Simple field extraction works."""
field = DjangoField('a')
self.assertEqual('aa', field.get_value(MockRecord()))
self.assertEqual('a', field.get_name())
