Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions datastore/infrastructure/engines/bigquery/export.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,16 +127,18 @@ async def dump(backend: Any, resource_id: str, fmt: str) -> list[str]:
prefix = f"{_EXPORT_ROOT}/{resource_id}/{fmt}/{rev}/"
source = f"`{backend.config.BIGQUERY_PROJECT}.{backend.config.BIGQUERY_DATASET}.{resource_id}`"

columns = _dump_column_order(table.schema)

async def build_export_sql(uri: str) -> tuple[str, bytes | None]:
# Schema is already known from get_table — no dry run needed.
sql = _export_data_sql(
uri,
fmt,
_export_select_list(table.schema, fmt),
_export_select_list(columns, fmt),
source=source,
suffix=" ORDER BY `_id`",
)
header = _csv_header_bytes(table.schema, fmt) if fmt in ("csv", "gzip") else None
header = _csv_header_bytes(columns, fmt) if fmt in ("csv", "gzip") else None
return sql, header

return await _prepare_download(
Expand Down Expand Up @@ -611,6 +613,19 @@ def _export_data_sql(
)


def _dump_column_order(schema: Any) -> list[Any]:
"""Table columns in download order: `_updated_at` moved last.

BigQuery stores the system columns first (`_id`, `_updated_at`,
then user fields), but `datastore_search` projects `_updated_at`
after the user columns; a whole-table dump follows the same order.
"""
fields = list(schema)
return [f for f in fields if f.name != "_updated_at"] + [
f for f in fields if f.name == "_updated_at"
]


def _export_select_list(schema: Any, fmt: str) -> str:
"""SELECT column list for an export.

Expand Down
24 changes: 24 additions & 0 deletions tests/test_datastore_dump.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@
from datastore.core.constants import DUMP_PREFIX
from datastore.infrastructure.engines.bigquery import BigQueryBackend
from datastore.infrastructure.engines.bigquery.export import (
_csv_header_bytes,
_dump_column_order,
_export_select_list,
_is_export_too_large,
)
Expand Down Expand Up @@ -190,6 +192,28 @@ def test_build_export_select_iso_casts_timestamp_and_datetime() -> None:
assert "`delivery_day`" in select


def test_dump_column_order_moves_updated_at_last() -> None:
"""Physical order is `_id, _updated_at, <user>`; the dump puts
`_updated_at` last, matching `datastore_search`."""
schema = [
_bq_field("_id", "INT64"),
_bq_field("_updated_at", "TIMESTAMP"),
_bq_field("auction_id", "INT64"),
_bq_field("product_code", "STRING"),
]
columns = _dump_column_order(schema)
assert [f.name for f in columns] == ["_id", "auction_id", "product_code", "_updated_at"]
assert _csv_header_bytes(columns, "csv") == b"_id,auction_id,product_code,_updated_at\n"
assert _export_select_list(columns, fmt="csv").endswith(
"FORMAT_TIMESTAMP('%Y-%m-%dT%H:%M:%S', `_updated_at`, 'UTC') AS `_updated_at`"
)


def test_dump_column_order_without_updated_at_is_unchanged() -> None:
schema = [_bq_field("_id", "INT64"), _bq_field("name", "STRING")]
assert [f.name for f in _dump_column_order(schema)] == ["_id", "name"]


def test_build_export_select_parquet_returns_star() -> None:
"""Nothing to cast → `*`, no column list."""
schema = [_bq_field("id", "INT64"), _bq_field("name", "STRING")]
Expand Down
Loading