11"""
22This module implements a backend for sample storage in ClickHouse.
33"""
4+
45import base64
56import logging
67import time
4546}
4647
4748
49+ class ClickHouseBackendError (Exception ):
50+ """Something bad happened in the ClickHouse backend."""
51+
52+
4853def create_runs_table (client : clickhouse_driver .Client ):
4954 query = """
5055 CREATE TABLE IF NOT EXISTS runs (
@@ -79,7 +84,7 @@ def create_chain_table(client: clickhouse_driver.Client, meta: ChainMeta, rmeta:
7984 # Check that it does not already exist
8085 cid = chain_id (meta )
8186 if client .execute (f"SHOW TABLES LIKE '{ cid } ';" ):
82- raise Exception (f"A table for { cid } already exists." )
87+ raise ClickHouseBackendError (f"A table for { cid } already exists." )
8388
8489 # Create a table with columns corresponding to the model variables
8590 columns = []
@@ -235,7 +240,7 @@ def _get_row_at(
235240 query = f"SELECT (`{ names } `,) FROM { self .cid } WHERE _draw_idx={ idx } ;"
236241 data = self ._client .execute (query )
237242 if not data :
238- raise Exception (f"No record found for draw index { idx } ." )
243+ raise ClickHouseBackendError (f"No record found for draw index { idx } ." )
239244 result = dict (zip (var_names , data [0 ][0 ]))
240245 return result
241246
@@ -363,7 +368,10 @@ def __init__(
363368 raise ValueError ("Either a `client` or a `client_fn` must be provided." )
364369
365370 if client_fn is None :
366- client_fn = lambda : client
371+
372+ def client_fn ():
373+ return client
374+
367375 if client is None :
368376 client = client_fn ()
369377
@@ -381,11 +389,11 @@ def init_run(self, meta: RunMeta) -> ClickHouseRun:
381389 else :
382390 created_at = datetime .now ().astimezone (timezone .utc )
383391 query = "INSERT INTO runs (created_at, rid, proto) VALUES"
384- params = dict (
385- created_at = created_at ,
386- rid = meta .rid ,
387- proto = base64 .encodebytes (bytes (meta )).decode ("ascii" ),
388- )
392+ params = {
393+ " created_at" : created_at ,
394+ " rid" : meta .rid ,
395+ " proto" : base64 .encodebytes (bytes (meta )).decode ("ascii" ),
396+ }
389397 self ._client .execute (query , [params ])
390398 return ClickHouseRun (meta , client_fn = self ._client_fn , created_at = created_at )
391399
@@ -407,7 +415,9 @@ def get_run(self, rid: str) -> ClickHouseRun:
407415 {"rid" : rid },
408416 )
409417 if len (rows ) != 1 :
410- raise Exception (f"Unexpected number of { len (rows )} results for rid='{ rid } '." )
418+ raise ClickHouseBackendError (
419+ f"Unexpected number of { len (rows )} results for rid='{ rid } '."
420+ )
411421 data = base64 .decodebytes (rows [0 ][2 ].encode ("ascii" ))
412422 meta = RunMeta ().parse (data )
413423 return ClickHouseRun (
0 commit comments