class PersistentBlueprintInstructionExecutionStore(BlueprintInstructionExecutionStore):
def __init__(self, manager: BlueprintManager, config):
super().__init__(manager, config)
self.config = config
self.db = PostgresqlExtDatabase(**config['db'])
self._initialize()
self._migrations()
self.receipthandle_by_instructionstateid = dict()
def remove_effects(self):
self.sqs.delete_queue(QueueUrl=self._queue_url)
self.db.drop_tables([BlueprintExecutionModel, BlueprintInstructionStateModel], safe=True)
def _initialize(self):
database_proxy.initialize(self.db)
sqs_region = self.config.get('sqs', {}).get('region_name')
kwargs = {} if not sqs_region else dict(region_name=sqs_region)
self.sqs = boto3.client('sqs', **kwargs)
def _get_queue_name(self):
queue_name = 'BlueprintInstructionExecutionStore'
return self.config['sqs'].get('prefix', '') + queue_name
def _migrations(self):
def does_queue_exist(queue_name):
response = self.sqs.list_queues()
for url in response.get('QueueUrls', []):
if queue_name in url:
self._queue_url = url
return True
return False
queue_name = self._get_queue_name()
self.db.create_tables([BlueprintExecutionModel, BlueprintInstructionStateModel], safe=True)
if not does_queue_exist(queue_name):
response = self.sqs.create_queue(QueueName=queue_name)
self._queue_url = response['QueueUrl']
def _store_blueprint_execution(self, blueprint_execution: BlueprintExecution):
bem = BlueprintExecutionModel(execution_id=blueprint_execution.execution_id, execution_context=blueprint_execution.execution_context,
blueprint=asdict(blueprint_execution.blueprint))
bem.save()
def _store_instruction_state(self, instruction_state: BlueprintInstructionState):
instruction_definition = asdict(instruction_state.instruction)
instr_state_model = BlueprintInstructionStateModel(instruction_state_id=instruction_state.id_,
blueprint_execution_id=instruction_state.blueprint_execution_id, instruction=instruction_definition,
status=instruction_state.status.value)
instr_state_model.save()
self.sqs.send_message(
QueueUrl=self._queue_url,
MessageBody=superjson(instruction_state)
)
def _get_instruction_to_process(self, worker_id) -> Optional[BlueprintInstructionState]:
response = self.sqs.receive_message(QueueUrl=self._queue_url, MaxNumberOfMessages=1)
messages = response.get('Messages')
if not messages:
return
b = json.loads(messages[0]['Body'])
self.receipthandle_by_instructionstateid[b['id_']] = messages[0]['ReceiptHandle']
try:
instruction_state_model = BlueprintInstructionStateModel.select().for_update().where(
(BlueprintInstructionStateModel.instruction_state_id == b['id_']) & (BlueprintInstructionStateModel.status == InstructionStatus.IDLE.value)).get()
except DoesNotExist:
instruction_state_model = None
if not instruction_state_model:
log.info("Got message with body {b} but did not get corresponding row in table. Might be a race condition.")
return
return BlueprintInstructionState(
instruction=self.manager.objectify_instruction(b['instruction']),
blueprint_execution_id=b['blueprint_execution_id'],
status=InstructionStatus(b['status']),
id_=b['id_']
)
def _remove_from_queue(self, instruction_state: BlueprintInstructionState):
self.sqs.delete_message(
QueueUrl=self._queue_url,
ReceiptHandle=self.receipthandle_by_instructionstateid[instruction_state.id_]
)
def _set_status_for_instruction(self, instruction_state: BlueprintInstructionState, status: InstructionStatus):
log.info(f"Setting instruction_state {instruction_state.id_}'s status to be {status}")
terminal_states = [InstructionStatus.SUCCESS, InstructionStatus.FAILED, InstructionStatus.END]
instruction_state.status = status
BlueprintInstructionStateModel.update(status=status.value).where(BlueprintInstructionStateModel.instruction_state_id == instruction_state.id_).execute()
if instruction_state.status in terminal_states:
self._remove_from_queue(instruction_state)
else:
# We're relying on the visibility timeout to retry
pass
def get_execution_context_from_id(self, blueprint_execution_id) -> Dict:
model: BlueprintExecutionModel = BlueprintExecutionModel.select().where(BlueprintExecutionModel.execution_id == blueprint_execution_id).get()
return dict(model.execution_context)
def get_instruction_to_process(self, worker_id=None) -> Optional[BlueprintInstructionState]:
with self.db.atomic():
return super().get_instruction_to_process(worker_id)
