def get_user(session, update):
"""Get the user from the update."""
user = None
# Check user permissions
if hasattr(update, 'message') and update.message:
user = User.get_or_create(session, update.message.from_user)
if hasattr(update, 'edited_message') and update.edited_message:
user = User.get_or_create(session, update.edited_message.from_user)
elif hasattr(update, 'inline_query') and update.inline_query:
user = User.get_or_create(session, update.inline_query.from_user)
elif hasattr(update, 'callback_query') and update.callback_query:
user = User.get_or_create(session, update.callback_query.from_user)
return user
def get_user(session, tg_user):
"""Get the user from the event."""
user = session.query(User).get(tg_user.id)
if user is not None and user.banned:
return user, None
if user is None:
user = User(tg_user.id, tg_user.username)
session.add(user)
try:
session.commit()
increase_stat(session, "new_users")
# Handle race condition for parallel user addition
# Return the user that has already been created
# in another session
except IntegrityError as e:
session.rollback()
user = session.query(User).get(tg_user.id)
if user is None:
raise e
if tg_user.username is not None:
user.username = tg_user.username.lower()
name = get_name_from_tg_user(tg_user)
user.name = name
# Ensure user statistics exist for this user
# We need to track at least some user activity, since there seem to be some users which
# abuse the bot by creating polls and spamming up to 1 million votes per day.
#
# I really hate doing this, but I don't see another way to prevent DOS attacks
# without tracking at least some numbers.
user_statistic = session.query(UserStatistic).get((date.today(), user.id))
if user_statistic is None:
user_statistic = UserStatistic(user)
session.add(user_statistic)
try:
session.commit()
# Handle race condition for parallel user statistic creation
# Return the statistic that has already been created in another session
except IntegrityError as e:
session.rollback()
user_statistic = session.query(UserStatistic).get(
(date.today(), user.id))
if user_statistic is None:
raise e
return user, user_statistic
def user_factory(session, user_id, name, admin=False):
"""Create a user."""
user = User(user_id, name)
session.add(user)
session.commit()
return user
