add configurable write_mode to BigQueryIOManager

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

@@ -45,6 +45,32 @@ 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,7 +1,8 @@
 from abc import abstractmethod
 from collections.abc import Generator, Sequence
 from contextlib import contextmanager
-from typing import Optional, cast
+from enum import Enum
+from typing import Optional, Union, cast
 
 from dagster import IOManagerDefinition, OutputContext, io_manager
 from dagster._config.pythonic_config import ConfigurableIOManagerFactory
@@ -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: Optional[type] = None
 ) -> IOManagerDefinition:
@@ -153,7 +160,7 @@ def bigquery_io_manager(init_context):
         """
         mgr = DbIOManager(
             type_handlers=type_handlers,
-            db_client=BigQueryClient(),
+            db_client=BigQueryClient(write_mode=init_context.resource_config.get("write_mode")),
             io_manager_name="BigQueryIOManager",
             database=init_context.resource_config["project"],
             schema=init_context.resource_config.get("dataset"),
@@ -258,6 +265,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 +321,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
@@ -311,7 +336,7 @@ def default_load_type() -> Optional[type]:
 
     def create_io_manager(self, context) -> Generator:
         mgr = DbIOManager(
-            db_client=BigQueryClient(),
+            db_client=BigQueryClient(write_mode=self.write_mode),
             io_manager_name="BigQueryIOManager",
             database=self.project,
             schema=self.dataset,
@@ -326,10 +351,38 @@ def create_io_manager(self, context) -> Generator:
 
 
 class BigQueryClient(DbClient):
-    @staticmethod
-    def delete_table_slice(context: OutputContext, table_slice: TableSlice, connection) -> None:
+    def __init__(self, write_mode: Optional[Union[BigQueryWriteMode, str]] = BigQueryWriteMode.TRUNCATE):
+        # Coerce string inputs (from raw config) to the enum, fallback to TRUNCATE on invalid values
+        if write_mode is None:
+            write_mode = BigQueryWriteMode.TRUNCATE
+
+        if isinstance(write_mode, str):
+            try:
+                write_mode = BigQueryWriteMode(write_mode)
+            except ValueError:
+                write_mode = BigQueryWriteMode.TRUNCATE
+
+        self.write_mode = write_mode
+
+    def delete_table_slice(self, 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
+            if self.write_mode == BigQueryWriteMode.TRUNCATE:
+                connection.query(
+                    f"TRUNCATE TABLE `{table_slice.database}.{table_slice.schema}.{table_slice.table}`"
+                ).result()
+            elif self.write_mode == BigQueryWriteMode.APPEND:
+                # Do nothing; preserve existing data and append
+                return
+            elif self.write_mode == BigQueryWriteMode.REPLACE:
+                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

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

@@ -0,0 +1,107 @@
+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(write_mode=BigQueryWriteMode.TRUNCATE)
+    mock_conn = MagicMock()
+    table_slice = TableSlice(database="my_project", schema="my_dataset", table="my_table")
+    context = MagicMock(spec=OutputContext)
+
+    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(write_mode=BigQueryWriteMode.REPLACE)
+    mock_conn = MagicMock()
+    table_slice = TableSlice(database="my_project", schema="my_dataset", table="my_table")
+    context = MagicMock(spec=OutputContext)
+
+    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)
+
+    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)
+
+    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")
+    iterator = manager_factory.create_io_manager(context)
+    next(iterator)
+
+    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="append")
+    iterator = manager_factory.create_io_manager(context)
+    next(iterator)
+
+    _, kwargs = MockDbIOManager.call_args
+    client = kwargs.get("db_client")
+
+    assert client is not None
+    assert client.write_mode == BigQueryWriteMode.APPEND