diff --git a/src/driutils/io/duckdb.py b/src/driutils/io/duckdb.py index fffa818..48b422c 100644 --- a/src/driutils/io/duckdb.py +++ b/src/driutils/io/duckdb.py @@ -51,7 +51,7 @@ 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: @@ -59,6 +59,7 @@ def __init__(self, auth_type: str, endpoint_url: Optional[str] = None, use_ssl: 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__() @@ -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