Skip to content

Commit

Permalink
add profiling option for duckdb
Browse files Browse the repository at this point in the history
  • Loading branch information
nkshaw23 committed Oct 16, 2024
1 parent 02a884e commit 0d7152f
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion src/driutils/io/duckdb.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,14 +51,15 @@ class DuckDBS3Reader(DuckDBReader):
"""Concrete Implementation of a DuckDB reader for reading
data from an S3 endpoint"""

def __init__(self, auth_type: str, endpoint_url: Optional[str] = None, use_ssl: bool = True) -> None:
def __init__(self, auth_type: str, endpoint_url: Optional[str] = None, use_ssl: bool = True, profiling = False) -> None:
"""Initializes
Args:
auth_type: The type of authentication to request. May
be one of ["auto", "sts", "custom_endpoint"]
endpoint_url: Custom s3 endpoint
use_ssl: Flag for using ssl (https connections).
profiling: Profile all duckdb queries. False by default.
"""

super().__init__()
Expand All @@ -76,6 +77,9 @@ def __init__(self, auth_type: str, endpoint_url: Optional[str] = None, use_ssl:

self._authenticate(auth_type, endpoint_url, use_ssl)

if profiling:
self._connection.execute("SET enable_profiling = query_tree;")

def _authenticate(self, method: str, endpoint_url: Optional[str] = None, use_ssl: bool = True) -> None:
"""Handles authentication selection
Expand Down

0 comments on commit 0d7152f

Please sign in to comment.