Skip to content

Commit

Permalink
Added support for providing a timeout option while setting up a conne…
Browse files Browse the repository at this point in the history
…ction (#85)
  • Loading branch information
aishikbh authored Nov 16, 2023
1 parent 2ee148a commit 11a8b66
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 1 deletion.
46 changes: 46 additions & 0 deletions examples/pinot_quickstart_timeout.py
Original file line number Diff line number Diff line change
@@ -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()
3 changes: 2 additions & 1 deletion pinotdb/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 11a8b66

Please sign in to comment.