How to structure in handling websocket and inserting data to databse in Python

Question

I connected to the websocket, and the data response in the callback function handle_message. I wish to insert all the data to influxdb using write_api.write(), since I need with InfluxClient() to wrap the whole CRUD function, I wrote like this:

from time import sleep
from pybit import usdt_perpetual
from dataclasses import dataclass
from influxdb_client import InfluxDBClient
from influxdb_client.client.write_api import SYNCHRONOUS
@dataclass
class Trade:
 tick_direction: str
 side: float
 is_block_trade: str
 price: float
 size: float
 timestamp: int
ws_linear = usdt_perpetual.WebSocket(
 test=False,
 ping_interval=30, 
 ping_timeout=10, 
 domain="bybit" 
)
with InfluxDBClient(
 url="http://localhost:8086",
 token="fake_token_id",
 org="organization") as client:
 write_api = client.write_api(write_options=SYNCHRONOUS)
 def handle_message(msg: dict) -> None:
 for data in msg["data"]:
 trade = Trade(data["tick_direction"], data["side"], data["is_block_trade"], float(data["price"]),
 float(data["size"]), data["timestamp"])
 write_api.write(bucket="test5",
 record=trade,
 record_measurement_name="tradeBTCUSDT",
 record_time_key="timestamp",
 record_tag_keys=["tick_direction", "side", "is_block_trade"],
 record_field_keys=["price", "size"])
ws_linear.trade_stream(
 handle_message, "BTCUSDT"
)
while True:
 sleep(1)

I think it is messy that I wrap handle_message by with InfluxClient, how can I improve it and the overall structure of the program?

Answer 1

I find it unlikely this program works as expected. You have written approximately:

1. Open a database connection, & assign that to client
2. Get write_api from client
3. Define a function which (at some point in the future) uses write_api
4. Call client.close() due to the exit of the with statement!
5. Ask another function to handle messages by calling your defined function which uses an api on a "now closed" client!

As you can see, the with statement is not doing anything but giving you a false sense of security. You need to either:

• open a new database connection inside handle_message, or
• indent the subsequent code such that client is not closed until after your while True: loop exits (presumably by a KeyboardInterrupt)

---

Here is a toy demonstration of the issue:

class Client:
 def __init__(self):
 self.opened = False
 
 def __enter__(self):
 self.opened = True
 print("Client opened")
 return self
 def __exit__(self, *args):
 self.opened = False
 print("Client closed")
 def write_api(self):
 return WriteApi(self)
class WriteApi():
 def __init__(self, client):
 self.client = client
 def use_api(self):
 print("Using API. Client open =", self.client.opened)
def trade_stream(handler):
 handler(None)
print("About to enter with-statement")
with Client() as client:
 print("Entered with-statement")
 write_api = client.write_api()
 
 def handle_message(msg):
 write_api.use_api()
 print("Exiting with-statement")
print("Exited with-statement")
trade_stream(handle_message)

Output:

About to enter with-statement
Client opened
Entered with-statement
Exiting with-statement
Client closed
Exited with-statement
Using API. Client open = False

Note that using the API happens after the client has been closed.

---

So why do you appear to have working code?

From the documentation ...

from influxdb_client import InfluxDBClient
# Initialize background batching instance of WriteApi
with InfluxDBClient(url="http://localhost:8086", token="my-token", org="my-org") as client:
 with client.write_api() as write_api:
 pass

... it appears client.write_api() returns an object which itself can behave as a context manager, and so should be used in a with-statement.

If the write_api is "open", and then the client is closed, whether or not the write_api is invalid will depend on how the library is written. It may internally hold the client open, until the write_api is explicitly closed (which you are not doing), or it may be Undefined Behaviour and might appear to work at the moment, and start failing in the future.

---

As stated earlier, the fix could be to indent the subsequent code, as in:

with InfluxDBClient(
 url="http://localhost:8086",
 token="fake_token_id",
 org="organization") as client:
 write_api = client.write_api(write_options=SYNCHRONOUS)
 def handle_message(msg: dict) -> None:
 ...
 # still indented, so still inside the `with client`
 ws_linear.trade_stream(
 handle_message, "BTCUSDT"
 )
 # still indented, so still inside the `with client`
 while True:
 sleep(1)
 # with statement ends here
# client will be closed here