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" })