def update_unit(
self, unit_id: str, agent_id: Optional[str] = None, status: Optional[str] = None
) -> None:
"""
Update the given task with the given parameters if possible, raise appropriate exception otherwise.
"""
if status not in AssignmentState.valid_unit():
raise MephistoDBException(f"Invalid status {status} for a unit")
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
if agent_id is not None:
c.execute(
"""
UPDATE units
SET agent_id = ?
WHERE unit_id = ?;
""",
(int(agent_id), int(unit_id)),
)
if status is not None:
c.execute(
"""
UPDATE units
SET status = ?
WHERE unit_id = ?;
""",
(status, int(unit_id)),
)
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException(
f"Given unit_id {unit_id} not found in the database"
)
raise MephistoDBException(e)
def new_worker(self, worker_name: str, provider_type: str) -> str:
"""
Create a new worker with the given name and provider type.
Raises EntryAlreadyExistsException
if there is already a worker with this name
worker_name should be the unique identifier by which the crowd provider
is using to keep track of this worker
"""
if worker_name == "":
raise MephistoDBException("Empty string is not a valid requester name")
assert_valid_provider(provider_type)
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"INSERT INTO workers(worker_name, provider_type) VALUES (?, ?);",
(worker_name, provider_type),
)
worker_id = str(c.lastrowid)
return worker_id
except sqlite3.IntegrityError as e:
if is_unique_failure(e):
raise EntryAlreadyExistsException()
raise MephistoDBException(e)
def new_onboarding_agent(
self, worker_id: str, task_id: str, task_run_id: str, task_type: str
) -> str:
"""
Create a new agent for the given worker id to assign to the given unit
Raises EntryAlreadyExistsException
"""
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"""INSERT INTO onboarding_agents(
worker_id,
task_id,
task_run_id,
task_type,
status
) VALUES (?, ?, ?, ?, ?);""",
(
int(worker_id),
int(task_id),
int(task_run_id),
task_type,
AgentState.STATUS_NONE,
),
)
return str(c.lastrowid)
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException(e)
raise MephistoDBException(e)
def assert_valid_provider(provider_type: str) -> None:
"""Throw an assertion error if the given provider type is not valid"""
valid_types = get_valid_provider_types()
if provider_type not in valid_types:
raise MephistoDBException(
f"Supplied provider {provider_type} is not in supported list of providers {valid_types}."
)
def new_agent(
self,
worker_id: str,
unit_id: str,
task_id: str,
task_run_id: str,
assignment_id: str,
task_type: str,
provider_type: str,
) -> str:
"""
Create a new agent with the given name and provider type.
Raises EntryAlreadyExistsException
if there is already a agent with this name
"""
assert_valid_provider(provider_type)
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"""INSERT INTO agents(
worker_id,
unit_id,
task_id,
task_run_id,
assignment_id,
task_type,
provider_type,
status
) VALUES (?, ?, ?, ?, ?, ?, ?, ?);""",
(
int(worker_id),
int(unit_id),
int(task_id),
int(task_run_id),
int(assignment_id),
task_type,
provider_type,
AgentState.STATUS_NONE,
),
)
agent_id = str(c.lastrowid)
c.execute(
"""
UPDATE units
SET status = ?, agent_id = ?, worker_id = ?
WHERE unit_id = ?;
""",
(
AssignmentState.ASSIGNED,
int(agent_id),
int(worker_id),
int(unit_id),
),
)
return agent_id
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException(e)
raise MephistoDBException(e)
def update_task(
self,
task_id: str,
task_name: Optional[str] = None,
project_id: Optional[str] = None,
) -> None:
"""
Update the given task with the given parameters if possible, raise appropriate exception otherwise.
Tasks can only be updated if no runs exist for this task yet, otherwise there's too much state
and we shouldn't make changes.
"""
if len(self.find_task_runs(task_id=task_id)) != 0:
raise MephistoDBException(
"Cannot edit a task that has already been run, for risk of data corruption."
)
if task_name in [""]:
raise MephistoDBException(f'Invalid task name "{task_name}')
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
if task_name is not None:
c.execute(
"""
UPDATE tasks
SET task_name = ?
WHERE task_id = ?;
""",
(task_name, int(task_id)),
)
if project_id is not None:
c.execute(
"""
UPDATE tasks
SET project_id = ?
WHERE task_id = ?;
""",
(int(project_id), int(task_id)),
)
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException(e)
elif is_unique_failure(e):
raise EntryAlreadyExistsException(
f"Task name {task_name} is already in use"
)
raise MephistoDBException(e)
def make_qualification(self, qualification_name: str) -> str:
"""
Make a new qualification, throws an error if a qualification by the given name
already exists. Return the id for the qualification.
"""
if qualification_name == "":
raise MephistoDBException("Empty string is not a valid qualification name")
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"INSERT INTO qualifications(qualification_name) VALUES (?);",
(qualification_name,),
)
qualification_id = str(c.lastrowid)
return qualification_id
except sqlite3.IntegrityError as e:
if is_unique_failure(e):
raise EntryAlreadyExistsException()
raise MephistoDBException(e)
def new_unit(
self,
task_id: str,
task_run_id: str,
requester_id: str,
assignment_id: str,
unit_index: int,
pay_amount: float,
provider_type: str,
task_type: str,
sandbox: bool = True,
) -> str:
"""
Create a new unit with the given index. Raises EntryAlreadyExistsException
if there is already a unit for the given assignment with the given index.
"""
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"""INSERT INTO units(
task_id,
task_run_id,
requester_id,
assignment_id,
unit_index,
pay_amount,
provider_type,
task_type,
sandbox,
status
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?);""",
(
int(task_id),
int(task_run_id),
int(requester_id),
int(assignment_id),
unit_index,
pay_amount,
provider_type,
task_type,
sandbox,
AssignmentState.CREATED,
),
)
unit_id = str(c.lastrowid)
return unit_id
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException(e)
elif is_unique_failure(e):
raise EntryAlreadyExistsException(e)
raise MephistoDBException(e)
def new_requester(self, requester_name: str, provider_type: str) -> str:
"""
Create a new requester with the given name and provider type.
Raises EntryAlreadyExistsException
if there is already a requester with this name
"""
if requester_name == "":
raise MephistoDBException("Empty string is not a valid requester name")
assert_valid_provider(provider_type)
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"INSERT INTO requesters(requester_name, provider_type) VALUES (?, ?);",
(requester_name, provider_type),
)
requester_id = str(c.lastrowid)
return requester_id
except sqlite3.IntegrityError as e:
if is_unique_failure(e):
raise EntryAlreadyExistsException()
raise MephistoDBException(e)
def new_task(
self,
task_name: str,
task_type: str,
project_id: Optional[str] = None,
parent_task_id: Optional[str] = None,
) -> str:
"""
Create a new task with the given task name. Raise EntryAlreadyExistsException if a task
with this name has already been created.
"""
if task_name in [""]:
raise MephistoDBException(f'Invalid task name "{task_name}')
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"""INSERT INTO tasks(
task_name,
task_type,
project_id,
parent_task_id
) VALUES (?, ?, ?, ?);""",
(
task_name,
task_type,
nonesafe_int(project_id),
nonesafe_int(parent_task_id),
),
)
task_id = str(c.lastrowid)
return task_id
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException(e)
elif is_unique_failure(e):
raise EntryAlreadyExistsException(e)
raise MephistoDBException(e)
def new_project(self, project_name: str) -> str:
"""
Create a new project with the given project name. Raise EntryAlreadyExistsException if a project
with this name has already been created.
"""
if project_name in [NO_PROJECT_NAME, ""]:
raise MephistoDBException(f'Invalid project name "{project_name}')
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"INSERT INTO projects(project_name) VALUES (?);", (project_name,)
)
project_id = str(c.lastrowid)
return project_id
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException()
elif is_unique_failure(e):
raise EntryAlreadyExistsException(
f"Project {project_name} already exists"
)
raise MephistoDBException(e)
def _get_connection(self) -> Connection:
"""Returns a singular database connection to be shared amongst all
calls for a given thread.
"""
# TODO(101) is there a problem with having just one db connection?
# Will this cause bugs with failed commits?
curr_thread = threading.get_ident()
if curr_thread not in self.conn or self.conn[curr_thread] is None:
try:
conn = sqlite3.connect(self.db_path)
conn.row_factory = StringIDRow
self.conn[curr_thread] = conn
except sqlite3.Error as e:
raise MephistoDBException(e)
return self.conn[curr_thread]
def update_agent(self, agent_id: str, status: Optional[str] = None) -> None:
"""
Update the given task with the given parameters if possible, raise appropriate exception otherwise.
"""
if status not in AgentState.valid():
raise MephistoDBException(f"Invalid status {status} for an agent")
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
c.execute(
"""
UPDATE agents
SET status = ?
WHERE agent_id = ?;
""",
(status, int(agent_id)),
)
def grant_qualification(
self, qualification_id: str, worker_id: str, value: int = 1
) -> None:
"""
Grant a worker the given qualification. Update the qualification value if it
already exists
"""
# Note that better syntax exists for python 3.8+, as described in PR #223
try:
# Update existing entry
qual_row = self.get_granted_qualification(qualification_id, worker_id)
with self.table_access_condition, self._get_connection() as conn:
if value != qual_row["value"]:
c = conn.cursor()
c.execute(
"""
UPDATE granted_qualifications
SET value = ?
WHERE (qualification_id = ?)
AND (worker_id = ?);
""",
(value, int(qualification_id), int(worker_id)),
)
conn.commit()
return None
except EntryDoesNotExistException:
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"""
INSERT INTO granted_qualifications(
qualification_id,
worker_id,
value
) VALUES (?, ?, ?);
""",
(int(qualification_id), int(worker_id), value),
)
qualification_id = str(c.lastrowid)
conn.commit()
return None
except sqlite3.IntegrityError as e:
if is_unique_failure(e):
raise EntryAlreadyExistsException()
raise MephistoDBException(e)
def update_task_run(self, task_run_id: str, is_completed: bool):
"""
Update a task run. At the moment, can only update completion status
"""
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"""
UPDATE task_runs
SET is_completed = ?
WHERE task_run_id = ?;
""",
(is_completed, int(task_run_id)),
)
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException(e)
raise MephistoDBException(e)
def new_task_run(
self,
task_id: str,
requester_id: str,
init_params: str,
provider_type: str,
task_type: str,
sandbox: bool = True,
) -> str:
"""Create a new task_run for the given task."""
with self.table_access_condition, self._get_connection() as conn:
# Ensure given ids are valid
c = conn.cursor()
try:
c.execute(
"""
INSERT INTO task_runs(
task_id,
requester_id,
init_params,
is_completed,
provider_type,
task_type,
sandbox
)
VALUES (?, ?, ?, ?, ?, ?, ?);""",
(
int(task_id),
int(requester_id),
init_params,
False,
provider_type,
task_type,
sandbox,
),
)
task_run_id = str(c.lastrowid)
return task_run_id
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException(e)
raise MephistoDBException(e)
def clear_unit_agent_assignment(self, unit_id: str) -> None:
"""
Update the given unit by removing the agent that is assigned to it, thus updating
the status to assignable.
"""
with self.table_access_condition, self._get_connection() as conn:
c = conn.cursor()
try:
c.execute(
"""
UPDATE units
SET agent_id = ?, worker_id = ?, status = ?
WHERE unit_id = ?;
""",
(None, None, AssignmentState.LAUNCHED, int(unit_id)),
)
except sqlite3.IntegrityError as e:
if is_key_failure(e):
raise EntryDoesNotExistException(
f"Given unit_id {unit_id} not found in the database"
)
raise MephistoDBException(e)