feat(pipeline)!: explicit BigQuery enable kwarg (#122)

Without this kwarg, a user is one command away from burning their daily
querying allowance (true story, happened yesterday). So, let us provide
an explicit knob to choose whether to enable or disable BigQuery without
any default. Why no default? Well, a true default would defeat the diff
purpose. A false default would hide that BigQuery does not run unless
you really want to, which could become confusing. An explicit knob,
instead, forces the developer to think about whether they really do want
to use BigQuery in a specific point in the codebase, and also allows to
carry over this decision directly from CLI flags (where it could default
to false).

I am using a kwarg for get_cache_entry instead of using one for the
whole pipeline because that enables creating a pipeline, doing a dry
run, then prompting the user whether to perform a real run all without
having to construct a new pipeline.

BREAKING CHANGE: `IQBPipeline.get_cache_entry` now requires an explicit
`enable_bigquery=True|False` kwarg.

Author: bassosimone

library/src/iqb/pipeline/pipeline.py

@@ -88,6 +88,7 @@ def get_cache_entry(
         self,
         *,
         dataset_name: str,
+        enable_bigquery: bool,
         start_date: str,
         end_date: str,
     ) -> PipelineCacheEntry:
@@ -99,6 +100,7 @@ def get_cache_entry(
 
         Args:
             dataset_name: Name for the dataset (e.g., "downloads_by_country")
+            enable_bigquery: Whether to enabled querying from BigQuery.
             start_date: Date when to start the query (included) -- format YYYY-MM-DD
             end_date: Date when to end the query (excluded) -- format YYYY-MM-DD
 
@@ -113,7 +115,8 @@ def get_cache_entry(
         )
 
         # 2. prepare for synching from BigQuery
-        entry.syncers.append(self._bq_syncer)
+        if enable_bigquery:
+            entry.syncers.append(self._bq_syncer)
 
         # 3. return the entry
         return entry

library/src/iqb/scripting/iqb_pipeline.py

@@ -17,14 +17,20 @@
 
 @dataclass(frozen=True, kw_only=True)
 class Pipeline:
-    """Wrapper for IQBPipeline providing convenience methods for scripting."""
+    """
+    Wrapper for IQBPipeline providing convenience methods for scripting.
+
+    Attributes:
+        pipeline: The IQBPipeline to use.
+    """
 
     pipeline: IQBPipeline
 
     def sync_mlab(
         self,
         granularity: str,
         *,
+        enable_bigquery: bool,
         end_date: str,
         start_date: str,
     ) -> None:
@@ -36,6 +42,7 @@ def sync_mlab(
 
         Arguments:
             end_date: exclusive end date as a YYYY-MM-DD string.
+            enable_bigquery: whether to enable querying from BigQuery.
             granularity: geographical granularity to use.
             start_date: incluive start date as a YYYY-MM-DD string.
 
@@ -57,6 +64,7 @@ def sync_mlab(
                     granularity=granularity,
                     table=table,
                 ),
+                enable_bigquery=enable_bigquery,
                 start_date=start_date,
                 end_date=end_date,
             )

library/tests/iqb/pipeline/pipeline_test.py

@@ -177,6 +177,7 @@ def test_get_cache_entry_when_exists(self, mock_client, tmp_path):
         # Get cache entry (should not execute query)
         entry = pipeline.get_cache_entry(
             dataset_name="downloads_by_country",
+            enable_bigquery=False,
             start_date="2024-10-01",
             end_date="2024-11-01",
         )
@@ -233,6 +234,7 @@ def mock_save_stats():
         # Get cache entry (BigQuery syncer is automatically added)
         entry = pipeline.get_cache_entry(
             dataset_name="downloads_by_country",
+            enable_bigquery=True,
             start_date="2024-10-01",
             end_date="2024-11-01",
         )
@@ -257,6 +259,25 @@ def mock_save_stats():
         assert entry.data_parquet_file_path().exists()
         assert entry.stats_json_file_path().exists()
 
+    @patch("iqb.pipeline.pipeline.PipelineBQPQClient")
+    def test_get_cache_entry_without_bigquery_syncer(self, mock_client, tmp_path):
+        """Test that get_cache_entry does not add BigQuery syncer when disabled."""
+        pipeline = IQBPipeline(project="test-project", data_dir=tmp_path)
+
+        entry = pipeline.get_cache_entry(
+            dataset_name="downloads_by_country",
+            enable_bigquery=False,
+            start_date="2024-10-01",
+            end_date="2024-11-01",
+        )
+
+        assert entry.syncers == []
+
+        with entry.lock(), pytest.raises(FileNotFoundError, match="Cache entry not found"):
+            entry.sync()
+
+        mock_client.return_value.execute_query.assert_not_called()
+
     @patch("iqb.pipeline.pipeline.PipelineBQPQClient")
     def test_bq_syncer_skip_when_exists(self, mock_client, tmp_path):
         """Test that _bq_syncer skips BigQuery when cache files already exist."""
@@ -280,6 +301,7 @@ def test_bq_syncer_skip_when_exists(self, mock_client, tmp_path):
         # Get cache entry (BigQuery syncer is automatically added)
         entry = pipeline.get_cache_entry(
             dataset_name="downloads_by_country",
+            enable_bigquery=True,
             start_date="2024-10-01",
             end_date="2024-11-01",
         )
@@ -310,6 +332,7 @@ def test_bq_syncer_failure(self, mock_client, tmp_path):
         # Get cache entry
         entry = pipeline.get_cache_entry(
             dataset_name="downloads_by_country",
+            enable_bigquery=True,
             start_date="2024-10-01",
             end_date="2024-11-01",
         )
@@ -343,6 +366,7 @@ def test_get_cache_entry_validation_checks(self, mock_client, tmp_path):
         with pytest.raises(ValueError, match="Invalid dataset name"):
             pipeline.get_cache_entry(
                 dataset_name="invalid_dataset-name",
+                enable_bigquery=False,
                 start_date="2024-10-01",
                 end_date="2024-11-01",
             )
@@ -351,6 +375,7 @@ def test_get_cache_entry_validation_checks(self, mock_client, tmp_path):
         with pytest.raises(ValueError, match="start_date must be < end_date"):
             pipeline.get_cache_entry(
                 dataset_name="downloads_by_country",
+                enable_bigquery=False,
                 start_date="2024-11-01",
                 end_date="2024-10-01",
             )

library/tests/iqb/scripting/iqb_pipeline_test.py

@@ -38,7 +38,12 @@ def test_syncs_missing_entries(self) -> None:
         pipeline.get_cache_entry.side_effect = [entry_download, entry_upload]
 
         wrapper = iqb_pipeline.Pipeline(pipeline=pipeline)
-        wrapper.sync_mlab("country", start_date="2024-01-01", end_date="2024-02-01")
+        wrapper.sync_mlab(
+            "country",
+            enable_bigquery=True,
+            start_date="2024-01-01",
+            end_date="2024-02-01",
+        )
 
         assert entry_download.synced is True
         assert entry_upload.synced is True
@@ -47,11 +52,13 @@ def test_syncs_missing_entries(self) -> None:
         assert pipeline.get_cache_entry.call_args_list == [
             call(
                 dataset_name="downloads_by_country",
+                enable_bigquery=True,
                 start_date="2024-01-01",
                 end_date="2024-02-01",
             ),
             call(
                 dataset_name="uploads_by_country",
+                enable_bigquery=True,
                 start_date="2024-01-01",
                 end_date="2024-02-01",
             ),
@@ -64,7 +71,12 @@ def test_skips_existing_entries(self) -> None:
         pipeline.get_cache_entry.side_effect = [entry_download, entry_upload]
 
         wrapper = iqb_pipeline.Pipeline(pipeline=pipeline)
-        wrapper.sync_mlab("country", start_date="2024-01-01", end_date="2024-02-01")
+        wrapper.sync_mlab(
+            "country",
+            enable_bigquery=False,
+            start_date="2024-01-01",
+            end_date="2024-02-01",
+        )
 
         assert entry_download.synced is False
         assert entry_upload.synced is False
@@ -74,7 +86,12 @@ def test_invalid_granularity_raises(self) -> None:
         wrapper = iqb_pipeline.Pipeline(pipeline=pipeline)
 
         with pytest.raises(ValueError, match="invalid granularity value"):
-            wrapper.sync_mlab("nope", start_date="2024-01-01", end_date="2024-02-01")
+            wrapper.sync_mlab(
+                "nope",
+                enable_bigquery=False,
+                start_date="2024-01-01",
+                end_date="2024-02-01",
+            )
 
         pipeline.get_cache_entry.assert_not_called()