eulaurarien/feed_dns.py

#!/usr/bin/env python3

import argparse
import database
import json
import logging
import sys
import typing
import multiprocessing

NUMBER_THREADS = 2
BLOCK_SIZE = 100

# select, confirm, write
FUNCTION_MAP: typing.Any = {
    'a': (
        database.Database.get_ip4,
        database.Database.get_domain_in_zone,
        database.Database.set_hostname,
    ),
    'cname': (
        database.Database.get_domain,
        database.Database.get_domain_in_zone,
        database.Database.set_hostname,
    ),
    'ptr': (
        database.Database.get_domain,
        database.Database.get_ip4_in_network,
        database.Database.set_ip4address,
    ),
}


class Reader(multiprocessing.Process):
    def __init__(self,
                 lines_queue: multiprocessing.Queue,
                 write_queue: multiprocessing.Queue,
                 index: int = 0):
        super(Reader, self).__init__()
        self.log = logging.getLogger(f'rd{index:03d}')
        self.lines_queue = lines_queue
        self.write_queue = write_queue
        self.index = index

    def run(self) -> None:
        self.db = database.Database(write=False)
        self.db.log = logging.getLogger(f'db{self.index:03d}')
        self.db.enter_step('line_wait')
        block: typing.List[str]
        try:
            for block in iter(self.lines_queue.get, None):
                for line in block:
                    dtype, updated, name, value = line
                    self.db.enter_step('feed_switch')
                    select, confirm, write = FUNCTION_MAP[dtype]
                    for rule in select(self.db, value):
                        if not any(confirm(self.db, name)):
                            self.db.enter_step('wait_put')
                            self.write_queue.put((write, name, updated))
                    self.db.enter_step('line_wait')
        except KeyboardInterrupt:
            self.log.error('Interrupted')

        self.db.enter_step('end')
        self.db.close()


class Writer(multiprocessing.Process):
    def __init__(self,
                 write_queue: multiprocessing.Queue,
                 ):
        super(Writer, self).__init__()
        self.log = logging.getLogger(f'wr   ')
        self.write_queue = write_queue

    def run(self) -> None:
        self.db = database.Database(write=True)
        self.db.log = logging.getLogger(f'dbw  ')
        self.db.enter_step('line_wait')
        block: typing.List[str]
        try:
            fun: typing.Callable
            name: str
            updated: int
            for fun, name, updated in iter(self.write_queue.get, None):
                self.db.enter_step('exec')
                fun(self.db, name, updated)
                self.db.enter_step('line_wait')
        except KeyboardInterrupt:
            self.log.error('Interrupted')

        self.db.enter_step('end')
        self.db.close()


if __name__ == '__main__':

    # Parsing arguments
    log = logging.getLogger('feed_dns')
    parser = argparse.ArgumentParser(
        description="TODO")
    parser.add_argument(
        # '-i', '--input', type=argparse.FileType('rb'), default=sys.stdin.buffer,
        '-i', '--input', type=argparse.FileType('r'), default=sys.stdin,
        help="TODO")
    args = parser.parse_args()

    DB = database.Database(write=False)  # Not needed, just for timing
    DB.log = logging.getLogger('db   ')

    lines_queue: multiprocessing.Queue = multiprocessing.Queue(maxsize=100)
    write_queue: multiprocessing.Queue = multiprocessing.Queue(maxsize=100)

    DB.enter_step('proc_create')
    readers: typing.List[Reader] = list()
    for w in range(NUMBER_THREADS):
        readers.append(Reader(lines_queue, write_queue, w))
    writer = Writer(write_queue)

    DB.enter_step('proc_start')
    for reader in readers:
        reader.start()
    writer.start()

    try:
        block: typing.List[str] = list()
        DB.enter_step('iowait')
        for line in args.input:
            DB.enter_step('block_append')
            DB.enter_step('feed_json_parse')
            data = json.loads(line)
            line = (data['type'],
                    int(data['timestamp']),
                    data['name'],
                    data['value'])
            block.append(line)
            if len(block) >= BLOCK_SIZE:
                DB.enter_step('wait_put')
                lines_queue.put(block)
                block = list()
            DB.enter_step('iowait')
        DB.enter_step('wait_put')
        lines_queue.put(block)

        DB.enter_step('end_put')
        for _ in range(NUMBER_THREADS):
            lines_queue.put(None)
        write_queue.put(None)

        DB.enter_step('proc_join')
        for reader in readers:
            reader.join()
        writer.join()
    except KeyboardInterrupt:
        log.error('Interrupted')

    DB.close()