def _check_command_count(self, final=False):
"""Check the command count matches expectations.
:param final: This is the final count check, not incremental
:raise CommandSequenceError: When performing final count with no
commands
:raise InvalidDeclarationError: When commands processed do not match
the declared amount
"""
total = self.config.total_instructions if self.config else None
if final:
if not self.command_count:
raise CommandSequenceError("No instructions received")
if self.command_count != total:
raise InvalidDeclarationError(
f"Expected more instructions. Found {self.command_count} expected {total}"
)
return
if not self.config:
return
if self.command_count > total:
raise InvalidDeclarationError(
f"More instructions ({self.command_count}) received than declared ({total})"
)
def _configure(self, config):
"""Configure this object and the executor."""
if self.config is not None:
raise CommandSequenceError("Cannot currently re-configure jobs")
self.executor.configure(config)
self.config = config
def execute(self, batch, effective_user_id=None, **_): # pylint: disable=arguments-differ
if effective_user_id is None:
raise CommandSequenceError(
"Effective user must be configured before upserting groups")
# Check that we can actually process this batch
self._check_upsert_queries(
batch, expected_keys=["authority", "authority_provided_id"])
static_values = {
# Set the group to be owned by the effective user
"creator_id": effective_user_id,
# Set the group to match the specified type (private in this case)
"joinable_by": self.type_flags.joinable_by,
"readable_by": self.type_flags.readable_by,
"writeable_by": self.type_flags.writeable_by,
}
# Prep the query
values = [command.body.attributes for command in batch]
for value in values:
value.update(static_values)
stmt = insert(Group).values(values)
stmt = stmt.on_conflict_do_update(
index_elements=["authority", "authority_provided_id"],
set_={
"name": stmt.excluded.name
},
).returning(Group.id, Group.authority, Group.authority_provided_id)
# Upsert the data
try:
group_rows = self._execute_statement(stmt).fetchall()
except ProgrammingError as err:
# https://www.postgresql.org/docs/9.4/errcodes-appendix.html
# 21000 == cardinality violation
if err.orig.pgcode == "21000":
raise ConflictingDataError(
"Attempted to create two groups with the same authority and id"
) from err
raise
# Report back
return [
Report(
id_,
public_id=Group(
authority=authority,
authority_provided_id=authority_provided_id).groupid,
) for id_, authority, authority_provided_id in group_rows
]
def _add_to_batch(self, command):
"""Add a single command to the batch.
This may cause the CommandBatcher to call the on_flush() callback that we passed to it
(self._execute_batch()) if it decides that it's time to execute the next batch.
"""
if self.config is None:
raise CommandSequenceError("Not configured yet")
with self.batcher.add(command):
# If we have any id references like `"id": {"$ref": ...}` we need
# to fill these out before we pass them to the executor. We do this
# now to get the earliest warning if we have any id references
# which don't match up
self.id_refs.fill_out(command.body.raw)
def execute(self, batch, effective_user_id=None, **_):
if effective_user_id is None:
raise CommandSequenceError(
"Effective user must be configured before upserting groups")
# Check that we can actually process this batch
self._check_upsert_queries(
batch, expected_keys=["authority", "authority_provided_id"])
static_values = {
# Set the group to be owned by the effective user
"creator_id": effective_user_id,
# Set the group to match the specified type (private in this case)
"joinable_by": self.type_flags.joinable_by,
"readable_by": self.type_flags.readable_by,
"writeable_by": self.type_flags.writeable_by,
}
# Prep the query
values = [command.body.attributes for command in batch]
for value in values:
value.update(static_values)
stmt = insert(Group).values(values)
stmt = stmt.on_conflict_do_update(
index_elements=["authority", "authority_provided_id"],
set_={
"name": stmt.excluded.name
},
).returning(Group.id, Group.authority, Group.authority_provided_id)
# Upsert the data
group_rows = self._execute_statement(stmt).fetchall()
# Report back
return [
Report(
id_,
public_id=Group(
authority=authority,
authority_provided_id=authority_provided_id).groupid,
) for id_, authority, authority_provided_id in group_rows
]
def _check_sequence(self, command):
"""Check to see if the command sequence is wrong / requires a flush."""
command_key = (command.type, command.body.type)
if self.current_task is None:
self.current_task = command_key
return
if command_key == self.current_task:
return
if command_key in self.finished_tasks:
raise CommandSequenceError("Cannot return to old task")
# Looks like we are switching to a new task
self.flush()
self.finished_tasks.add(self.current_task)
self.current_task = command_key