Simpler API

Author: dforsber

README.md

@@ -14,22 +14,14 @@ A python client for sending SQL queries and receive results from BoilingData Web
 
 ```python
 import asyncio
-from pprint import pprint
 from py_boilingdata import BoilingData
 
-
 async def main():
     boiling = BoilingData()
     await boiling.connect()
-    await boiling.populate()
-
-    def cb(resp):
-        pprint(resp)
-
     q = "SELECT first_name, email FROM parquet_scan('s3://boilingdata-demo/test.parquet') LIMIT 10"
-    await boiling.execute(q, cb)
-    await asyncio.sleep(10) # There is no way to await for the cb call..
-    print("DONE.")
+    resp = await boiling.execute(q)
+    print(resp)
 
 loop = asyncio.get_event_loop()
 loop.run_until_complete(main())

main.py

@@ -1,20 +1,20 @@
 import asyncio
-from pprint import pprint
 from py_boilingdata import BoilingData
 
+boiling = BoilingData()
+
 
 async def main():
-    boiling = BoilingData()
     await boiling.connect()
-    await boiling.populate()
-
-    def cb(resp):
-        pprint(resp)
-
-    q = "SELECT first_name, email FROM parquet_scan('s3://boilingdata-demo/test.parquet') LIMIT 10"
-    await boiling.execute(q, cb)
-    await asyncio.sleep(10)
-    print("DONE.")
+    results = await boiling.execute(
+        """
+        SELECT first_name, email 
+          FROM parquet_scan('s3://boilingdata-demo/test.parquet') 
+         LIMIT 2
+        """
+    )
+    print(results)
+    await boiling.close()
 
 
 loop = asyncio.get_event_loop()

py_boilingdata/__init__.py

@@ -1,13 +1,12 @@
 """BoilingData Client"""
-import os, json, uuid, time
-import rel
-import duckdb
+import os, json, uuid
+import logging
 import threading
+import duckdb
 import websocket
 import boto3
 import asyncio
 import botocore.auth
-from pprint import pprint
 from botocore.exceptions import NoCredentialsError
 from botocore.awsrequest import AWSRequest
 from botocore.credentials import Credentials
@@ -39,22 +38,22 @@
 class BoilingData:
     """Run SQL with BoilingData and local DuckDB"""
 
-    def __init__(self):
-        self.bd_conn = BoilingDataConnection()
+    def __init__(self, log_level=logging.ERROR):
+        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):
+    async def _populate(self):
         self.conn.execute("ATTACH ':memory:' AS boilingdata;")
         self.conn.execute("SET search_path='memory,boilingdata';")
-        # Boiling specific "information_schema" table
+        # Boiling specific table, contains data shares
         q = "SELECT * FROM information_schema.create_tables_statements"
-
-        def cb(bd_tables):
-            if bd_tables:
-                for table in bd_tables:
-                    self.conn.execute(table)
-
-        await self.bd_conn.bd_execute(q, cb)
+        tables = await self.execute(q, None, True)
+        if tables:
+            for table in tables:
+                self.conn.execute(table)
 
     def _is_boiling_execute(self, sql):
         ## 1) Get all Boiling tables so we know what to intercept
@@ -91,22 +90,32 @@ def _is_boiling_execute(self, sql):
     async def connect(self):
         """Connect to BoilingData"""
         await self.bd_conn.connect()
+        await self._populate()  # get catalog entries
 
     async def close(self):
         """Close WebSocket connection to Boiling"""
         await self.bd_conn.close()
 
-    async def execute(self, sql, cb):
+    async def execute(self, sql, cb=None, force_boiling=False):
         """Send SQL Query to Boiling or run locally"""
-        if not self._is_boiling_execute(sql):
+        if not force_boiling and not self._is_boiling_execute(sql):
             return self.conn.execute(sql).fetchall()
-        return await self.bd_conn.bd_execute(sql, cb)
+        fut = await self.bd_conn.bd_execute(sql, cb)
+        if cb is not None:
+            return
+        # TODO: Get rid of this while loop?!
+        while not fut.done():
+            await asyncio.sleep(0.005)
+        return fut.result()
 
 
 class BoilingDataConnection:
     """Create authenticated WebSocket connection to BoilingData"""
 
-    def __init__(self, region=AWS_REGION):
+    def __init__(self, region=AWS_REGION, log_level=logging.ERROR):
+        self.log_level = log_level
+        self.logger = logging.getLogger("BoilingDataConnection")
+        self.logger.setLevel(self.log_level)
         self.region = region
         self.username = os.getenv("BD_USERNAME", "")
         self.password = os.getenv("BD_PASSWORD", "")
@@ -115,6 +124,7 @@ def __init__(self, region=AWS_REGION):
                 "Missing username (BD_USERNAME) and/or "
                 + "password (BD_PASSWORD) environment variable(s)"
             )
+        self.wsConnectTimeoutS = 10
         self.websocket = None
         self.aws_creds = None
         self.ws_app = None
@@ -150,10 +160,10 @@ def _get_cognito_tokens(self, username, password):
             )
             return response["AuthenticationResult"]
         except self.idp_client.exceptions.NotAuthorizedException as e:
-            print("The username or password is incorrect.")
+            self.logger.error("The username or password is incorrect.")
             raise e
         except NoCredentialsError as e:
-            print("Credentials not available.")
+            self.logger.error("Credentials not available.")
             raise e
 
     def _get_credentials(self):
@@ -175,20 +185,21 @@ def _get_credentials(self):
         return self.aws_creds
 
     async def _ws_send(self, msg):
-        # print(f"> {msg}")
+        self.logger.debug(f"> {msg}")
         return self.ws_app.send(msg)
 
     def _on_open(self, ws_app):
-        print("WS OPEN")
+        self.logger.info("WS OPEN")
         self.bd_is_open = True
 
     def _on_msg(self, ws_app, data):
-        # print(f"< {data}")
+        self.logger.debug(f"< {data}")
         msg = json.loads(data)
         reqId = msg.get("requestId")
         if not reqId:
             return
         msg_type = msg.get("messageType")
+        # TODO: Store statistics sent from Boiling (INFO messages)
         if msg_type != "DATA":
             return
         req = self.requests.get(reqId)
@@ -201,32 +212,41 @@ def _on_msg(self, ws_app, data):
             del self.requests[reqId]
 
     def _on_error(self, ws_app, error):
-        print(f"WS ERROR: {error}")
+        self.logger.error(f"WS ERROR: {error}")
 
     def _on_close(self, ws_app, code, msg):
-        print(f"WS CLOSE: {code} {msg}")
-
-    ##
-    ## public
-    ##
+        self.logger.info(f"WS CLOSE: {code} {msg}")
+        self.is_open = False
 
     def _all_messages_received(self, event):
         requestId = event["requestId"]
         data = event["data"]
         cb = self.requests.get(requestId)
-        cb["callback"](data)
+        cb.get("callback")(data) if cb.get("callback") else None
+        cb.get("fut").set_result(data) if cb.get("fut") else None
+
+    ##
+    ## public
+    ##
 
     async def bd_execute(self, sql, cb):
-        if self.bd_is_open is not True:
+        if not self.bd_is_open:
+            await self.connect()
+        if not self.bd_is_open:
             raise Exception("No Boiling connection")
         reqId = uuid.uuid4().hex
         body = '{"sql":"' + sql + '","requestId":"' + reqId + '"}'
+        loop = asyncio.get_running_loop()
+        fut = loop.create_future()
         self.requests[reqId] = {
-            "q": DataQueue(reqId, self._all_messages_received),
+            "q": DataQueue(reqId, self._all_messages_received, fut),
+            "sql": sql,
             "reqId": reqId,
             "callback": cb,
+            "future": fut,
         }
         await self._ws_send(body)
+        return fut
 
     async def connect(self):
         """Connect to BoilingData WebSocket API"""
@@ -247,7 +267,7 @@ async def connect(self):
         wst.daemon = True
         wst.start()
         timeoutS = 1
-        while self.bd_is_open is not True and timeoutS < 10:
+        while not self.bd_is_open and timeoutS < self.wsConnectTimeoutS:
             await asyncio.sleep(1)
             timeoutS = timeoutS + 1
 

py_boilingdata/data_queue.py

@@ -2,17 +2,18 @@
 
 
 class DataQueue:
-    def __init__(self, requestId, callback):
+    def __init__(self, requestId, callback, fut=None):
         self.requestId = requestId
         self.callback = callback
+        self.fut = fut
         self.data = []
         self.batchCounters = dict()
         self.parts_done = False
         self.have_new_messages = False
         self.is_deleted = False
         return
 
-    def have_all_parts(self):
+    def _have_all_parts(self):
         if self.is_deleted:
             raise Exception("Deleted queue!")
         if not self.have_new_messages:
@@ -75,29 +76,7 @@ def have_all_parts(self):
         self.parts_done = True
         return True
 
-    def push(self, msg):
-        if self.is_deleted:
-            raise Exception("Deleted queue!")
-        self.data.append(msg)
-        self.have_new_messages = True
-        if self.have_all_parts():
-            self.notify()
-        return
-
-    # Prefer the callback method
-    def get_data(self):
-        if self.is_deleted:
-            raise Exception("Deleted queue!")
-        if self.is_done():
-            return self.compile()
-        return None
-
-    def is_done(self):
-        if self.is_deleted:
-            raise Exception("Deleted queue!")
-        return self.parts_done == True or self.have_all_parts()
-
-    def compile(self):
+    def _compile(self):
         if self.is_deleted:
             raise Exception("Deleted queue!")
         data = []
@@ -115,12 +94,42 @@ def compile(self):
                                 data.extend(value3["data"])
         return data
 
-    def notify(self):
+    def _notify(self):
         if self.is_deleted:
             raise Exception("Deleted queue!")
-        data = self.compile()
+        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)
+
+    ##
+    ## public
+    ##
+
+    def push(self, msg):
+        if self.is_deleted:
+            raise Exception("Deleted queue!")
+        self.data.append(msg)
+        self.have_new_messages = True
+        if self._have_all_parts():
+            self._notify()  # callback and future
+        return
+
+    # Prefer the callback method
+    def get_data(self):
+        if self.is_deleted:
+            raise Exception("Deleted queue!")
+        if self.is_done():
+            return self._compile()
+        return None
+
+    def is_done(self):
+        if self.is_deleted:
+            raise Exception("Deleted queue!")
+        return self.parts_done == True or self._have_all_parts()
 
     def delete(self):
         self.is_deleted = True

pyproject.toml

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