Skip to content

Commit 79f9a5c

Browse files
ahuang11ahuang11
authored
Add async source methods (#1257)
1 parent 78f7db7 commit 79f9a5c

4 files changed

Lines changed: 359 additions & 2 deletions

File tree

‎lumen/sources/base.py‎

Lines changed: 101 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import asyncio
34
import fnmatch
45
import hashlib
56
import json
@@ -25,15 +26,20 @@
2526
import param # type: ignore
2627
import requests
2728

29+
try:
30+
import aiohttp
31+
except ImportError:
32+
aiohttp = None
33+
2834
from panel.io.cache import _generate_hash
2935

3036
from ..base import MultiTypeComponent
3137
from ..filters.base import Filter
3238
from ..state import state
3339
from ..transforms.base import Filter as FilterTransform, Transform
3440
from ..transforms.sql import (
35-
SQLCount, SQLDistinct, SQLLimit, SQLMinMax, SQLSample, SQLSelectFrom,
36-
SQLTransform,
41+
SQLCount, SQLDistinct, SQLFilter, SQLLimit, SQLMinMax, SQLSample,
42+
SQLSelectFrom, SQLTransform,
3743
)
3844
from ..util import get_dataframe_schema, is_ref, merge_schemas
3945
from ..validation import ValidationError, match_suggestion_message
@@ -635,6 +641,24 @@ def get(self, table: str, **query) -> DataFrame:
635641
A DataFrame containing the queried table.
636642
"""
637643

644+
async def get_async(self, table: str, **query) -> DataFrame:
645+
"""
646+
Return a table asynchronously; optionally filtered by the given query.
647+
648+
Parameters
649+
----------
650+
table : str
651+
The name of the table to query
652+
query : dict
653+
A dictionary containing all the query parameters
654+
655+
Returns
656+
-------
657+
DataFrame
658+
A DataFrame containing the queried table.
659+
"""
660+
return await asyncio.to_thread(self.get, table, **query)
661+
638662
def __str__(self) -> str:
639663
return self.name
640664

@@ -668,6 +692,32 @@ def get(self, table: str, **query) -> pd.DataFrame:
668692
df = pd.DataFrame(r.json())
669693
return df
670694

695+
async def get_async(self, table: str, **query) -> pd.DataFrame:
696+
"""
697+
Return a table asynchronously; optionally filtered by the given query.
698+
699+
Parameters
700+
----------
701+
table : str
702+
The name of the table to query
703+
query : dict
704+
A dictionary containing all the query parameters
705+
706+
Returns
707+
-------
708+
DataFrame
709+
A pandas DataFrame containing the queried table.
710+
"""
711+
if aiohttp is None:
712+
return super().get_async(table, **query)
713+
714+
query = dict(table=table, **query)
715+
async with aiohttp.ClientSession() as session:
716+
async with session.get(self.url+'/data', params=query) as response:
717+
data = await response.json()
718+
df = pd.DataFrame(data)
719+
return df
720+
671721

672722
class InMemorySource(Source):
673723
"""
@@ -978,6 +1028,55 @@ def execute(self, sql_query: str, *args, **kwargs) -> pd.DataFrame:
9781028
"""
9791029
raise NotImplementedError
9801030

1031+
async def execute_async(self, sql_query: str, *args, **kwargs) -> pd.DataFrame:
1032+
"""
1033+
Executes a SQL query asynchronously and returns the result as a DataFrame.
1034+
1035+
This default implementation runs the synchronous execute() method in a thread
1036+
to avoid blocking the event loop. Subclasses can override this method
1037+
to provide truly asynchronous implementations.
1038+
1039+
Arguments
1040+
---------
1041+
sql_query : str
1042+
The SQL Query to execute
1043+
*args : list
1044+
Positional arguments to pass to the SQL query
1045+
**kwargs : dict
1046+
Keyword arguments to pass to the SQL query
1047+
1048+
Returns
1049+
-------
1050+
pd.DataFrame
1051+
The result as a pandas DataFrame
1052+
"""
1053+
return await asyncio.to_thread(self.execute, sql_query, *args, **kwargs)
1054+
1055+
async def get_async(self, table: str, **query) -> DataFrame:
1056+
"""
1057+
Return a table asynchronously; optionally filtered by the given query.
1058+
1059+
Parameters
1060+
----------
1061+
table : str
1062+
The name of the table to query
1063+
query : dict
1064+
A dictionary containing all the query parameters
1065+
1066+
Returns
1067+
-------
1068+
DataFrame
1069+
A DataFrame containing the queried table.
1070+
"""
1071+
sql_expr = self.get_sql_expr(table)
1072+
1073+
conditions = list(query.items())
1074+
if conditions:
1075+
sql_filter = SQLFilter(conditions=conditions)
1076+
sql_expr = sql_filter.apply(sql_expr)
1077+
1078+
return await self.execute_async(sql_expr)
1079+
9811080
@cached_schema
9821081
def get_schema(
9831082
self, table: str | None = None, limit: int | None = None, shuffle: bool = False

‎lumen/sources/bigquery.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import asyncio
12
import json
23
import threading
34

@@ -119,6 +120,28 @@ def _authorize(self) -> None:
119120
def execute(self, sql_query: str) -> pd.DataFrame:
120121
return self._sql_client.query_and_wait(sql_query).to_dataframe()
121122

123+
async def execute_async(self, sql_query: str) -> pd.DataFrame:
124+
"""
125+
Execute a BigQuery SQL query asynchronously and return the result as a DataFrame.
126+
127+
Parameters
128+
----------
129+
sql_query : str
130+
The SQL query to execute
131+
132+
Returns
133+
-------
134+
pd.DataFrame
135+
The query result as a pandas DataFrame
136+
"""
137+
job = self._sql_client.query(sql_query)
138+
139+
while not job.done():
140+
await asyncio.sleep(0.1)
141+
job.reload()
142+
143+
return job.to_dataframe()
144+
122145
def get_tables(self) -> list[str]:
123146
"""Get a list of available tables for the project.
124147
@@ -377,6 +400,33 @@ def get(self, table, **query):
377400
sql_expr = st.apply(sql_expr)
378401
return self.execute(sql_expr)
379402

403+
async def get_async(self, table, **query):
404+
"""
405+
Retrieve a table from BigQuery asynchronously with optional filtering.
406+
407+
Parameters
408+
----------
409+
table : str
410+
The table name to query
411+
**query : dict
412+
Query parameters and filters to apply
413+
414+
Returns
415+
-------
416+
pd.DataFrame
417+
The filtered table data as a pandas DataFrame
418+
"""
419+
query.pop("__dask", None)
420+
sql_expr = self.get_sql_expr(table)
421+
sql_transforms = query.pop("sql_transforms", [])
422+
conditions = list(query.items())
423+
filter_in_sql = bool(self.filter_in_sql)
424+
if filter_in_sql:
425+
sql_transforms = [SQLFilter(conditions=conditions)] + sql_transforms
426+
for st in sql_transforms:
427+
sql_expr = st.apply(sql_expr)
428+
return await self.execute_async(sql_expr)
429+
380430
def close(self):
381431
"""
382432
Close the BigQuery client connections, releasing associated resources.

‎lumen/sources/snowflake.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import asyncio
34
import contextlib
45
import datetime
56
import decimal
@@ -16,6 +17,7 @@
1617
from cryptography.hazmat.primitives.serialization import (
1718
Encoding, NoEncryption, PrivateFormat, load_pem_private_key,
1819
)
20+
from snowflake.connector.constants import QueryStatus
1921

2022
from ..transforms.sql import SQLFilter
2123
from .base import BaseSQLSource, cached, cached_schema
@@ -259,6 +261,42 @@ def execute(self, sql_query: str, *args, **kwargs):
259261
df = self._cursor.execute(sql_query, *args, **kwargs).fetch_pandas_all()
260262
return self._cast_to_supported_dtypes(df)
261263

264+
async def execute_async(self, sql_query: str, *args, **kwargs):
265+
"""
266+
Execute a Snowflake SQL query asynchronously and return the result as a DataFrame.
267+
268+
Parameters
269+
----------
270+
sql_query : str
271+
The SQL query to execute
272+
*args : tuple
273+
Positional arguments to pass to the query
274+
**kwargs : dict
275+
Keyword arguments to pass to the query
276+
277+
Returns
278+
-------
279+
pd.DataFrame
280+
The query result as a pandas DataFrame with supported dtypes
281+
"""
282+
self._cursor.execute_async(sql_query, *args, **kwargs)
283+
query_id = self._cursor.sfqid
284+
285+
while True:
286+
status = self._cursor.get_query_status(query_id)
287+
if status in (QueryStatus.SUCCESS, QueryStatus.FAILED_WITH_ERROR, QueryStatus.ABORTED, QueryStatus.FAILED_WITH_INCIDENT):
288+
break
289+
await asyncio.sleep(0.1) # Check every 100ms
290+
291+
if status != QueryStatus.SUCCESS:
292+
raise snowflake.connector.errors.ProgrammingError(
293+
f"Query failed with status: {status}")
294+
295+
self._cursor.get_results_from_sfqid(query_id)
296+
297+
df = self._cursor.fetch_pandas_all()
298+
return self._cast_to_supported_dtypes(df)
299+
262300
def get_tables(self) -> list[str]:
263301
# limited set of tables was provided
264302
if isinstance(self.tables, dict | list):
@@ -283,6 +321,32 @@ def get(self, table, **query):
283321
sql_expr = st.apply(sql_expr)
284322
return self.execute(sql_expr)
285323

324+
async def get_async(self, table, **query):
325+
"""
326+
Retrieve data from a table asynchronously using the same query logic as get().
327+
328+
Parameters
329+
----------
330+
table : str
331+
The name of the table to query
332+
**query : dict
333+
Query parameters and filters to apply
334+
335+
Returns
336+
-------
337+
pd.DataFrame
338+
The query result as a pandas DataFrame
339+
"""
340+
query.pop('__dask', None)
341+
sql_expr = self.get_sql_expr(table)
342+
sql_transforms = query.pop('sql_transforms', [])
343+
conditions = list(query.items())
344+
if self.filter_in_sql:
345+
sql_transforms = [SQLFilter(conditions=conditions)] + sql_transforms
346+
for st in sql_transforms:
347+
sql_expr = st.apply(sql_expr)
348+
return await self.execute_async(sql_expr)
349+
286350
@contextlib.contextmanager
287351
def _timeout_context(self, seconds=None):
288352
"""Context manager to apply a timeout for operations using Snowflake's

0 commit comments

Comments
 (0)