data-factory-utils 3.2.2__tar.gz → 3.4.0__tar.gz

{data_factory_utils-3.2.2 → data_factory_utils-3.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: data-factory-utils
-Version: 3.2.2
+Version: 3.4.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-3.2.2 → data_factory_utils-3.4.0}/pyproject.toml

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

{data_factory_utils-3.2.2 → data_factory_utils-3.4.0}/pyproject.toml.orig

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

{data_factory_utils-3.2.2 → data_factory_utils-3.4.0}/src/data_factory_utils/athena.py

@@ -200,24 +200,23 @@ class AthenaQuery:
         # Get the raw query results from the API
         responses = self._get_responses()
 
-        # Parse the raw row data from the results
-        rows = []
+        column_names: list[str | None] | None = None
+        columns: list[list[str | None]] = []
+
         for response in responses:
             for row in response["ResultSet"]["Rows"]:
                 row_data = [column.get("VarCharValue") for column in row["Data"]]
-                rows.append(row_data)
-
-        if len(rows) == 0:
-            return pl.DataFrame({})
 
-        # The first row returned by the response is the names of the columns.
-        column_names, values_in_rows = rows[0], rows[1:]
+                # The first row returned by the response is the names of the columns.
+                if column_names is None:
+                    column_names = row_data
+                    columns = [[] for _ in column_names]
+                    continue
 
-        if len(values_in_rows) == 0:
-            return pl.DataFrame(data={col: [] for col in column_names})
+                for column, value in zip(columns, row_data, strict=True):
+                    column.append(value)
 
-        # Convert a list of rows into a list of columns
-        values_in_columns = list(zip(*values_in_rows, strict=True))
+        if column_names is None:
+            return pl.DataFrame({})
 
-        # Package it up into a dictionary of columns
-        return pl.DataFrame(dict(zip(column_names, values_in_columns, strict=True)))
+        return pl.DataFrame(dict(zip(column_names, columns, strict=True)))

{data_factory_utils-3.2.2 → data_factory_utils-3.4.0}/src/data_factory_utils/query.py

@@ -79,6 +79,23 @@ class WhereBasicComparison(NamedTuple):
     brackets: Literal["open", "close", "wrap"] | None = None
 
 
+class WhereExistsComparison(NamedTuple):
+    """Config for WHERE EXISTS comparisons."""
+
+    subquery: Subquery
+    negative: bool = False
+    prefix: Literal["AND", "OR", "NOT"] | None = "AND"
+    brackets: Literal["open", "close", "wrap"] | None = None
+
+
+class WhereNotExistsComparison(NamedTuple):
+    """Config for WHERE NOT EXISTS comparisons."""
+
+    subquery: Subquery
+    prefix: Literal["AND", "OR", "NOT"] | None = "AND"
+    brackets: Literal["open", "close", "wrap"] | None = None
+
+
 class WhereRangeComparison(NamedTuple):
     """Config for WHERE BETWEEN comparisons."""
 
@@ -120,6 +137,17 @@ class WhereLikeRegex(NamedTuple):
     brackets: Literal["open", "close", "wrap"] | None = None
 
 
+WhereCondition = (
+    WhereBasicComparison
+    | WhereExistsComparison
+    | WhereNotExistsComparison
+    | WhereRangeComparison
+    | WhereSetMembership
+    | WherePatternMatch
+    | WhereLikeRegex
+)
+
+
 class Subquery:
     """Method to wrap queries as subqueries."""
 
@@ -537,6 +565,30 @@ def _build_where_between(config: WhereRangeComparison) -> str:
     return query_str
 
 
+def _build_where_exists_comparison(config: WhereExistsComparison | WhereNotExistsComparison) -> str:
+    """Build a statement for EXISTS / NOT EXISTS.
+
+    Accepted styles:
+        - WHERE EXISTS (subquery)
+        - WHERE NOT EXISTS (subquery)
+        - AND EXISTS (subquery)
+        - AND NOT EXISTS (subquery)
+        - OR EXISTS (subquery)
+        - OR NOT EXISTS (subquery)
+    """
+    if not isinstance(config.subquery, Subquery):
+        msg = "EXISTS condition requires a Subquery instance"
+        raise TypeError(msg)
+
+    negative = getattr(config, "negative", False) or isinstance(config, WhereNotExistsComparison)
+    op = "NOT EXISTS" if negative else "EXISTS"
+    query_str = f"{op} {config.subquery}"
+
+    if config.brackets:
+        query_str = _wrap(loc=config.brackets, text=query_str)
+    return query_str
+
+
 def _build_where_set_membership(config: WhereSetMembership) -> str:
     """Build a statement for set membership.
 
@@ -607,9 +659,7 @@ def _build_where_regex_like(config: WhereLikeRegex) -> str:
 
 
 def _validate_combined_where_conditions(
-    conditions: list[
-        WhereBasicComparison | WhereRangeComparison | WhereSetMembership | WherePatternMatch | WhereLikeRegex
-    ],
+    conditions: list[WhereCondition],
 ) -> bool:
     """Validate WHERE conditions when used as a list."""
     open_bracket = []
@@ -803,9 +853,7 @@ class Query:
 
     def WHERE(  # noqa: N802
         self,
-        conditions: list[
-            WhereBasicComparison | WhereRangeComparison | WhereSetMembership | WherePatternMatch | WhereLikeRegex
-        ],
+        conditions: list[WhereCondition],
     ) -> Self:
         """Write WHERE conditions."""
         query_str = []
@@ -824,6 +872,8 @@ class Query:
                     cond_str = _build_where_pattern_match(condition)
                 case WhereLikeRegex():
                     cond_str = _build_where_regex_like(condition)
+                case WhereExistsComparison() | WhereNotExistsComparison():
+                    cond_str = _build_where_exists_comparison(condition)
                 case _:
                     msg = "WHERE statement has no valid conditions."
                     raise ValueError(msg)