add configurable write_mode to BigQueryIOManager

docs/docs/integrations/libraries/gcp/bigquery/reference.md

@@ -45,6 +45,36 @@ In this example, we only use the columns containing sepal data from the `IRIS_DA
 
 When Dagster materializes `sepal_data` and loads the `iris_data` asset using the BigQuery I/O manager, it will only fetch the `sepal_length_cm` and `sepal_width_cm` columns of the `IRIS.IRIS_DATA` table and pass them to `sepal_data` as a Pandas DataFrame.
 
+## Configuring write modes
+
+By default, the BigQuery I/O manager performs a `TRUNCATE` operation when writing to an existing table. This deletes all rows but preserves the table schema. If the schema of your DataFrame changes (e.g., adding a new column), the insertion will fail.
+
+You can configure the `write_mode` to handle these scenarios:
+
+- `"truncate"` (default): Truncates the table before inserting data. Schema is preserved.
+- `"replace"`: Drops the table and recreates it. This allows schema evolution.
+- `"append"`: Inserts data without deleting existing rows.
+
+:::note
+
+For partitioned assets, the write mode is ignored, and the I/O manager will replace data in the targeted partitions.
+
+:::
+
+```python
+from dagster_gcp import BigQueryIOManager
+from dagster import Definitions, EnvVar
+
+defs = Definitions(
+    resources={
+        "io_manager": BigQueryIOManager(
+            project=EnvVar("GCP_PROJECT"),
+            write_mode="replace"
+        )
+    }
+)
+```
+
 ## Storing partitioned assets
 
 The BigQuery I/O manager supports storing and loading partitioned data. In order to correctly store and load data from the BigQuery table, the BigQuery I/O manager needs to know which column contains the data defining the partition bounds. The BigQuery I/O manager uses this information to construct the correct queries to select or replace the data. In the following sections, we describe how the I/O manager constructs these queries for different types of partitions.

docs/docs/integrations/libraries/gcp/bigquery/using-bigquery-with-dagster.md

@@ -68,6 +68,29 @@ You can also specify a `location` where computation should take place.
   endBefore="end_example"
 />
 
+#### Configuring write modes
+
+By default, the BigQuery I/O manager truncates data when writing to an existing table. You can change this behavior using the `write_mode` configuration option:
+
+- `truncate` (default): Deletes all rows in the table but keeps the schema.
+- `replace`: Drops the table and creates a new one. Useful when the schema changes.
+- `append`: Appends data to the existing table without deleting rows.
+
+```python
+from dagster_gcp import BigQueryIOManager
+from dagster import Definitions, EnvVar
+
+defs = Definitions(
+    resources={
+        "io_manager": BigQueryIOManager(
+            project=EnvVar("GCP_PROJECT"),
+            dataset="my_dataset",
+            write_mode="replace"  # Change write mode here
+        )
+    }
+)
+```
+
 ### Step 2: Create tables in BigQuery
 
 <Tabs>

python_modules/libraries/dagster-gcp/dagster_gcp/bigquery/io_manager.py

@@ -1,6 +1,7 @@
 from abc import abstractmethod
-from collections.abc import Generator, Sequence
+from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
+from enum import Enum
 from typing import cast
 
 from dagster import IOManagerDefinition, OutputContext, io_manager
@@ -23,6 +24,12 @@
 BIGQUERY_DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
 
 
+class BigQueryWriteMode(str, Enum):
+    TRUNCATE = "truncate"
+    APPEND = "append"
+    REPLACE = "replace"
+
+
 def build_bigquery_io_manager(
     type_handlers: Sequence[DbTypeHandler], default_load_type: type | None = None
 ) -> IOManagerDefinition:
@@ -136,7 +143,7 @@ def my_table_a(my_table: pd.DataFrame) -> pd.DataFrame:
     """
 
     @dagster_maintained_io_manager
-    @io_manager(config_schema=BigQueryIOManager.to_config_schema())  # pyright: ignore[reportArgumentType]
+    @io_manager(config_schema=BigQueryIOManager.to_config_schema())  # type: ignore
     def bigquery_io_manager(init_context):
         """I/O Manager for storing outputs in a BigQuery database.
 
@@ -151,16 +158,26 @@ def bigquery_io_manager(init_context):
         * dataset -> schema
         * table -> table
         """
+        resource_config = init_context.resource_config or {}
+        write_mode = resource_config.get("write_mode") or "truncate"
+        project = resource_config.get("project")
+        dataset = resource_config.get("dataset")
+
+        if project is None:
+            raise ValueError("Missing 'project' in configuration")
+
         mgr = DbIOManager(
             type_handlers=type_handlers,
-            db_client=BigQueryClient(),
+            db_client=BigQueryClient(write_mode=write_mode),
             io_manager_name="BigQueryIOManager",
-            database=init_context.resource_config["project"],
-            schema=init_context.resource_config.get("dataset"),
+            database=str(project),
+            schema=str(dataset) if dataset else None,
             default_load_type=default_load_type,
         )
-        if init_context.resource_config.get("gcp_credentials"):
-            with setup_gcp_creds(init_context.resource_config.get("gcp_credentials")):
+
+        gcp_creds = resource_config.get("gcp_credentials")
+        if gcp_creds:
+            with setup_gcp_creds(gcp_creds):
                 yield mgr
         else:
             yield mgr
@@ -258,6 +275,20 @@ def my_table_a(my_table: pd.DataFrame) -> pd.DataFrame:
         After the run completes, the file will be deleted, and ``GOOGLE_APPLICATION_CREDENTIALS`` will be
         unset. The key must be base64 encoded to avoid issues with newlines in the keys. You can retrieve
         the base64 encoded with this shell command: ``cat $GOOGLE_APPLICATION_CREDENTIALS | base64``
+    To change the write mode (default is "truncate"), you can set the ``write_mode`` configuration.
+        Supported modes: "truncate", "replace", "append".
+
+        .. code-block:: python
+
+            defs = Definitions(
+                assets=[my_table],
+                resources={
+                    "io_manager": BigQueryIOManager(
+                        project=EnvVar("GCP_PROJECT"),
+                        write_mode="replace"
+                    )
+                }
+            )
     """
 
     project: str = Field(description="The GCP project to use.")
@@ -300,6 +331,10 @@ def my_table_a(my_table: pd.DataFrame) -> pd.DataFrame:
             " queries (loading and reading from tables)."
         ),
     )
+    write_mode: BigQueryWriteMode = Field(
+        default=BigQueryWriteMode.TRUNCATE,
+        description="Write mode to use for non-partitioned table cleanup: truncate, append, or replace.",
+    )
 
     @staticmethod
     @abstractmethod
@@ -309,15 +344,21 @@ def type_handlers() -> Sequence[DbTypeHandler]: ...
     def default_load_type() -> type | None:
         return None
 
-    def create_io_manager(self, context) -> Generator:
-        mgr = DbIOManager(
-            db_client=BigQueryClient(),
+    def create_io_manager(self, context) -> DbIOManager:
+        return DbIOManager(
+            db_client=BigQueryClient(
+                write_mode=self.write_mode, gcp_credentials=self.gcp_credentials
+            ),
             io_manager_name="BigQueryIOManager",
             database=self.project,
             schema=self.dataset,
             type_handlers=self.type_handlers(),
             default_load_type=self.default_load_type(),
         )
+
+    @contextmanager
+    def yield_for_execution(self, context) -> Iterator[DbIOManager]:
+        mgr = self.create_io_manager(context)
         if self.gcp_credentials:
             with setup_gcp_creds(self.gcp_credentials):
                 yield mgr
@@ -326,10 +367,39 @@ def create_io_manager(self, context) -> Generator:
 
 
 class BigQueryClient(DbClient):
+    def __init__(
+        self,
+        write_mode: BigQueryWriteMode | str = BigQueryWriteMode.TRUNCATE,
+        gcp_credentials: str | None = None,
+    ):
+        if isinstance(write_mode, str):
+            write_mode = BigQueryWriteMode(write_mode)
+
+        self.write_mode = write_mode
+        self.gcp_credentials = gcp_credentials
+
     @staticmethod
     def delete_table_slice(context: OutputContext, table_slice: TableSlice, connection) -> None:
         try:
-            connection.query(_get_cleanup_statement(table_slice)).result()
+            # If partitioned, keep existing behavior (delete matching partitions)
+            if table_slice.partition_dimensions:
+                connection.query(_get_cleanup_statement(table_slice)).result()
+                return
+
+            # Non-partitioned tables: behavior depends on configured write_mode
+            resource_config = getattr(context, "resource_config", {}) or {}
+            write_mode = resource_config.get("write_mode")
+            if write_mode == BigQueryWriteMode.TRUNCATE.value or write_mode is None:
+                connection.query(
+                    f"TRUNCATE TABLE `{table_slice.database}.{table_slice.schema}.{table_slice.table}`"
+                ).result()
+            elif write_mode == BigQueryWriteMode.APPEND.value:
+                # Do nothing; preserve existing data and append
+                return
+            elif write_mode == BigQueryWriteMode.REPLACE.value:
+                connection.query(
+                    f"DROP TABLE IF EXISTS `{table_slice.database}.{table_slice.schema}.{table_slice.table}`"
+                ).result()
         except NotFound:
             # table doesn't exist yet, so ignore the error
             pass
@@ -353,10 +423,15 @@ def ensure_schema_exists(context: OutputContext, table_slice: TableSlice, connec
 
     @staticmethod
     @contextmanager
-    def connect(context, _):  # pyright: ignore[reportIncompatibleMethodOverride]
+    def connect(context, table_slice):
+        config = context.resource_config or {}
+
+        project_val = config.get("project")
+        location_val = config.get("location")
+
         conn = bigquery.Client(
-            project=context.resource_config.get("project"),
-            location=context.resource_config.get("location"),
+            project=str(project_val) if project_val else None,
+            location=str(location_val) if location_val else None,
         )
 
         yield conn

python_modules/libraries/dagster-gcp/dagster_gcp_tests/bigquery_tests/test_bigquery_write_modes.py

@@ -0,0 +1,132 @@
+import base64
+import json
+from unittest.mock import MagicMock, patch
+
+from dagster import OutputContext, build_init_resource_context
+from dagster._core.storage.db_io_manager import TableSlice
+from dagster_gcp.bigquery.io_manager import BigQueryClient, BigQueryIOManager, BigQueryWriteMode
+
+
+def test_delete_table_slice_truncate():
+    """Test TRUNCATE mode: Should execute TRUNCATE TABLE statement."""
+    client = BigQueryClient()
+    mock_conn = MagicMock()
+    table_slice = TableSlice(database="my_project", schema="my_dataset", table="my_table")
+
+    context = MagicMock(spec=OutputContext)
+    context.resource_config = {"write_mode": "truncate"}
+
+    client.delete_table_slice(context, table_slice, mock_conn)
+
+    mock_conn.query.assert_called_once_with("TRUNCATE TABLE `my_project.my_dataset.my_table`")
+
+
+def test_delete_table_slice_replace():
+    """Test REPLACE mode: Should execute DROP TABLE statement."""
+    client = BigQueryClient()
+    mock_conn = MagicMock()
+    table_slice = TableSlice(database="my_project", schema="my_dataset", table="my_table")
+    context = MagicMock(spec=OutputContext)
+    context.resource_config = {"write_mode": "replace"}
+
+    client.delete_table_slice(context, table_slice, mock_conn)
+
+    mock_conn.query.assert_called_once_with("DROP TABLE IF EXISTS `my_project.my_dataset.my_table`")
+
+
+def test_delete_table_slice_append():
+    """Test APPEND mode: Should do NOTHING (no deletion)."""
+    client = BigQueryClient(write_mode=BigQueryWriteMode.APPEND)
+    mock_conn = MagicMock()
+    table_slice = TableSlice(database="my_project", schema="my_dataset", table="my_table")
+    context = MagicMock(spec=OutputContext)
+    context.resource_config = {"write_mode": "append"}
+
+    client.delete_table_slice(context, table_slice, mock_conn)
+
+    mock_conn.query.assert_not_called()
+
+
+def test_partitioned_table_ignores_write_mode():
+    """Test that partitioned tables ignore the write mode and use legacy cleanup logic."""
+    client = BigQueryClient(write_mode=BigQueryWriteMode.REPLACE)
+    mock_conn = MagicMock()
+
+    mock_partition = MagicMock()
+    mock_partition.partitions = ["some_value"]
+    mock_partition.partition_expr = "my_partition_col"
+
+    table_slice = TableSlice(
+        database="my_project",
+        schema="my_dataset",
+        table="my_table",
+        partition_dimensions=[mock_partition],
+    )
+    context = MagicMock(spec=OutputContext)
+    context.resource_config = {"write_mode": "replace"}
+
+    client.delete_table_slice(context, table_slice, mock_conn)
+
+    args, _ = mock_conn.query.call_args
+    query_str = args[0]
+
+    assert "DROP TABLE" not in query_str
+    assert "DELETE FROM" in query_str
+
+
+class TestBigQueryIOManager(BigQueryIOManager):
+    @staticmethod
+    def type_handlers():
+        return []
+
+
+@patch("dagster_gcp.bigquery.io_manager.DbIOManager")
+def test_default_write_mode_in_factory(MockDbIOManager):
+    """Test that the default write mode propagates correctly from config to client."""
+    context = build_init_resource_context(config={"project": "test-project"})
+
+    manager_factory = TestBigQueryIOManager(project="test-project")
+    with manager_factory.yield_for_execution(context) as _:
+        assert MockDbIOManager.called
+        _, kwargs = MockDbIOManager.call_args
+        client = kwargs.get("db_client")
+
+        assert client is not None
+        assert client.write_mode == BigQueryWriteMode.TRUNCATE
+
+
+@patch("dagster_gcp.bigquery.io_manager.DbIOManager")
+def test_explicit_write_mode_in_factory(MockDbIOManager):
+    """Test that explicit write mode propagates correctly."""
+    context = build_init_resource_context(
+        config={"project": "test-project", "write_mode": "append"}
+    )
+
+    manager_factory = TestBigQueryIOManager(
+        project="test-project", write_mode=BigQueryWriteMode.APPEND
+    )
+    with manager_factory.yield_for_execution(context) as _:
+        _, kwargs = MockDbIOManager.call_args
+        client = kwargs.get("db_client")
+
+        assert client is not None
+        assert client.write_mode == BigQueryWriteMode.APPEND
+
+
+@patch("dagster_gcp.bigquery.io_manager.DbIOManager")
+def test_gcp_credentials_propagation(MockDbIOManager):
+    """Test that gcp_credentials are correctly passed to the client.
+    This ensures PySpark/Pandas won't fail with 'keyfile must not be null'.
+    """
+    dummy_json = json.dumps({"type": "service_account", "project_id": "test"})
+    creds_value = base64.b64encode(dummy_json.encode("utf-8")).decode("utf-8")
+
+    context = build_init_resource_context(config={"project": "test-project"})
+
+    manager_factory = TestBigQueryIOManager(project="test-project", gcp_credentials=creds_value)
+    with manager_factory.yield_for_execution(context) as _:
+        _, kwargs = MockDbIOManager.call_args
+        client = kwargs.get("db_client")
+
+        assert client is not None
+        assert client.gcp_credentials == creds_value