Tweak a couple of bigquery functions

The-Academic-Observatory · Sep 25, 2024 · 38329f7 · 38329f7
1 parent d666bbf
commit 38329f7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/observatory_platform/google/bigquery.py b/observatory_platform/google/bigquery.py
@@ -146,9 +146,8 @@ def bq_table_exists(table_id: str, client: Optional[bigquery.Client] = None) ->
 
 
 def bq_select_table_shard_dates(
-    *,
     table_id: str,
-    end_date: Union[pendulum.DateTime, pendulum.Date],
+    end_date: Union[pendulum.DateTime, pendulum.Date, None] = None,
     limit: int = 1,
     client: Optional[bigquery.Client] = None,
 ) -> List[pendulum.Date]:
@@ -167,7 +166,7 @@ def bq_select_table_shard_dates(
     query = render_template(
         template_path,
         table_id=table_id,
-        end_date=end_date.strftime("%Y-%m-%d"),
+        end_date=end_date if end_date is None else end_date.strftime("%Y-%m-%d"),
         limit=limit,
     )
     if client is None:
@@ -182,7 +181,10 @@ def bq_select_table_shard_dates(
 
 
 def bq_select_latest_table(
-    *, table_id: str, end_date: Union[pendulum.DateTime, pendulum.Date], sharded: bool, client: bigquery.Client = None
+    table_id: str,
+    end_date: Union[pendulum.DateTime, pendulum.Date, None],
+    sharded: bool = True,
+    client: bigquery.Client = None,
 ):
     """Select the latest fully qualified BigQuery table identifier.
 
@@ -197,7 +199,7 @@ def bq_select_latest_table(
         if client is None:
             client = bigquery.Client()
         table_date = bq_select_table_shard_dates(
-            table_id=table_id,
+            table_id,
             end_date=end_date,
             client=client,
         )[0]

diff --git a/observatory_platform/sql/select_table_shard_dates.sql.jinja2 b/observatory_platform/sql/select_table_shard_dates.sql.jinja2
@@ -18,6 +18,10 @@ SELECT DISTINCT
   PARSE_DATE('%Y%m%d', _TABLE_SUFFIX) as suffix
 FROM
   `{{ table_id }}*`
-WHERE CHAR_LENGTH(_TABLE_SUFFIX) = 8 AND REGEXP_CONTAINS(_TABLE_SUFFIX, r"\d{8}") AND PARSE_DATE('%Y%m%d', _TABLE_SUFFIX) <= DATE('{{ end_date }}')
+WHERE CHAR_LENGTH(_TABLE_SUFFIX) = 8
+    AND REGEXP_CONTAINS(_TABLE_SUFFIX, r"\d{8}")
+    {% if end_date is not none %}
+    AND PARSE_DATE('%Y%m%d', _TABLE_SUFFIX) <= DATE('{{ end_date }}')
+    {% endif %}
 ORDER BY suffix DESC
 LIMIT {{ limit }};