diff --git a/oonidata/src/oonidata/models/base.py b/oonidata/src/oonidata/models/base.py index 7356e896..bfb6d4b3 100644 --- a/oonidata/src/oonidata/models/base.py +++ b/oonidata/src/oonidata/models/base.py @@ -14,10 +14,13 @@ class Config(BaseConfig): code_generation_options = [TO_DICT_ADD_OMIT_NONE_FLAG] -def table_model(table_name: str, table_index: Tuple[str, ...]): +def table_model( + table_name: str, table_index: Tuple[str, ...], partition_key: Optional[str] = None +): def decorator(cls): cls.__table_name__ = table_name cls.__table_index__ = table_index + cls.__partition_key__ = partition_key return cls return decorator @@ -28,6 +31,7 @@ def decorator(cls): class TableModelProtocol(Protocol): __table_name__: str __table_index__: Tuple[str, ...] + __partition_key__: Optional[str] probe_meta: Any measurement_meta: Any diff --git a/oonidata/src/oonidata/models/observations.py b/oonidata/src/oonidata/models/observations.py index b5a3bac5..cce96c82 100644 --- a/oonidata/src/oonidata/models/observations.py +++ b/oonidata/src/oonidata/models/observations.py @@ -187,6 +187,7 @@ class TCPObservation: "measurement_uid", "observation_idx", ), + partition_key="concat(substring(bucket_date, 1, 4), substring(bucket_date, 6, 2))", ) @dataclass class WebControlObservation: @@ -231,6 +232,7 @@ class WebControlObservation: "measurement_uid", "observation_idx", ), + partition_key="concat(substring(bucket_date, 1, 4), substring(bucket_date, 6, 2))", ) @dataclass class WebObservation: @@ -351,6 +353,7 @@ class WebObservation: "measurement_uid", "observation_idx", ), + partition_key="concat(substring(bucket_date, 1, 4), substring(bucket_date, 6, 2))", ) @dataclass class HTTPMiddleboxObservation: diff --git a/oonipipeline/src/oonipipeline/db/create_tables.py b/oonipipeline/src/oonipipeline/db/create_tables.py index 3356bdcf..1a42d86a 100644 --- a/oonipipeline/src/oonipipeline/db/create_tables.py +++ b/oonipipeline/src/oonipipeline/db/create_tables.py @@ -152,7 +152,9 @@ def format_create_query( index_str = ",\n".join(model.__table_index__) extra_str = "" if extra: - extra_str = f"ORDER BY ({index_str}) SETTINGS index_granularity = 8192;" + if model.__partition_key__: + extra_str = f"PARTITION BY ({model.__partition_key__})\n" + extra_str += f"ORDER BY ({index_str}) SETTINGS index_granularity = 8192;" return ( f""" CREATE TABLE IF NOT EXISTS {table_name} (