bin.Databases.ClickHouseDB

  1import collections
  2import threading
  3from Databases.database import Database
  4from clickhouse_driver import Client
  5from queue import Queue
  6
  7
  8class ClickHouseDB(Database):
  9    name = "clickhouse"
 10
 11    def __init__(self, config):
 12        """
 13        Args:
 14            config (dict): Format :
 15                            {"host":"127.0.0.1", "port":9000, "batch_size":15000}
 16        """
 17        super().__init__()
 18        self.started = False
 19        self.BATCH_SIZE = int(config["batch_size"]) if "batch_size" in config else 15000
 20        self.queue = Queue()
 21        self.client = Client(
 22            host=config["host"], port=int(config["port"]), compression="lz4"
 23        )
 24        self.client.execute(
 25            "CREATE TABLE IF NOT EXISTS bgp ("
 26            "time DateTime,"
 27            "type FixedString(1),"
 28            "peer Int32,"
 29            "collector String,"
 30            "project String,"
 31            "country String,"
 32            "source String,"
 33            "prefix String,"
 34            "path String"
 35            ") ENGINE = MergeTree ORDER BY (time, prefix, path)"
 36            "SETTINGS old_parts_lifetime=10"
 37        )
 38        print("Clickhouse : Generated BGP table :D")
 39
 40    def start(self):
 41        """
 42        Create bgp table
 43        (time DateTime, type, peer, collector, country, source, prefix, path)
 44        Start thread for batch inserts
 45        """
 46        # self.client.execute("DROP TABLE IF EXISTS bgp")
 47        self.started = True
 48        threading.Thread(
 49            target=self.insert_batches, daemon=True, name="BGP monitor - clickhouse"
 50        ).start()
 51
 52    def stop(self):
 53        """Stop inserts"""
 54        if self.started:
 55            self.started = False
 56
 57    ###############
 58    #   INSERTS   #
 59    ###############
 60
 61    def save(self, data):
 62        """Input data in queue for processing
 63
 64        Args:
 65            data (BGPElem): bgp element to save
 66        """
 67        self.queue.put(data)
 68
 69    def get_data(self):
 70        """Retrieve data for inserts
 71        Yields:
 72            dict: Data to insert
 73        """
 74        for idx in range(self.BATCH_SIZE):
 75            rec = self.queue.get()
 76            yield {
 77                "time": int(rec["time"]),
 78                "type": rec["type"],
 79                "peer": rec["peer_asn"],
 80                "collector": rec["collector"],
 81                "country": rec["country_code"] or "",
 82                "source": rec["source"] or "",
 83                "prefix": rec["prefix"],
 84                "path": rec.get("as-path", ""),
 85                "project": rec["project"],
 86            }
 87            if self.queue.qsize() == 0 and not self.started:
 88                print("Clickhouse : Inserting last batch")
 89                return
 90
 91    def insert_batches(self):
 92        """
 93        Insert BATCH_SIZE(15000 lines) batches
 94        https://clickhouse.com/docs/en/about-us/performance/"""
 95        while self.started:
 96            self.client.execute("INSERT INTO bgp VALUES ", self.get_data())
 97
 98    ##############
 99    #   GETTER   #
100    ##############
101
102    var_names = {
103        "time_start": "time_start",
104        "time_end": "time_end",
105        "record_type": "type",
106        "peer_asn": "peer",
107        "collectors": "collector",
108        "countries": "country",
109        "as_numbers": "source",
110        "prefixes": "prefix",
111        "as_paths": "path",
112    }
113    # Used in `ClickHouseDB.get` function as arg_name:db_column_name
114
115    def get(
116        self,
117        time_start,
118        time_end,
119        record_type=None,
120        peer_asn=None,
121        collectors=None,
122        countries=None,
123        as_numbers=None,
124        prefixes=None,
125        as_paths=None,
126    ):
127        query = "SELECT * FROM bgp WHERE"
128        params = {key: value for key, value in locals() if value is not None}
129
130        if len(params) > 0:
131            for k, v in params:
132                query += (
133                    f" %({ClickHouseDB.var_names[k]}) "
134                    + ("IN " if isinstance(v, collections.Sequence) else "== ")
135                    + f"%({k}) and"
136                )
137            query += query.rsplit(" ", 1)[0]
138        return self.client.execute_iter(query, params)