Skip to content

Commit

Permalink
Add partition_key to create query
Browse files Browse the repository at this point in the history
  • Loading branch information
hellais committed Sep 11, 2024
1 parent d199c93 commit 3e162cb
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 2 deletions.
6 changes: 5 additions & 1 deletion oonidata/src/oonidata/models/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Check warning on line 23 in oonidata/src/oonidata/models/base.py

View check run for this annotation

Codecov / codecov/patch

oonidata/src/oonidata/models/base.py#L23

Added line #L23 was not covered by tests
return cls

return decorator
Expand All @@ -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
3 changes: 3 additions & 0 deletions oonidata/src/oonidata/models/observations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
4 changes: 3 additions & 1 deletion oonipipeline/src/oonipipeline/db/create_tables.py
Original file line number Diff line number Diff line change
Expand Up @@ -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} (
Expand Down

0 comments on commit 3e162cb

Please sign in to comment.