improve logging

Author: dforsber

main.py

@@ -1,10 +1,9 @@
 import asyncio
 from py_boilingdata import BoilingData
 
-boiling = BoilingData()
-
 
 async def main():
+    boiling = BoilingData()
     await boiling.connect()
     results = await boiling.execute(
         """
@@ -17,5 +16,4 @@ async def main():
     await boiling.close()
 
 
-loop = asyncio.get_event_loop()
-loop.run_until_complete(main())
+asyncio.new_event_loop().run_until_complete(main())

py_boilingdata/__init__.py

@@ -38,21 +38,24 @@
 class BoilingData:
     """Run SQL with BoilingData and local DuckDB"""
 
-    def __init__(self, log_level=logging.ERROR):
+    def __init__(self, log_level=logging.DEBUG):
+        logging.basicConfig()
         self.log_level = log_level
         self.logger = logging.getLogger("BoilingData")
         self.logger.setLevel(self.log_level)
         self.bd_conn = BoilingDataConnection(log_level=self.log_level)
         self.conn = duckdb.connect(":memory:")
 
     async def _populate(self):
+        self.logger.debug("Creating local boilingdata data catalog")
         self.conn.execute("ATTACH ':memory:' AS boilingdata;")
         self.conn.execute("SET search_path='memory,boilingdata';")
         # Boiling specific table, contains data shares
         q = "SELECT * FROM information_schema.create_tables_statements"
         tables = await self.execute(q, None, True)
         if tables:
             for table in tables:
+                self.logger.debug(f"Creating table {table}")
                 self.conn.execute(table)
 
     def _is_boiling_execute(self, sql):
@@ -98,6 +101,7 @@ async def close(self):
 
     async def execute(self, sql, cb=None, force_boiling=False):
         """Send SQL Query to Boiling or run locally"""
+        sql = sql.replace("\n", " ")
         if not force_boiling and not self._is_boiling_execute(sql):
             return self.conn.execute(sql).fetchall()
         fut = await self.bd_conn.bd_execute(sql, cb)
@@ -199,7 +203,10 @@ def _on_msg(self, ws_app, data):
         if not reqId:
             return
         msg_type = msg.get("messageType")
-        # TODO: Store statistics sent from Boiling (INFO messages)
+        if msg_type == "LOG_MESSAGE":
+            log_level = msg.get("logLevel")
+            if log_level == "ERROR":
+                raise Exception(msg.get("logMessage"))
         if msg_type != "DATA":
             return
         req = self.requests.get(reqId)
@@ -252,6 +259,7 @@ async def connect(self):
         """Connect to BoilingData WebSocket API"""
         if self.websocket is not None:
             raise Exception("WebSocket already exists")
+        self.logger.info("Connecting")
         self.websocket = websocket.WebSocket()
         websocket.enableTrace(self.ws_trace)
         auth_headers = self._get_auth_headers()

py_boilingdata/data_queue.py

@@ -1,6 +1,3 @@
-from pprint import pprint
-
-
 class DataQueue:
     def __init__(self, requestId, callback, fut=None):
         self.requestId = requestId
@@ -21,7 +18,6 @@ def _have_all_parts(self):
         if self.parts_done == True:
             return True
         msg = self.data[0]
-        # print(msg)
         totalBatches = msg.get("totalBatches")
         totalSubBatches = msg.get("totalSubBatches")
         totalSplitSerials = msg.get("totalSplitSerials")
@@ -72,7 +68,6 @@ def _have_all_parts(self):
                         # print(f"\tsplits: {len(value.get('splits'))}/{totalSplits}")
                         if len(value.get("splits")) < totalSplits:
                             return False
-        # pprint(self.batchCounters)
         self.parts_done = True
         return True
 
@@ -99,10 +94,8 @@ def _notify(self):
             raise Exception("Deleted queue!")
         data = self._compile()
         if self.callback:
-            # print(f"CALLING CALLBACK {data}")
             self.callback({"data": data, "requestId": self.requestId})
         if self.fut:
-            # print(f"SETTING FUT RESULT: {data}")
             self.fut.set_result(data)
 
     ##

pyproject.toml

@@ -3,7 +3,7 @@ asyncio_mode = "auto"
 
 [tool.poetry]
 name = "py-boilingdata"
-version = "0.2.3"
+version = "0.2.4"
 description = "BoilingData (websockets) client for Python"
 authors = ["Dan Forsberg <dan@boilingdata.com>"]
 readme = "README.md"

tests/test_connection.py

@@ -39,8 +39,16 @@ async def test_local_tables():
 
 
 @pytest.mark.asyncio
-async def test_bd_query_cb():
+async def test_bd_query():
     """Small response from Boiling"""
+    q = "SELECT first_name, email FROM parquet_scan('s3://boilingdata-demo/test.parquet') LIMIT 1"
+    data = await boiling.execute(q)  # awaits as lont as the result is available
+    assert data == [{"email": "ajordan0@com.com", "first_name": "Amanda"}]
+
+
+@pytest.mark.asyncio
+async def test_bd_query_cb():
+    """Small response from Boiling with callback"""
     global data
     data = None
 
@@ -49,17 +57,11 @@ def cb(resp):
         global data
         data = resp
 
-    q = "SELECT first_name, email FROM parquet_scan('s3://boilingdata-demo/test.parquet') LIMIT 1"
+    q = "SELECT email FROM parquet_scan('s3://boilingdata-demo/test.parquet') LIMIT 2"
     await boiling.execute(q, cb)  # only awaits as long as the request is dispatched
-    await asyncio.sleep(5)
-    assert data == [{"email": "ajordan0@com.com", "first_name": "Amanda"}]
-
-
-@pytest.mark.asyncio
-async def test_bd_query():
-    """Small response from Boiling"""
     q = "SELECT first_name, email FROM parquet_scan('s3://boilingdata-demo/test.parquet') LIMIT 1"
-    data = await boiling.execute(q)  # awaits as lont as the result is available
+    await boiling.execute(q, cb)  # only awaits as long as the request is dispatched
+    await asyncio.sleep(5)  # let them race...
     assert data == [{"email": "ajordan0@com.com", "first_name": "Amanda"}]