data-factory-utils 2.0.0__tar.gz → 2.2.0__tar.gz

{data_factory_utils-2.0.0 → data_factory_utils-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: data-factory-utils
-Version: 2.0.0
+Version: 2.2.0
 Summary: Utility functions for interacting with data factories.
 Requires-Dist: boto3>=1.42.8
 Requires-Dist: boto3-stubs>=1.42.89

{data_factory_utils-2.0.0 → data_factory_utils-2.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "data-factory-utils"
-version = "2.0.0"
+version = "2.2.0"
 description = "Utility functions for interacting with data factories."
 readme = "README.md"
 requires-python = ">=3.12"

{data_factory_utils-2.0.0 → data_factory_utils-2.2.0}/src/data_factory_utils/athena.py

@@ -8,7 +8,11 @@ from collections.abc import Generator
 import boto3
 import polars as pl
 from mypy_boto3_athena.client import AthenaClient
-from mypy_boto3_athena.type_defs import GetQueryResultsOutputTypeDef, QueryExecutionContextTypeDef
+from mypy_boto3_athena.type_defs import (
+    GetQueryResultsOutputTypeDef,
+    QueryExecutionContextTypeDef,
+    QueryExecutionStatisticsTypeDef,
+)
 
 from data_factory_utils.environment import Environment
 from data_factory_utils.query import Query
@@ -76,10 +80,12 @@ class AthenaQuery:
         self.config = config
         self.has_succeeded = False
 
-    def run_and_await(self) -> None:
+    def run_and_await(self, *, stats: bool = False) -> QueryExecutionStatisticsTypeDef | None:
         """Run query, wait for it to complete."""
         self.execution_id = self.start_execution()
-        self.await_query_completion()
+        if stats:
+            return self.await_query_completion(stats=True)
+        return self.await_query_completion()
 
     def start_execution(self) -> str:
         """Begin an Athena query with standard settings. Return the query ID."""
@@ -117,10 +123,8 @@ class AthenaQuery:
         self.config.athena_client.stop_query_execution(QueryExecutionId=self.execution_id)
 
     def await_query_completion(
-        self,
-        max_iterations: int = 20,
-        max_interval_seconds: int = 120,
-    ) -> None:
+        self, max_iterations: int = 20, max_interval_seconds: int = 120, *, stats: bool = False
+    ) -> QueryExecutionStatisticsTypeDef | None:
         """Wait for a managed query with the given ID to finish and return the result."""
         if self.execution_id is None:
             msg = "No query id available. Make sure that the query has been startedbefore querying the result."
@@ -145,7 +149,9 @@ class AthenaQuery:
                 case "SUCCEEDED":
                     self.has_succeeded = True
                     logger.info("Query succeeded")
-                    return
+                    if stats:
+                        return response["QueryExecution"]["Statistics"]
+                    return None
                 case "FAILED" | "CANCELLED":
                     msg = f"Query execution failed: {execution_status}"
                     raise AthenaQueryError(msg)

{data_factory_utils-2.0.0 → data_factory_utils-2.2.0}/src/data_factory_utils/query.py

@@ -17,6 +17,7 @@ DISALLOWED_CHARACTERS = re.compile(r"[;\n\r]")
 MAX_LIST_LENGTH = 1000
 S3_PATH_FORMAT = re.compile(r"^s3://[a-z0-9.-]{3,63}/.+$")
 PERMITTED_IDENTIFIERS = ("NULL", "True", "False")
+PERMITTED_TYPE_LITERALS = ("TIMESTAMP",)
 
 Column = str
 OrderByClause = tuple[Column, Literal["ASC", "DESC"]]
@@ -229,7 +230,7 @@ def validate_qualified_names(full_name: str) -> str:
 
 def validate_expression(expression: str) -> str:
     """Ensure that an expression is valid."""
-    if expression.startswith("COALESCE"):
+    if expression.upper().startswith(("COALESCE", "CAST")):
         return expression
 
     match = EXPRESSION_IDENTIFIER.fullmatch(expression)
@@ -265,6 +266,9 @@ def validate_identifier(identifier: str | Subquery | datetime) -> str | Subquery
     if identifier in PERMITTED_IDENTIFIERS:
         return identifier
 
+    if isinstance(identifier, str) and identifier.upper().startswith(PERMITTED_TYPE_LITERALS):
+        return identifier
+
     try:
         # allow numeric values
         float(identifier)