def ensure_indices(mongo_db: pymongo.database.Database) -> None:
"""Ensure that indices exist on relevant collections."""
db_collections = set(mongo_db.list_collection_names())
for collection, field in _INDICES:
if collection not in db_collections:
logging.error('The collection "%s" does not exist.', collection)
continue
mongo_db.get_collection(collection).create_index({field: 1})
def _iterate_all_records(mongo_db: pymongo.database.Database,
mongo_fields: Iterable[_MongoField],
field_type: str) -> Iterator[MongoReference]:
db_collections = set(mongo_db.list_collection_names())
for collection, field_name in mongo_fields:
if collection not in db_collections:
logging.error('The collection "%s" does not exist.', collection)
continue
records = mongo_db.get_collection(collection).\
find({field_name: {'$exists': True}}, {field_name: 1})
has_at_least_one_field = False
has_at_least_one_value = False
for record in records:
field_value = record[field_name]
has_at_least_one_field = True
if not field_value:
continue
has_at_least_one_value = True
yield MongoReference(collection, field_name, field_value,
record['_id'])
if not has_at_least_one_field:
logging.error('The collection "%s" has no field "%s".', collection,
field_name)
continue
if not has_at_least_one_value:
logging.error('The collection "%s" has no %ss in its field "%s"',
collection, field_type, field_name)
def _iterate_all_records(
mongo_db: pymongo.database.Database,
mongo_fields: Iterable[_MongoField], field_type: str) \
-> Iterator[Tuple[str, Dict[str, Any], str]]:
db_collections = set(mongo_db.list_collection_names())
for collection, field_name in mongo_fields:
if collection not in db_collections:
logging.error('The collection "%s" does not exist.', collection)
continue
records = mongo_db.get_collection(collection).\
find({field_name: {'$exists': True}}, {field_name: 1})
if not records.count():
logging.error('The collection "%s" has no field "%s".', collection, field_name)
continue
has_at_least_one_field = False
for record in records:
field_value = record[field_name]
if not field_value:
continue
has_at_least_one_field = True
yield field_value, record, collection
if not has_at_least_one_field:
logging.error(
'The collection "%s" has no %ss in its field "%s"',
collection, field_type, field_name)
def get_data_from_database(
database: pymongo.database.Database) -> pymongo.cursor.Cursor:
collection = database.get_collection(COLLECTION_NAME)
collection_contents = collection.find()
return collection_contents
def send_mail(
self, campaign_id: str, user: _UserProto, *, database: pymongo.database.Database,
users_database: pymongo.database.Database, now: datetime.datetime,
action: 'Action' = 'dry-run',
dry_run_email: Optional[str] = None,
mongo_user_update: Optional[Dict[str, Any]] = None) -> bool:
"""Send an email for this campaign."""
template_vars = self._get_vars(
user, database=database, users_database=users_database, now=now)
if not template_vars:
return False
collection = self._users_collection
if action == 'list':
user_id = collection.get_id(user)
logging.info('%s: %s %s', campaign_id, user_id, collection.get_profile(user).email)
return True
if action == 'dry-run':
collection.get_profile(user).email = dry_run_email or '*****@*****.**'
if action == 'ghost':
email_sent = user.emails_sent.add()
email_sent.sent_at.FromDatetime(now)
email_sent.sent_at.nanos = 0
email_sent.subject = get_campaign_subject(self._mailjet_template) or ''
else:
res = mail.send_template(
self._mailjet_template, collection.get_profile(user), template_vars,
sender_email=self._sender_email, sender_name=self._sender_name,
campaign_id=campaign_id)
logging.info('Email sent to %s', collection.get_profile(user).email)
res.raise_for_status()
maybe_email_sent = mail.create_email_sent_proto(res)
if not maybe_email_sent:
logging.warning('Impossible to retrieve the sent email ID:\n%s', res.json())
return False
if action == 'dry-run':
return True
email_sent = maybe_email_sent
email_sent.mailjet_template = self._mailjet_template
email_sent.campaign_id = campaign_id
if mongo_user_update and '$push' in mongo_user_update: # pragma: no-cover
raise ValueError(
f'$push operations are not allowed in mongo_user_update:\n{mongo_user_update}')
user_id = collection.get_id(user)
if user_id and action != 'ghost':
users_database.get_collection(collection.mongo_collection).update_one(
{'_id': objectid.ObjectId(user_id)},
dict(mongo_user_update or {}, **{'$push': {
'emailsSent': json_format.MessageToDict(email_sent),
}}))
# TODO(pascal): Clean that up or make it work in ghost mode.
if self._on_email_sent:
self._on_email_sent(
user, email_sent=email_sent, template_vars=template_vars,
database=database, user_database=users_database)
return True