|
| 1 | +from functools import partial |
| 2 | + |
| 3 | +import ray |
| 4 | + |
| 5 | +import btrdb |
| 6 | +from btrdb.conn import BTrDB |
| 7 | + |
| 8 | +def register_serializer(conn_str=None, apikey=None, profile=None): |
| 9 | + """ |
| 10 | + Register serializer for BTrDB Object |
| 11 | + Parameters |
| 12 | + ---------- |
| 13 | + conn_str: str, default=None |
| 14 | + The address and port of the cluster to connect to, e.g. `192.168.1.1:4411`. |
| 15 | + If set to None, will look in the environment variable `$BTRDB_ENDPOINTS` |
| 16 | + (recommended). |
| 17 | + apikey: str, default=None |
| 18 | + The API key used to authenticate requests (optional). If None, the key |
| 19 | + is looked up from the environment variable `$BTRDB_API_KEY`. |
| 20 | + profile: str, default=None |
| 21 | + The name of a profile containing the required connection information as |
| 22 | + found in the user's predictive grid credentials file |
| 23 | + `~/.predictivegrid/credentials.yaml`. |
| 24 | + """ |
| 25 | + ray.register_custom_serializer( |
| 26 | + BTrDB, serializer=btrdb_serializer, deserializer=partial(btrdb_deserializer, conn_str=conn_str, apikey=apikey, profile=profile)) |
| 27 | + |
| 28 | +def btrdb_serializer(_): |
| 29 | + """ |
| 30 | + sererialize function |
| 31 | + """ |
| 32 | + return None |
| 33 | + |
| 34 | +def btrdb_deserializer(_, conn_str=None, apikey=None, profile=None): |
| 35 | + """ |
| 36 | + deserialize function |
| 37 | + |
| 38 | + Parameters |
| 39 | + ---------- |
| 40 | + conn_str: str, default=None |
| 41 | + The address and port of the cluster to connect to, e.g. `192.168.1.1:4411`. |
| 42 | + If set to None, will look in the environment variable `$BTRDB_ENDPOINTS` |
| 43 | + (recommended). |
| 44 | + apikey: str, default=None |
| 45 | + The API key used to authenticate requests (optional). If None, the key |
| 46 | + is looked up from the environment variable `$BTRDB_API_KEY`. |
| 47 | + profile: str, default=None |
| 48 | + The name of a profile containing the required connection information as |
| 49 | + found in the user's predictive grid credentials file |
| 50 | + `~/.predictivegrid/credentials.yaml`. |
| 51 | + Returns |
| 52 | + ------- |
| 53 | + db : BTrDB |
| 54 | + An instance of the BTrDB context to directly interact with the database. |
| 55 | + """ |
| 56 | + return btrdb.connect(conn_str=conn_str, apikey=apikey, profile=profile) |
0 commit comments