async def merge_graph(self, db: DbAccess) -> GraphUpdate: # type: ignore
model = Model.from_kinds([kind async for kind in db.model_db.all()])
builder = GraphBuilder(model)
nxt = self.next_action()
while isinstance(nxt, ReadElement):
for element in nxt.jsons():
builder.add_from_json(element)
log.debug(f"Read {int(BatchSize / 1000)}K elements in process")
nxt = self.next_action()
if isinstance(nxt, PoisonPill):
log.debug("Got poison pill - going to die.")
shutdown_process(0)
elif isinstance(nxt, MergeGraph):
log.debug("Graph read into memory")
builder.check_complete()
graphdb = db.get_graph_db(nxt.graph)
outer_edge_db = db.get_pending_outer_edge_db()
_, result = await graphdb.merge_graph(builder.graph, model,
nxt.change_id, nxt.is_batch)
if nxt.task_id and builder.deferred_edges:
await outer_edge_db.update(
PendingDeferredEdges(nxt.task_id, nxt.graph,
builder.deferred_edges))
log.debug(
f"Updated {len(builder.deferred_edges)} pending outer edges for collect task {nxt.task_id}"
)
return result
def test_not_existing(system_db: StandardDatabase, test_db: StandardDatabase) -> None:
access = DbAccess(test_db, NoEventSender(), NoAdjust(), empty_config())
# foo db and pass does not exist
foodb = ["--graphdb-username", "foo", "--graphdb-password", "test", "--graphdb-database", "foo"]
system_db.delete_user("foo", ignore_missing=True)
system_db.delete_database("foo", ignore_missing=True)
access.connect(parse_args(foodb), timedelta(seconds=5), sleep_time=0.1)
assert system_db.has_user("foo")
assert system_db.has_database("foo")
def test_already_existing(test_db: StandardDatabase) -> None:
access = DbAccess(test_db, NoEventSender(), NoAdjust())
# test db and user already exist
testdb = [
"--graphdb-username", "test", "--graphdb-password", "test",
"--graphdb-database", "test"
]
access.connect(parse_args(testdb),
timedelta(seconds=0.01),
sleep_time=0.01)
def test_not_existing_and_default_root_account(
local_client: ArangoClient, system_db: StandardDatabase, test_db: StandardDatabase
) -> None:
access = DbAccess(test_db, NoEventSender(), NoAdjust(), empty_config())
# foo db and pass does not exist
foodb = ["--graphdb-username", "foo", "--graphdb-password", "bombproof", "--graphdb-database", "foo"]
system_db.delete_user("foo", ignore_missing=True)
system_db.delete_database("foo", ignore_missing=True)
access.connect(parse_args(foodb), timedelta(seconds=5), sleep_time=0.1)
# The default root account is used and a valid password is given -> also the root account uses this password
changed_root = local_client.db(username="******", password="******")
# Rest the password to the default one, to reset the state before the test
changed_root.replace_user("root", "", True)
async def cli_deps(
filled_graph_db: ArangoGraphDB,
message_bus: MessageBus,
event_sender: InMemoryEventSender,
foo_model: Model,
task_queue: WorkerTaskQueue,
worker: Tuple[WorkerTaskDescription, WorkerTaskDescription, WorkerTaskDescription],
expander: TemplateExpander,
config_handler: ConfigHandler,
cert_handler: CertificateHandler,
) -> AsyncIterator[CLIDependencies]:
db_access = DbAccess(filled_graph_db.db.db, event_sender, NoAdjust(), empty_config())
model_handler = ModelHandlerStatic(foo_model)
config = empty_config(["--graphdb-database", "test", "--graphdb-username", "test", "--graphdb-password", "test"])
deps = CLIDependencies(
message_bus=message_bus,
event_sender=event_sender,
db_access=db_access,
model_handler=model_handler,
worker_task_queue=task_queue,
config=config,
template_expander=expander,
forked_tasks=Queue(),
config_handler=config_handler,
cert_handler=cert_handler,
)
yield deps
await deps.stop()
async def setup_and_merge(self) -> GraphUpdate:
sender = InMemoryEventSender()
_, _, sdb = DbAccess.connect(self.args, timedelta(seconds=3))
db = db_access(self.args, sdb, sender)
result = await self.merge_graph(db)
for event in sender.events:
self.write_queue.put(EmitAnalyticsEvent(event))
return result
def run_process(args: Namespace) -> None:
with TemporaryDirectory() as temp_name:
temp = Path(temp_name)
with warnings.catch_warnings(): # ignore ssl errors during setup
warnings.simplefilter("ignore", HTTPWarning)
# wait here for an initial connection to the database before we continue. blocking!
created, system_data, sdb = DbAccess.connect(
args, timedelta(seconds=120), verify=False)
config = config_from_db(args, sdb)
cert_handler = CertificateHandler.lookup(config, sdb, temp)
verify: Union[bool,
str] = False if args.graphdb_no_ssl_verify else str(
cert_handler.ca_bundle)
config = replace(config, run=RunConfig(temp, verify))
# in case of tls: connect again with the correct certificate settings
use_tls = args.graphdb_server.startswith("https://")
db_client = DbAccess.connect(
args, timedelta(seconds=30), verify=verify)[2] if use_tls else sdb
with_config(created, system_data, db_client, cert_handler, config)
def db_access(config: Namespace, db: StandardDatabase, event_sender: AnalyticsEventSender) -> DbAccess:
adjuster = DirectAdjuster()
return DbAccess(db, event_sender, adjuster, update_outdated=config.graph_updates_abort_after)
def run(arguments: List[str]) -> None:
"""
Run application. When this method returns, the process is done.
:param arguments: the arguments provided to this process.
Note: this method is used in tests to specify arbitrary arguments.
"""
args = parse_args(arguments)
setup_process(args)
# after setup, logging is possible
info = system_info()
log.info(
f"Starting up version={info.version} on system with cpus={info.cpus}, "
f"available_mem={info.mem_available}, total_mem={info.mem_total}")
# wait here for an initial connection to the database before we continue. blocking!
created, system_data, sdb = DbAccess.connect(args, timedelta(seconds=60))
event_sender = NoEventSender(
) if args.analytics_opt_out else PostHogEventSender(system_data)
db = db_access(args, sdb, event_sender)
cert_handler = CertificateHandler.lookup(args, sdb)
message_bus = MessageBus()
scheduler = Scheduler()
worker_task_queue = WorkerTaskQueue()
model = ModelHandlerDB(db.get_model_db(), args.plantuml_server)
template_expander = DBTemplateExpander(db.template_entity_db)
config_handler = ConfigHandlerService(
db.config_entity_db,
db.config_validation_entity_db,
db.configs_model_db,
worker_task_queue,
message_bus,
)
cli_deps = CLIDependencies(
message_bus=message_bus,
event_sender=event_sender,
db_access=db,
model_handler=model,
worker_task_queue=worker_task_queue,
args=args,
template_expander=template_expander,
config_handler=config_handler,
)
default_env = {
"graph": args.cli_default_graph,
"section": args.cli_default_section
}
cli = CLI(cli_deps, all_commands(cli_deps), default_env, aliases())
subscriptions = SubscriptionHandler(db.subscribers_db, message_bus)
task_handler = TaskHandlerService(db.running_task_db, db.job_db,
message_bus, event_sender, subscriptions,
scheduler, cli, args)
cli_deps.extend(task_handler=task_handler)
api = Api(
db,
model,
subscriptions,
task_handler,
message_bus,
event_sender,
worker_task_queue,
cert_handler,
config_handler,
cli,
template_expander,
args,
)
event_emitter = emit_recurrent_events(event_sender, model, subscriptions,
worker_task_queue, message_bus,
timedelta(hours=1),
timedelta(hours=1))
async def on_start() -> None:
# queue must be created inside an async function!
cli_deps.extend(forked_tasks=Queue())
await db.start()
await event_sender.start()
await subscriptions.start()
await scheduler.start()
await worker_task_queue.start()
await event_emitter.start()
await cli.start()
await task_handler.start()
await api.start()
if created:
await event_sender.core_event(CoreEvent.SystemInstalled)
await event_sender.core_event(
CoreEvent.SystemStarted,
{
"version": version(),
"created_at": to_json(system_data.created_at),
"system": platform.system(),
"platform": platform.platform(),
"inside_docker": info.inside_docker,
},
cpu_count=info.cpus,
mem_total=info.mem_total,
mem_available=info.mem_available,
)
async def on_stop() -> None:
duration = utc() - info.started_at
await api.stop()
await task_handler.stop()
await cli.stop()
await event_sender.core_event(CoreEvent.SystemStopped,
total_seconds=int(
duration.total_seconds()))
await event_emitter.stop()
await worker_task_queue.stop()
await scheduler.stop()
await subscriptions.stop()
await db.stop()
await event_sender.stop()
async def async_initializer() -> Application:
async def on_start_stop(_: Application) -> AsyncIterator[None]:
await on_start()
log.info("Initialization done. Starting API.")
yield
log.info("Shutdown initiated. Stop all tasks.")
await on_stop()
api.app.cleanup_ctx.append(on_start_stop)
return api.app
tls_context: Optional[SSLContext] = None
if args.tls_cert:
tls_context = SSLContext(ssl.PROTOCOL_TLS)
tls_context.load_cert_chain(args.tls_cert, args.tls_key,
args.tls_password)
runner.run_app(async_initializer(),
api.stop,
host=args.host,
port=args.port,
ssl_context=tls_context)
def db_access(config: CoreConfig, db: StandardDatabase,
event_sender: AnalyticsEventSender) -> DbAccess:
adjuster = DirectAdjuster()
return DbAccess(db, event_sender, adjuster, config)