def update_properties(
self,
database_client:
DatabaseClient, # pylint: disable = too-many-arguments
worker: dict,
version: Optional[str] = None,
display_name: Optional[str] = None,
properties: Optional[dict] = None) -> None:
now = self.date_time_provider.now()
update_data = {
"version": version,
"display_name": display_name,
"properties": properties,
"update_date": self.date_time_provider.serialize(now),
}
update_data = {
key: value
for key, value in update_data.items() if value is not None
}
worker.update(update_data)
database_client.update_one(self.table,
{"identifier": worker["identifier"]},
update_data)
def create(
self,
database_client:
DatabaseClient, # pylint: disable = too-many-arguments
worker_identifier: str,
owner: str,
version: str,
display_name: str) -> dict:
now = self.date_time_provider.now()
worker = {
"identifier": worker_identifier,
"owner": owner,
"version": version,
"display_name": display_name,
"properties": {},
"is_enabled": True,
"is_active": False,
"should_disconnect": False,
"creation_date": self.date_time_provider.serialize(now),
"update_date": self.date_time_provider.serialize(now),
}
database_client.insert_one(self.table, worker)
return worker
def create_or_update(self, database_client: DatabaseClient, # pylint: disable = too-many-arguments
schedule_identifier: str, project: str, display_name: str, job: str, parameters: dict, expression: str) -> dict:
now = self.date_time_provider.now()
schedule = self.get(database_client, project, schedule_identifier)
if schedule is None:
schedule = {
"project": project,
"identifier": schedule_identifier,
"display_name": display_name,
"job": job,
"parameters": parameters,
"expression": expression,
"is_enabled": False,
"last_run": None,
"creation_date": self.date_time_provider.serialize(now),
"update_date": self.date_time_provider.serialize(now),
}
database_client.insert_one(self.table, schedule)
else:
update_data = {
"display_name": display_name,
"job": job,
"parameters": parameters,
"expression": expression,
"update_date": self.date_time_provider.serialize(now),
}
schedule.update(update_data)
database_client.update_one(self.table, { "project": project, "identifier": schedule_identifier }, update_data)
return schedule
def update_status(
self,
database_client:
DatabaseClient, # pylint: disable = too-many-arguments
worker: dict,
is_active: Optional[bool] = None,
is_enabled: Optional[bool] = None,
should_disconnect: Optional[bool] = None) -> None:
now = self.date_time_provider.now()
update_data = {
"is_active": is_active,
"is_enabled": is_enabled,
"should_disconnect": should_disconnect,
"update_date": self.date_time_provider.serialize(now),
}
update_data = {
key: value
for key, value in update_data.items() if value is not None
}
worker.update(update_data)
database_client.update_one(self.table,
{"identifier": worker["identifier"]},
update_data)
def create(self, database_client: DatabaseClient, # pylint: disable = too-many-arguments
project: str, job: str, parameters: dict, source: dict) -> dict:
identifier = self.create_identifier(database_client)
now = self.date_time_provider.now()
run = {
"identifier": identifier,
"project": project,
"job": job,
"parameters": parameters,
"source": source,
"worker": None,
"status": "pending",
"start_date": None,
"completion_date": None,
"results": None,
"should_cancel": False,
"should_abort": False,
"creation_date": self.date_time_provider.serialize(now),
"update_date": self.date_time_provider.serialize(now),
}
database_client.insert_one(self.table, run)
return run
def create_token(self, database_client: DatabaseClient, user: str, description: str, expiration: Optional[datetime.timedelta]) -> dict:
now = self.date_time_provider.now()
token = {
"identifier": str(uuid.uuid4()),
"user": user,
"type": "token",
"description": description,
"hash_function": self.token_hash_function,
"hash_function_parameters": self.token_hash_function_parameters,
"hash_function_salt": None,
"expiration_date": None,
"creation_date": self.date_time_provider.serialize(now),
"update_date": self.date_time_provider.serialize(now),
}
if expiration is not None:
token["expiration_date"] = self.date_time_provider.serialize(now + expiration)
secret = secrets.token_hex(self.token_size)
token["secret"] = self.hash_token(secret, token["hash_function"], token["hash_function_parameters"])
database_client.insert_one(self.table, token)
result = self.convert_to_public(token)
result["secret"] = secret
return result
def set_password(self, database_client: DatabaseClient, user: str, password: str) -> dict:
now = self.date_time_provider.now()
authentication = database_client.find_one(self.table, { "user": user, "type": "password" })
if authentication is None:
authentication = {
"identifier": str(uuid.uuid4()),
"user": user,
"type": "password",
"creation_date": self.date_time_provider.serialize(now),
"update_date": self.date_time_provider.serialize(now),
}
database_client.insert_one(self.table, authentication)
authentication.update({
"hash_function": self.password_hash_function,
"hash_function_parameters": self.password_hash_function_parameters,
"hash_function_salt": secrets.token_hex(self.password_salt_size),
"update_date": self.date_time_provider.serialize(now),
})
authentication["secret"] = self.hash_password(password, authentication["hash_function_salt"], authentication["hash_function"], authentication["hash_function_parameters"])
database_client.update_one(self.table, { "identifier": authentication["identifier"] }, authentication)
return self.convert_to_public(authentication)
def update_roles(self, database_client: DatabaseClient, user: dict, roles: List[str]) -> None:
now = self.date_time_provider.now()
update_data = {
"roles": roles,
"update_date": self.date_time_provider.serialize(now),
}
user.update(update_data)
database_client.update_one(self.table, { "identifier": user["identifier"] }, update_data)
def set_results(self, database_client: DatabaseClient, run: dict, results: dict) -> None:
now = self.date_time_provider.now()
update_data = {
"results": results,
"update_date": self.date_time_provider.serialize(now),
}
run.update(update_data)
database_client.update_one(self.table, { "project": run["project"], "identifier": run["identifier"] }, update_data)
def update_identity(self, database_client: DatabaseClient, user: dict, display_name: Optional[str] = None) -> None:
now = self.date_time_provider.now()
update_data = {
"display_name": display_name,
"update_date": self.date_time_provider.serialize(now),
}
update_data = { key: value for key, value in update_data.items() if value is not None }
user.update(update_data)
database_client.update_one(self.table, { "identifier": user["identifier"] }, update_data)
def update_status(self, database_client: DatabaseClient, user: dict, is_enabled: Optional[bool] = None) -> None:
now = self.date_time_provider.now()
update_data = {
"is_enabled": is_enabled,
"update_date": self.date_time_provider.serialize(now),
}
update_data = { key: value for key, value in update_data.items() if value is not None }
user.update(update_data)
database_client.update_one(self.table, { "identifier": user["identifier"] }, update_data)
def set_token_expiration(self, database_client: DatabaseClient, user_identifier: str, token_identifier: str, expiration: datetime.timedelta) -> None:
token = database_client.find_one(self.table, { "identifier": token_identifier, "user": user_identifier, "type": "token" })
if token["expiration_date"] is None:
raise ValueError("Token '%s' does not expire" % token_identifier)
now = self.date_time_provider.now()
update_data = {
"expiration_date": self.date_time_provider.serialize(now + expiration),
"update_date": self.date_time_provider.serialize(now),
}
database_client.update_one(self.table, { "identifier": token_identifier, "user": user_identifier, "type": "token" }, update_data)
def import_table(database_client: DatabaseClient,
table: str,
source_directory: str,
simulate: bool = False) -> None:
logger.info("Importing table '%s'", table)
source_file_path = os.path.join(source_directory, table + ".json")
with open(source_file_path, mode="r", encoding="utf-8") as source_file:
dataset = json.load(source_file)
if not simulate:
if len(dataset) > 0:
database_client.insert_many(table, dataset)
def update_status(self, database_client: DatabaseClient,
schedule: dict, is_enabled: Optional[bool] = None, last_run: Optional[str] = None) -> None:
now = self.date_time_provider.now()
update_data = {
"is_enabled": is_enabled,
"last_run": last_run,
"update_date": self.date_time_provider.serialize(now),
}
update_data = { key: value for key, value in update_data.items() if value is not None }
schedule.update(update_data)
database_client.update_one(self.table, { "project": schedule["project"], "identifier": schedule["identifier"] }, update_data)
def delete(self, database_client: DatabaseClient,
user_identifier: str, authentication_provider: AuthenticationProvider, worker_provider: WorkerProvider) -> None:
user_record = self.get(database_client, user_identifier)
if user_record is None:
raise ValueError("User '%s' does not exist" % user_identifier)
if user_record["is_enabled"]:
raise ValueError("User '%s' is enabled" % user_identifier)
if worker_provider.count(database_client, owner = user_identifier) > 0:
raise ValueError("User '%s' owns workers" % user_identifier)
authentication_provider.remove_password(database_client, user_identifier)
for token in authentication_provider.get_token_list(database_client, user_identifier):
authentication_provider.delete_token(database_client, user_identifier, token["identifier"])
database_client.delete_one(self.table, { "identifier": user_identifier })
def create_or_update(
self,
database_client:
DatabaseClient, # pylint: disable = too-many-arguments
job_identifier: str,
project: str,
display_name: str,
description: str,
definition: dict,
parameters: list,
properties: dict) -> dict:
now = self.date_time_provider.now()
job = self.get(database_client, project, job_identifier)
if job is None:
job = {
"project": project,
"identifier": job_identifier,
"display_name": display_name,
"description": description,
"definition": definition,
"parameters": parameters,
"properties": properties,
"is_enabled": True,
"creation_date": self.date_time_provider.serialize(now),
"update_date": self.date_time_provider.serialize(now),
}
database_client.insert_one(self.table, job)
else:
update_data = {
"display_name": display_name,
"description": description,
"definition": definition,
"parameters": parameters,
"properties": properties,
"update_date": self.date_time_provider.serialize(now),
}
job.update(update_data)
database_client.update_one(self.table, {
"project": project,
"identifier": job_identifier
}, update_data)
return job
def get_list(self, database_client: DatabaseClient, # pylint: disable = too-many-arguments
project: Optional[str] = None, job: Optional[str] = None,
skip: int = 0, limit: Optional[int] = None, order_by: Optional[List[Tuple[str,str]]] = None) -> List[dict]:
filter = { "project": project, "job": job } # pylint: disable = redefined-builtin
filter = { key: value for key, value in filter.items() if value is not None }
return database_client.find_many(self.table, filter, skip = skip, limit = limit, order_by = order_by)
def get_list(self, database_client: DatabaseClient, # pylint: disable = too-many-arguments
project: Optional[str] = None, job: Optional[str] = None, worker: Optional[str] = None, status: Optional[str] = None,
skip: int = 0, limit: Optional[int] = None, order_by: Optional[List[Tuple[str,str]]] = None) -> List[dict]:
filter = { "project": project, "job": job, "worker": worker, "status": status } # pylint: disable = redefined-builtin
filter = { key: value for key, value in filter.items() if value is not None }
run_collection = database_client.find_many(self.table, filter, skip = skip, limit = limit, order_by = order_by)
return [ self.convert_to_public(run) for run in run_collection ]
def get_token_list(self, # pylint: disable = too-many-arguments
database_client: DatabaseClient, user: Optional[str] = None,
skip: int = 0, limit: Optional[int] = None, order_by: Optional[List[Tuple[str,str]]] = None) -> List[dict]:
filter = { "user": user, "type": "token" } # pylint: disable = redefined-builtin
filter = { key: value for key, value in filter.items() if value is not None }
token_list = database_client.find_many(self.table, filter, skip = skip, limit = limit, order_by = order_by)
return [ self.convert_to_public(token) for token in token_list ]
def create(self, database_client: DatabaseClient, user_identifier: str, display_name: str) -> dict:
if self.user_identifier_regex.search(user_identifier) is None:
raise ValueError("User identifier is invalid: '%s'" % user_identifier)
now = self.date_time_provider.now()
user = {
"identifier": user_identifier,
"display_name": display_name,
"roles": [],
"is_enabled": True,
"creation_date": self.date_time_provider.serialize(now),
"update_date": self.date_time_provider.serialize(now),
}
database_client.insert_one(self.table, user)
return user
def count(self,
database_client: DatabaseClient,
owner: Optional[str] = None) -> int:
filter = {"owner": owner} # pylint: disable = redefined-builtin
filter = {
key: value
for key, value in filter.items() if value is not None
}
return database_client.count(self.table, filter)
def delete(self, database_client: DatabaseClient, worker_identifier: str,
run_provider: RunProvider) -> None:
worker_record = self.get(database_client, worker_identifier)
if worker_record is None:
raise ValueError("Worker '%s' does not exist" % worker_identifier)
if worker_record["is_enabled"]:
raise ValueError("Worker '%s' is enabled" % worker_identifier)
if worker_record["is_active"]:
raise ValueError("Worker '%s' is active" % worker_identifier)
if run_provider.count(worker=worker_identifier, status="pending") > 0:
raise ValueError("Worker '%s' has pending runs" %
worker_identifier)
if run_provider.count(worker=worker_identifier, status="running") > 0:
raise ValueError("Worker '%s' has running runs" %
worker_identifier)
database_client.delete_one(self.table,
{"identifier": worker_identifier})
def check_if_empty(database_client: DatabaseClient,
all_tables: List[str]) -> None:
is_empty = True
for table in all_tables:
if database_client.find_one(table, {}) is not None:
is_empty = False
logger.error("Table '%s' is not empty", table)
if not is_empty:
raise ValueError("Database is not empty")
def authenticate_with_token(self, database_client: DatabaseClient, user_identifier: str, secret: str) -> bool:
now = self.date_time_provider.serialize(self.date_time_provider.now())
user_tokens = database_client.find_many(self.table, { "user": user_identifier, "type": "token" })
for token in user_tokens:
if token["expiration_date"] is None or token["expiration_date"] > now:
hashed_secret = self.hash_token(secret, token["hash_function"], token["hash_function_parameters"])
if secrets.compare_digest(hashed_secret, token["secret"]):
return True
return False
def update_status(self, database_client: DatabaseClient, # pylint: disable = too-many-arguments
run: dict, worker: Optional[str] = None, status: Optional[str] = None,
start_date: Optional[str] = None, completion_date: Optional[str] = None,
should_cancel: Optional[bool] = None, should_abort: Optional[bool] = None) -> None:
now = self.date_time_provider.now()
update_data = {
"worker": worker,
"status": status,
"start_date": start_date,
"completion_date": completion_date,
"should_cancel": should_cancel,
"should_abort": should_abort,
"update_date": self.date_time_provider.serialize(now),
}
update_data = { key: value for key, value in update_data.items() if value is not None }
run.update(update_data)
database_client.update_one(self.table, { "project": run["project"], "identifier": run["identifier"] }, update_data)
def export_table(database_client: DatabaseClient,
table: str,
output_directory: str,
simulate: bool = False) -> None:
logger.info("Exporting table '%s'", table)
dataset = database_client.find_many(table, {})
output_file_path = os.path.join(output_directory, table + ".json")
if not simulate:
with open(output_file_path + ".tmp", mode="w",
encoding="utf-8") as output_file:
json.dump(dataset, output_file, indent=4)
os.replace(output_file_path + ".tmp", output_file_path)
def get_archive(self, database_client: DatabaseClient, project: str, run_identifier: str) -> dict:
run = database_client.find_one(self.table, { "project": project, "identifier": run_identifier })
if run is None:
return None
file_name = run_identifier + ".zip"
now = time.gmtime()
with io.BytesIO() as file_object:
with zipfile.ZipFile(file_object, mode = "w", compression = zipfile.ZIP_DEFLATED) as archive:
entry_info = zipfile.ZipInfo("run.json", now[0:6])
entry_info.external_attr = 0o644 << 16
archive.writestr(entry_info, json.dumps(run, indent = 4))
entry_info = zipfile.ZipInfo("run.log", now[0:6])
entry_info.external_attr = 0o644 << 16
archive.writestr(entry_info, self.get_log(project, run_identifier)[0])
return { "file_name": file_name, "data": file_object.getvalue(), "type": "zip" }
def count_tokens(self, database_client: DatabaseClient, user: Optional[str] = None) -> int:
filter = { "user": user, "type": "token" } # pylint: disable = redefined-builtin
filter = { key: value for key, value in filter.items() if value is not None }
return database_client.count(self.table, filter)
def authenticate_with_password(self, database_client: DatabaseClient, user_identifier: str, password: str) -> bool:
authentication = database_client.find_one(self.table, { "user": user_identifier, "type": "password" })
if authentication is None:
return False
hashed_password = self.hash_password(password, authentication["hash_function_salt"], authentication["hash_function"], authentication["hash_function_parameters"])
return secrets.compare_digest(hashed_password, authentication["secret"])
def remove_password(self, database_client: DatabaseClient, user: str) -> None:
database_client.delete_one(self.table, { "user": user, "type": "password" })