lookyloo/bin/run_backend.py

#!/usr/bin/env python3

from __future__ import annotations

import argparse
import os
import time
from pathlib import Path
from subprocess import Popen

from redis import Redis
from redis.exceptions import ConnectionError

from lookyloo.default import get_homedir, get_socket_path, get_config


def check_running(name: str) -> bool:
    socket_path = get_socket_path(name)
    if not os.path.exists(socket_path):
        return False
    try:
        r = Redis(unix_socket_path=socket_path)
        return True if r.ping() else False
    except ConnectionError:
        return False


def launch_cache(storage_directory: Path | None=None) -> None:
    if not storage_directory:
        storage_directory = get_homedir()
    if not check_running('cache'):
        Popen(["./run_redis.sh"], cwd=(storage_directory / 'cache'))


def shutdown_cache(storage_directory: Path | None=None) -> None:
    if not storage_directory:
        storage_directory = get_homedir()
    r = Redis(unix_socket_path=get_socket_path('cache'))
    r.shutdown(save=True)
    print('Redis cache database shutdown.')


def launch_indexing(storage_directory: Path | None=None) -> None:
    if not storage_directory:
        storage_directory = get_homedir()
    if not check_running('indexing'):
        Popen(["./run_redis.sh"], cwd=(storage_directory / 'indexing'))


def shutdown_indexing(storage_directory: Path | None=None) -> None:
    if not storage_directory:
        storage_directory = get_homedir()
    r = Redis(unix_socket_path=get_socket_path('indexing'))
    r.shutdown(save=True)
    print('Redis indexing database shutdown.')


def launch_full_index(storage_directory: Path | None=None) -> None:
    if not storage_directory:
        storage_directory = get_homedir()
    if not check_running('full_index'):
        Popen(["./run_kvrocks.sh"], cwd=(storage_directory / 'full_index'))


def shutdown_full_index(storage_directory: Path | None=None) -> None:
    if not storage_directory:
        storage_directory = get_homedir()
    r = Redis(unix_socket_path=get_socket_path('full_index'))
    r.shutdown()
    print('Kvrocks full indexing database shutdown.')


def launch_all() -> None:
    launch_cache()
    launch_indexing()
    if get_config('generic', 'index_everything'):
        launch_full_index()


def check_all(stop: bool=False) -> None:
    backends: dict[str, bool] = {'cache': False, 'indexing': False}
    if get_config('generic', 'index_everything'):
        backends['full_index'] = False
    while True:
        for db_name in backends.keys():
            try:
                backends[db_name] = check_running(db_name)
            except Exception:
                backends[db_name] = False
        if stop:
            if not any(running for running in backends.values()):
                break
        else:
            if all(running for running in backends.values()):
                break
        for db_name, running in backends.items():
            if not stop and not running:
                print(f"Waiting on {db_name} to start")
            if stop and running:
                print(f"Waiting on {db_name} to stop")
        time.sleep(1)


def stop_all() -> None:
    shutdown_cache()
    shutdown_indexing()
    if get_config('generic', 'index_everything'):
        shutdown_full_index()


def main() -> None:
    parser = argparse.ArgumentParser(description='Manage backend DBs.')
    parser.add_argument("--start", action='store_true', default=False, help="Start all")
    parser.add_argument("--stop", action='store_true', default=False, help="Stop all")
    parser.add_argument("--status", action='store_true', default=True, help="Show status")
    args = parser.parse_args()

    if args.start:
        launch_all()
    if args.stop:
        stop_all()
    if not args.stop and args.status:
        check_all()


if __name__ == '__main__':
    main()