def handle_telegram(telegram_payload):
update = Update.de_json(telegram_payload, bot)
if update.callback_query:
return handle_callback_query(update.callback_query)
message = update.message
if not message:
return
if message.text == "/start":
bot.sendMessage(message.chat.id, """This is a poll bot. To create a poll, add this bot to a group and send a poll question with "/new " prefix. For more information check out this page:\nhttps://itpp.dev/chat/opinions-bot/index.html""")
return
if DEBUG:
# Create tables
with CreateTableIfNotExists():
Poll.create_table(read_capacity_units=5, write_capacity_units=5, wait=True)
with CreateTableIfNotExists():
ddb_client = boto3.client('dynamodb')
DynamoDBLockClient.create_dynamodb_table(ddb_client)
poll_message = message.reply_to_message
if poll_message:
if poll_message.forward_from and poll_message.forward_from.id == BOT_ID:
# This may happen only is bot has access to all message.
# Reply to forwarded message, rather than original one.
bot.sendMessage(
message.chat.id,
"<em>%s</em>" % ON_RESPONSE_TO_FORWARDED_MESSAGE,
parse_mode='HTML',
reply_to_message_id=message.message_id
)
elif poll_message.from_user.id != BOT_ID:
# Reply to third-party message.
# Ignore.
logger.debug("Reply to third-party message: %s", [BOT_ID, poll_message.from_user.id])
else:
poll_key = message2poll_key(poll_message)
set_vote(message.from_user, poll_key, option_text=message.text, reply=message.message_id)
return
command, question = get_command_and_text(message.text or '')
if not command:
logger.debug("Message without command")
return
if not question:
logger.debug("Message without text")
return
create_poll(message, question)
def update_poll_message(poll):
# get a reference to the DynamoDB resource
dynamodb_resource = boto3.resource('dynamodb')
# create the lock-client
lock_client = DynamoDBLockClient(dynamodb_resource)
with lock_client.acquire_lock(
DYNAMO_DB_TABLE_NAME + ":" + poll.key,
retry_period=timedelta(seconds=UPDATING_POLL_MESSAGE_DELAY/3),
retry_timeout=timedelta(seconds=UPDATING_POLL_MESSAGE_DELAY*5),
raise_context_exception=True,
):
poll.refresh()
if poll.version == poll.telegram_version:
# no need to update
logger.debug("Poll was already updated. Exit")
return
now = get_now()
logger.debug("Checking for sleep: now is %s, last update is %s", now, poll.telegram_datetime)
since_last_update = (now - poll.telegram_datetime).total_seconds()
if since_last_update < UPDATING_POLL_MESSAGE_DELAY:
sleep = UPDATING_POLL_MESSAGE_DELAY - since_last_update
logger.debug("Sleep for %ssec", poll)
time.sleep(sleep)
poll.refresh()
chat_id, message_id = poll2chat_message_ids(poll)
# Update text
try:
bot.editMessageText(
poll2text(poll),
chat_id,
message_id,
parse_mode='HTML',
)
except telegram.error.BadRequest as e:
if e.message != "Message is not modified: specified new message content and reply markup are exactly the same as a current content and reply markup of the message":
raise
# Update Markup
bot.editMessageReplyMarkup(
chat_id,
message_id,
reply_markup=poll2markup(poll),
)
poll.telegram_version = poll.version + 1
poll.telegram_datetime = get_now()
poll.save()
def _setup_kv(**kwargs):
"""setup lock table and job completion tables in dynamodb"""
ddb = boto3.client("dynamodb")
# https://python-dynamodb-lock.readthedocs.io/en/latest/python_dynamodb_lock.html
log.info("creating dynamodb_lock table")
try:
DynamoDBLockClient.create_dynamodb_table(ddb)
except ClientError as clierr:
if clierr.response['Error']['Code'] == 'ResourceInUseException':
log.info("table already exists. skipping...")
pass
else:
raise
_create_job_table(ddb)
def with_distributed_lock(enabled: bool, lock_key: str):
# https://python-dynamodb-lock.readthedocs.io/en/latest/usage.html
if not enabled:
return
import boto3
from python_dynamodb_lock.python_dynamodb_lock import DynamoDBLockClient
with DynamoDBLockClient(boto3.resource("dynamodb")).acquire_lock(lock_key):
yield
def __init__(self, s3_bucket, s3_key, *, locking=True):
self._s3_bucket_name = s3_bucket
self._s3_key = s3_key
# self._timeout = timeout
self._opened = False
self._s3_bucket = boto3.resource('s3').Bucket(self._s3_bucket_name)
if self._s3_bucket.creation_date is None:
raise KeyError("No such bucket: %s" % self._s3_bucket_name)
ddb_client = boto3.client('dynamodb')
try:
DynamoDBLockClient.create_dynamodb_table(ddb_client)
except ddb_client.exceptions.ResourceInUseException:
# the table already exists -- ok with us
pass
if locking:
self._lock_client = DynamoDBLockClient(boto3.resource('dynamodb'))
self._lock = None
else:
self._lock_client = None
def lock_client():
# this is called only during build.
# not used at runtime.
from python_dynamodb_lock.python_dynamodb_lock import DynamoDBLockClient
if not lock_client.client:
# get a reference to the DynamoDB resource
dynamodb_resource = boto3.resource('dynamodb')
# create the lock-client
lock_client.client = DynamoDBLockClient(dynamodb_resource)
# FIXME close the lock_client on shutdown
#lock_client.close()
return lock_client.client
def main(repo, output_folder, user):
global INSTANCE_ID
# Notify the ASG LifeCycle hook that we are now In Service and ready to
# process requests/safe to be shut down
# Fetch current instance ID from where cloutinit writes it to
if not INSTANCE_ID:
with open('/var/lib/cloud/data/instance-id') as fh:
INSTANCE_ID = fh.readline().strip()
log.info("Starting on %s...", INSTANCE_ID)
output_folder = os.path.expanduser(output_folder)
short_time = datetime.timedelta(microseconds=1)
dynamodb = boto3.resource('dynamodb')
client = DynamoDBLockClient(
dynamodb,
table_name='GitHubRunnerLocks',
expiry_period=datetime.timedelta(0, 300),
heartbeat_period=datetime.timedelta(seconds=10),
)
# Just keep trying until we get some credentials.
while True:
# Have each runner try to get a credential in a random order.
possibles = get_possible_credentials(repo)
random.shuffle(possibles)
log.info("Trying to get a set of credentials in this order: %r",
possibles)
notify = get_sd_notify_func()
for index in possibles:
try:
lock = client.acquire_lock(
f'{repo}/{index}',
retry_period=short_time,
retry_timeout=short_time,
raise_context_exception=True,
)
except DynamoDBLockError as e:
log.info("Could not lock %s (%s)", index, e)
continue
with lock:
log.info("Obtained lock on %s", index)
write_credentials_to_files(repo, index, output_folder, user)
merge_in_settings(repo, output_folder)
notify(f"STATUS=Obtained lock on {index}")
if get_lifecycle_state() == "Pending:Wait":
complete_asg_lifecycle_hook()
notify("READY=1")
log.info("Watching for Runner.Worker processes")
ProcessWatcher().run()
client.close()
exit()
from typing import Optional
from datetime import timedelta
import boto3 # type: ignore
from contextlib import contextmanager
from src.config import LOCK_TABLE
from python_dynamodb_lock.python_dynamodb_lock import DynamoDBLockClient # type: ignore
dynamodb_resource = boto3.resource("dynamodb")
lock_client = DynamoDBLockClient(
dynamodb_resource,
table_name=LOCK_TABLE,
expiry_period=timedelta(
minutes=2), # The Lambda function has a 120 second timeout by default
# partition_key_name="key",
# sort_key_name="key",
)
@contextmanager
def dynamodb_lock(lock_name: str,
retry_timeout: Optional[timedelta] = timedelta(seconds=20)):
# TODO: Make this match get-lock-client in the clojure code
lock = lock_client.acquire_lock(lock_name,
sort_key=lock_name,
retry_timeout=retry_timeout)
try:
yield lock
finally:
lock.release(best_effort=True)
def provide_lock_client(self):
resource = boto3.resource("dynamodb",
endpoint_url=os.environ.get("DDB_ENDPOINT"))
return DynamoDBLockClient(resource,
table_name=os.environ["LOCK_TABLE_NAME"])
class AwsStateStorage:
def __init__(self, s3_bucket, s3_key, *, locking=True):
self._s3_bucket_name = s3_bucket
self._s3_key = s3_key
# self._timeout = timeout
self._opened = False
self._s3_bucket = boto3.resource('s3').Bucket(self._s3_bucket_name)
if self._s3_bucket.creation_date is None:
raise KeyError("No such bucket: %s" % self._s3_bucket_name)
ddb_client = boto3.client('dynamodb')
try:
DynamoDBLockClient.create_dynamodb_table(ddb_client)
except ddb_client.exceptions.ResourceInUseException:
# the table already exists -- ok with us
pass
if locking:
self._lock_client = DynamoDBLockClient(boto3.resource('dynamodb'))
self._lock = None
else:
self._lock_client = None
def open_and_read(self, read_cb):
assert not self._opened
if self._lock_client:
self._lock = self._lock_client.acquire_lock("{}/{}".format(
self._s3_bucket_name, self._s3_key))
try:
with io.BytesIO() as f:
try:
self._s3_bucket.download_fileobj(self._s3_key, f)
f.seek(0)
read_cb(f)
except botocore.exceptions.ClientError as e:
if e.response["Error"]["Code"] == "404":
read_cb(None)
else:
raise
self._opened = True
return self
except: # noqa: E722
if self._lock_client:
self._lock.release()
raise
def close(self):
assert self._opened
self._opened = False
if self._lock_client:
self._lock.release()
self._lock_client.close()
def write(self, write_cb):
assert self._opened
with io.StringIO() as f:
write_cb(f)
self._s3_bucket.put_object(Key=self._s3_key,
Body=f.getvalue().encode())