Repository: dlt-hub/dlt
Issue: dlt deletes pipeline state table

Hello! 👋 

Ran into quite a nasty bug,

Running a pipeline with `replace_strategy = "staging-optimized"` would consistently nuke the state table in BigQuery. Naturally I have months worth of state, luckily BigQuery snapshots everything so I was able to manually rollback. I confirmed the behavior across multiple invocations. 



Changed files:
dlt/load/load.py, tests/load/test_dummy_client.py, tests/load/pipeline/utils.py, tests/load/pipeline/test_replace_disposition.py

diff --git a/dlt/load/load.py b/dlt/load/load.py
index 643d1bbf..5dc89b5d 100644
--- a/dlt/load/load.py
+++ b/dlt/load/load.py
@@ -176,20 +176,22 @@ class Load(Runnable[ThreadPool]):
         jobs_info: List[ParsedLoadJobFileName] = []
         new_job_files = self.load_storage.list_new_jobs(load_id)
         for job_file in new_job_files:
-            if not dispositions or self.get_load_table(schema, job_file)["write_disposition"] in dispositions:
+            if dispositions is None or self.get_load_table(schema, job_file)["write_disposition"] in dispositions:
                 jobs_info.append(LoadStorage.parse_job_file_name(job_file))
         return jobs_info
 
+    def get_completed_table_chain(self, load_id: str, schema: Schema, top_merged_table: TTableSchema, being_completed_job_id: str = None) -> List[TTableSchema]:
+        """Gets a table chain starting from the `top_merged_table` containing only tables with completed/failed jobs. None is returned if there's any job that is not completed
 
-    def get_completed_table_chain(self, load_id: str, schema: Schema, top_merged_table: TTableSchema, starting_job_id: str) -> List[TTableSchema]:
-        """Gets a table chain starting from the `top_merged_table` containing only tables with completed/failed jobs. None is returned if there's any job that is not completed"""
+           Optionally `being_completed_job_id` can be passed that is considered to be completed before job itself moves in storage
+        """
         # returns ordered list of tables from parent to child leaf tables
         table_chain: List[TTableSchema] = []
         # make sure all the jobs for the table chain is completed
         for table in get_child_tables(schema.tables, top_merged_table["name"]):
             table_jobs = self.load_storage.list_jobs_for_table(load_id, table["name"])
             # all jobs must be completed in order for merge to be created
-            if any(job.state not in ("failed_jobs", "completed_jobs") and job.job_file_info.job_id() != starting_job_id for job in table_jobs):
+            if any(job.state not in ("failed_jobs", "completed_jobs") and job.job_file_info.job_id() != being_completed_job_id for job in table_jobs):
                 return None
             # if there are no jobs for the table, skip it, unless the write disposition is replace, as we need to create and clear the child tables
             if not table_jobs and top_merged_table["write_disposition"] != "replace":
diff --git a/tests/load/pipeline/test_replace_disposition.py b/tests/load/pipeline/test_replace_disposition.py
index 545b4601..f65e3de8 100644
--- a/tests/load/pipeline/test_replace_disposition.py
+++ b/tests/load/pipeline/test_replace_disposition.py
@@ -1,3 +1,4 @@
+from typing import Dict
 import yaml
 import dlt, os, pytest
 from dlt.common.utils import uniq_id
@@ -16,6 +17,18 @@ def test_replace_disposition(destination_config: DestinationTestConfiguration, r
     os.environ['DATA_WRITER__FILE_MAX_ITEMS'] = "40"
     # use staging tables for replace
     os.environ['DESTINATION__REPLACE_STRATEGY'] = replace_strategy
+    # make duckdb to reuse database in working folder
+    os.environ["DESTINATION__DUCKDB__CREDENTIALS"] = "duckdb:///test_replace_disposition.duckdb"
+
+    # TODO: start storing _dlt_loads with right json content
+    increase_loads = lambda x: x if destination_config.destination == "filesystem" else x + 1
+    increase_state_loads = lambda info: len([job for job in info.load_packages[0].jobs["completed_jobs"] if job.job_file_info.table_name == "_dlt_pipeline_state" and job.job_file_info.file_format != "reference"])
+
+    # filesystem does not have versions and child tables
+    def norm_table_counts(counts: Dict[str, int], *child_tables: str) -> Dict[str, int]:
+        if destination_config.destination != "filesystem":
+            return counts
+        return {**{"_dlt_version": 0}, **{t:0 for t in child_tables}, **counts}
 
     dataset_name = "test_replace_strategies_ds" + uniq_id()
     pipeline = destination_config.setup_pipeline("test_replace_strategies", dataset_name=dataset_name, full_refresh=False)
@@ -49,18 +62,29 @@ def test_replace_disposition(destination_config: DestinationTestConfiguration, r
 
     # first run with offset 0
     info = pipeline.run(load_items, loader_file_format=destination_config.file_format)
+    assert_load_info(info)
+    # count state records that got extracted
+    state_records = increase_state_loads(info)
+    dlt_loads: int = increase_loads(0)
+    dlt_versions: int = increase_loads(0)
 
     # second run with higher offset so we can check the results
     offset = 1000
     info = pipeline.run(load_items, loader_file_format=destination_config.file_format)
     assert_load_info(info)
+    state_records += increase_state_loads(info)
+    dlt_loads = increase_loads(dlt_loads)
 
     # we should have all items loaded
-    table_counts = load_table_counts(pipeline, *[t["name"] for t in pipeline.default_schema.data_tables()])
-
-    assert table_counts["items"] == 120
-    assert table_counts["items__sub_items"] == 240
-    assert table_counts["items__sub_items__sub_sub_items"] == 120
+    table_counts = load_table_counts(pipeline, *pipeline.default_schema.tables.keys())
+    assert norm_table_counts(table_counts) == {
+        "items": 120,
+        "items__sub_items": 240,
+        "items__sub_items__sub_sub_items": 120,
+        "_dlt_pipeline_state": state_records,
+        "_dlt_loads": dlt_loads,
+        "_dlt_version": dlt_versions
+    }
 
     # check we really have the replaced data in our destination
     table_dicts = load_tables_to_dicts(pipeline, *[t["name"] for t in pipeline.default_schema.data_tables()])
@@ -68,27 +92,60 @@ def test_replace_disposition(destination_config: DestinationTestConfiguration, r
     assert {x for i,x in enumerate(range(2000, 2000+120), 1)}.union({x for i,x in enumerate(range(3000, 3000+120), 1)}) == {int(x["id"]) for x in table_dicts["items__sub_items"]}
     assert {x for i,x in enumerate(range(4000, 4120), 1)} == {int(x["id"]) for x in table_dicts["items__sub_items__sub_sub_items"]}
 
-
     # we need to test that destination tables including child tables are cleared when we yield none from the resource
     @dlt.resource(name="items", write_disposition="replace", primary_key="id")
     def load_items_none():
         yield
     info = pipeline.run(load_items_none, loader_file_format=destination_config.file_format)
-
-    table_counts = load_table_counts(pipeline, *[t["name"] for t in pipeline.default_schema.data_tables()])
+    assert_load_info(info)
+    state_records += increase_state_loads(info)
+    dlt_loads = increase_loads(dlt_loads)
 
     # table and child tables should be cleared
-    assert table_counts["items"] == 0
-    assert table_counts.get("items__sub_items", 0) == 0
-    assert table_counts.get("items__sub_items__sub_sub_items", 0) == 0
+    table_counts = load_table_counts(pipeline, *pipeline.default_schema.tables.keys())
+    assert norm_table_counts(table_counts, "items__sub_items", "items__sub_items__sub_sub_items") == {
+        "items": 0,
+        "items__sub_items": 0,
+        "items__sub_items__sub_sub_items": 0,
+        "_dlt_pipeline_state": state_records,
+        "_dlt_loads": dlt_loads,
+        "_dlt_version": dlt_versions
+    }
 
     # drop and deactivate existing pipeline
     # drop_active_pipeline_data()
 
     # create a pipeline with different name but loading to the same dataset as above - this is to provoke truncating non existing tables
-    pipeline = destination_config.setup_pipeline("test_replace_strategies_2", dataset_name=dataset_name, full_refresh=False)
-    info = pipeline.run(load_items, table_name="items_copy", loader_file_format=destination_config.file_format)
+    pipeline_2 = destination_config.setup_pipeline("test_replace_strategies_2", dataset_name=dataset_name, full_refresh=False)
+    info = pipeline_2.run(load_items, table_name="items_copy", loader_file_format=destination_config.file_format)
+    assert_load_info(info)
+
+    new_state_records = increase_state_loads(info)
+    assert new_state_records == 1
+    dlt_loads = increase_loads(dlt_loads)
+
+    # new pipeline
+    table_counts = load_table_counts(pipeline_2, *pipeline_2.default_schema.tables.keys())
+    assert norm_table_counts(table_counts) == {
+        "items_copy": 120,
+        "items_copy__sub_items": 240,
+        "items_copy__sub_items__sub_sub_items": 120,
+        "_dlt_pipeline_state": state_records + 1,
+        "_dlt_loads": dlt_loads,
+        "_dlt_version": increase_loads(dlt_versions)
+    }
+
 
+    # old pipeline -> shares completed loads and versions table
+    table_counts = load_table_counts(pipeline, *pipeline.default_schema.tables.keys())
+    assert norm_table_counts(table_counts, "items__sub_items", "items__sub_items__sub_sub_items") == {
+        "items": 0,
+        "items__sub_items": 0,
+        "items__sub_items__sub_sub_items": 0,
+        "_dlt_pipeline_state": state_records + 1,
+        "_dlt_loads": dlt_loads,  #  next load
+        "_dlt_version": increase_loads(dlt_versions)  # new table name -> new schema
+    }
 
 @pytest.mark.parametrize("destination_config", destinations_configs(local_filesystem_configs=True, default_staging_configs=True, default_non_staging_configs=True), ids=lambda x: x.name)
 @pytest.mark.parametrize("replace_strategy", REPLACE_STRATEGIES)
diff --git a/tests/load/pipeline/utils.py b/tests/load/pipeline/utils.py
index 2158a05c..8746b0d1 100644
--- a/tests/load/pipeline/utils.py
+++ b/tests/load/pipeline/utils.py
@@ -160,7 +160,7 @@ def drop_active_pipeline_data() -> None:
                 for schema_name in p.schema_names:
                     drop_func(schema_name)
 
-        p._wipe_working_folder()
+        # p._wipe_working_folder()
         # deactivate context
         Container()[PipelineContext].deactivate()
 
diff --git a/tests/load/test_dummy_client.py b/tests/load/test_dummy_client.py
index eaba5287..2e6e35fd 100644
--- a/tests/load/test_dummy_client.py
+++ b/tests/load/test_dummy_client.py
@@ -20,6 +20,7 @@ from dlt.destinations import dummy
 from dlt.destinations.dummy import dummy as dummy_impl
 from dlt.destinations.dummy.configuration import DummyClientConfiguration
 from dlt.load.exceptions import LoadClientJobFailed, LoadClientJobRetry
+from dlt.common.schema.utils import get_top_level_table
 
 from tests.utils import clean_test_storage, init_test_logging, TEST_DICT_CONFIG_PROVIDER, preserve_environ
 from tests.load.utils import prepare_load_package
@@ -89,6 +90,46 @@ def test_unsupported_write_disposition() -> None:
     assert "LoadClientUnsupportedWriteDisposition" in load.load_storage.storage.load(exception)
 
 
+def test_get_new_jobs_info() -> None:
+    load = setup_loader()
+    load_id, schema = prepare_load_package(
+        load.load_storage,
+        NORMALIZED_FILES
+    )
+
+    # no write disposition specified - get all new jobs
+    assert len(load.get_new_jobs_info(load_id, schema)) == 2
+    # empty list - none
+    assert len(load.get_new_jobs_info(load_id, schema, [])) == 0
+    # two appends
+    assert len(load.get_new_jobs_info(load_id, schema, ["append"])) == 2
+    assert len(load.get_new_jobs_info(load_id, schema, ["replace"])) == 0
+    assert len(load.get_new_jobs_info(load_id, schema, ["replace", "append"])) == 2
+
+    load.load_storage.start_job(load_id, "event_loop_interrupted.839c6e6b514e427687586ccc65bf133f.0.jsonl")
+    assert len(load.get_new_jobs_info(load_id, schema, ["replace", "append"])) == 1
+
+
+def test_get_completed_table_chain_single_job_per_table() -> None:
+    load = setup_loader()
+    load_id, schema = prepare_load_package(
+        load.load_storage,
+        NORMALIZED_FILES
+    )
+
+    top_job_table = get_top_level_table(schema.tables, "event_user")
+    assert load.get_completed_table_chain(load_id, schema, top_job_table) is None
+    # fake being completed
+    assert len(load.get_completed_table_chain(load_id, schema, top_job_table, "event_user.839c6e6b514e427687586ccc65bf133f.0.jsonl")) == 1
+    # actually complete
+    loop_top_job_table = get_top_level_table(schema.tables, "event_loop_interrupted")
+    load.load_storage.start_job(load_id, "event_loop_interrupted.839c6e6b514e427687586ccc65bf133f.0.jsonl")
+    assert load.get_completed_table_chain(load_id, schema, loop_top_job_table) is None
+    load.load_storage.complete_job(load_id, "event_loop_interrupted.839c6e6b514e427687586ccc65bf133f.0.jsonl")
+    assert load.get_completed_table_chain(load_id, schema, loop_top_job_table) == [schema.get_table("event_loop_interrupted")]
+    assert load.get_completed_table_chain(load_id, schema, loop_top_job_table, "event_user.839c6e6b514e427687586ccc65bf133f.0.jsonl") == [schema.get_table("event_loop_interrupted")]
+
+
 def test_spool_job_failed() -> None:
     # this config fails job on start
     load = setup_loader(client_config=DummyClientConfiguration(fail_prob=1.0))