Added support for providing a timeout option while setting up a connection (#85)

Author: aishikbh

examples/pinot_quickstart_timeout.py

@@ -0,0 +1,46 @@
+from pinotdb import connect
+import time
+import pytest
+import httpx
+
+## Start Pinot Quickstart Batch
+## docker run --name pinot-quickstart -p 2123:2123 -p 9000:9000 -p 8000:8000 \
+##    -d apachepinot/pinot:latest QuickStart -type hybrid
+
+def run_pinot_quickstart_timeout_example() -> None:
+
+    #Test 1 : Try without timeout. The request should succeed.
+
+    conn = connect(host="localhost", port=8000, path="/query/sql", scheme="http")
+    curs = conn.cursor()
+    sql = "SELECT * FROM airlineStats LIMIT 5"
+    print(f"Sending SQL to Pinot: {sql}")
+    curs.execute(sql)
+    conn.close()
+
+    #Test 2 : Try with timeout=None. The request should succeed.
+
+    conn = connect(host="localhost", port=8000, path="/query/sql", scheme="http", timeout=None)
+    curs = conn.cursor()
+    sql = "SELECT count(*) FROM airlineStats LIMIT 5"
+    print(f"Sending SQL to Pinot: {sql}")
+    curs.execute(sql)
+    conn.close()
+
+    #Test 3 : Try with a really small timeout. The query should raise an exception.
+
+    conn = connect(host="localhost", port=8000, path="/query/sql", scheme="http", timeout=0.001)
+    curs = conn.cursor()
+    sql = "SELECT AirlineID, sum(Cancelled) FROM airlineStats WHERE Year > 2010 GROUP BY AirlineID LIMIT 5"
+    print(f"Sending SQL to Pinot: {sql}")
+    with pytest.raises(httpx.ReadTimeout):
+        curs.execute(sql)
+    conn.close()
+
+
+def run_main():
+    run_pinot_quickstart_timeout_example()
+
+
+if __name__ == '__main__':
+    run_main()

pinotdb/db.py

@@ -179,7 +179,8 @@ def cursor(self):
         """Return a new Cursor Object using the connection."""
         if not self.session or self.session.is_closed:
             self.session = httpx.Client(
-                verify=self._kwargs.get('verify_ssl'))
+                verify=self._kwargs.get('verify_ssl'),
+                timeout=self._kwargs.get('timeout'))
 
         self._kwargs['session'] = self.session
         cursor = Cursor(*self._args, **self._kwargs)