Suppress along all dimensions

dependencies/requirements.test.txt

@@ -4,3 +4,6 @@ ipython==8.23.0
 moto[server]==5.1.1
 pytest==8.2.0
 pytest-asyncio==0.24.0
+pytest-timeout==2.4.0
+pyyaml==6.0.2
+shandy-sqlfmt[jinjafmt]==0.27.0

dependencies/requirements.txt

@@ -1,5 +1,5 @@
 boto3==1.37.7
-duckdb==1.3.1
+duckdb==1.3.2
 jinja2==3.1.5
 pandas==2.2.3
 pydantic==2.10.6

ostrich_egg/config.py

@@ -4,7 +4,14 @@
 
 from typing import Any, Union, Annotated, List, Sequence, Optional, Literal
 
-from pydantic import BaseModel, Field, AliasChoices, TypeAdapter
+from pydantic import (
+    BaseModel,
+    Field,
+    AliasChoices,
+    TypeAdapter,
+    AfterValidator,
+    field_serializer,
+)
 from enum import StrEnum
 
 from ostrich_egg.utils import identifier, get_logger, DEFAULT_MASKING_VALUE
@@ -36,6 +43,23 @@ class ReplaceWithRedactedParameters(BaseModel):
             ),
         ]
     ] = None
+    non_summable_dimensions: Optional[
+        Annotated[
+            Union[List[str], None],
+            Field(
+                description="List of dimensions that are part of the dataset but will not ever actually be aggregated. For example, if you have unrelated indicators in a column or years when you won't ever sum the total number of incidences across time (preventing revelation through subtraction). Use with caution."
+            ),
+        ]
+    ] = None
+    first_order_only: Optional[
+        Annotated[
+            bool,
+            Field(
+                description="Whether to only redact cells to prevent latent revelation along a single axis of dimensions. Enabling reduces the total number of cells suppressed but creates a risk of transitive revelation across dimensions.",
+                default=False,
+            ),
+        ]
+    ] = False
 
 
 class ReplaceWithRedacted(Strategy):
@@ -58,6 +82,15 @@ class MarkRedactedParameters(BaseModel):
             ),
         ]
     ] = None
+    first_order_only: Optional[
+        Annotated[
+            bool,
+            Field(
+                description="Whether to only redact cells to prevent latent revelation along a single axis of dimensions. Enabling reduces the total number of cells suppressed but creates a risk of transitive revelation across dimensions.",
+                default=False,
+            ),
+        ]
+    ] = False
 
 
 class MarkRedacted(Strategy):
@@ -224,6 +257,37 @@ def should_include_in_initial_state(self, initial: bool = False):
         return initial_conditions_match or counter_condition_does_not_exclude
 
 
+class DimensionOrder(BaseModel):
+    dimension: str
+    direction: Literal["asc", "desc"] = "asc"
+
+    @property
+    def sql_expression(self) -> str:
+        """
+        For use in templates so you can just | join(', ') the list.
+        """
+        return f"{identifier(self.dimension)} {self.direction}"
+
+
+def validate_dimension_orders(
+    dimension_orders: List[Union[str, DimensionOrder]] | None,
+):
+    """
+    Partially for backwards compatibility but also ease-of-use, the user could just plausibly specify as list of dimensions and we'll assume they meant to order them in ascending order.
+    """
+
+    valid_dimension_orders = []
+    if not dimension_orders:
+        return valid_dimension_orders
+
+    for dimension_order in dimension_orders:
+        if isinstance(dimension_order, str):
+            valid_dimension_orders.append(DimensionOrder(dimension=dimension_order))
+        elif isinstance(dimension_order, DimensionOrder):
+            valid_dimension_orders.append(dimension_order)
+    return valid_dimension_orders
+
+
 class Dataset(BaseModel):
 
     name: Optional[
@@ -268,13 +332,26 @@ class Dataset(BaseModel):
     ] = None
     redaction_order_dimensions: Optional[
         Annotated[
-            Union[List[str], None],
+            Union[List[DimensionOrder], List[str], None],
+            AfterValidator(validate_dimension_orders),
             Field(
                 default_factory=list,
-                description="The dimensions to order the redaction by. If not specified, the dimensions will be ordered by the order they are specified in the dataset.",
+                description="The dimensions to order the redaction by. If not specified, the dimensions will be ordered by the order they are specified in the dataset ascending.",
             ),
         ]
-    ] = None
+    ] = Field(
+        default_factory=list,
+    )
+
+    @field_serializer("redaction_order_dimensions")
+    def serialize_redaction_order_dimensions(
+        self, dimension_orders: List[Union[str, DimensionOrder]] | None
+    ):
+        return (
+            None
+            if not dimension_orders
+            else validate_dimension_orders(dimension_orders)
+        )
 
 
 class DataSource(BaseModel):

ostrich_egg/connectors/base.py

@@ -3,6 +3,8 @@
 
 import duckdb
 
+from ostrich_egg.utils import should_redact_along_axis
+
 DEFAULT_TABLE_NAME = "dataset"
 DEFAULT_RESULT_NAME = "result"
 
@@ -28,6 +30,7 @@ class BaseConnector(AbstractBaseClass):
 
     def __init__(self, table_name=DEFAULT_TABLE_NAME, **kwargs):
         self.table_name = table_name
+        self.init_duckdb()
 
     def __exit__(self):
         self.duckdb_connection.close()
@@ -36,10 +39,19 @@ def __exit__(self):
     def duckdb_connection(self) -> duckdb.DuckDBPyConnection:
         return duckdb.connect()
 
+    @property
+    def db(self):
+        return self.duckdb_connection
+
     def init_duckdb(self):
+        try:
+            del self.duckdb_connection
+        except AttributeError:
+            pass
         for extension in self.extensions:
             self.duckdb_connection.install_extension(extension)
             self.duckdb_connection.load_extension(extension)
+        self.load_custom_functions()
 
     @abstractmethod
     def load_source_table(self, *args, **kwargs):
@@ -48,3 +60,15 @@ def load_source_table(self, *args, **kwargs):
         If this were postgres, you wouldn't actually need to create a table.
         """
         raise NotImplementedError("Connectors must implement a create table interface")
+
+    def load_custom_functions(self):
+        should_redact_function_exists_count = (
+            self.db.sql(
+                "from duckdb_functions() where function_name = 'should_redact_along_axis'"
+            )
+            .count("*")
+            .fetchone()[0]
+        )
+        if should_redact_function_exists_count == 1:
+            self.db.remove_function("should_redact_along_axis")
+        self.db.create_function("should_redact_along_axis", should_redact_along_axis)

ostrich_egg/connectors/file_system.py

@@ -15,7 +15,6 @@ def load_source_table(
         """
         table_name = table_name or self.table_name
         file_path = source_file or self.file_path
-        self.init_duckdb()
         self.duckdb_connection.execute(
             f"create or replace table {table_name} as (select * from '{file_path}')"
         )

ostrich_egg/connectors/s3.py

@@ -5,7 +5,7 @@
 import duckdb
 from jinja2 import Template
 
-from ostrich_egg.connectors.base import BaseConnector
+from ostrich_egg.connectors.base import BaseConnector, DEFAULT_TABLE_NAME
 
 # refer to https://duckdb.org/docs/configuration/secrets_manager
 DEFAULT_S3_SECRET_NAME = "__default_s3"
@@ -45,13 +45,14 @@ def __init__(
         use_ssl: bool = True,
         url_style: Literal["vhost", "path"] = "vhost",
         chain: str = None,
+        table_name: str = DEFAULT_TABLE_NAME,
         *args,
         **kwargs,
     ):
         """
         See https://duckdb.org/docs/extensions/httpfs/s3api
         """
-        super().__init__(*args, **kwargs)
+        self.table_name = table_name
         self.bucket = bucket
         self.key = key
         self.region = region
@@ -63,6 +64,7 @@ def __init__(
         self.endpoint = endpoint
         self.use_ssl = use_ssl
         self.url_style = url_style
+        self.init_duckdb()
 
     def __exit__(self):
         self.duckdb_connection.close()
@@ -151,7 +153,6 @@ def load_source_table(
         bucket = bucket or self.bucket
         key = source_file or self.key
         key = key_as_s3_uri(bucket=bucket, key=key)
-        self.init_duckdb()
         self.duckdb_connection.execute(
             f"""
             create or replace table "{table_name}" as (