Fix airflow_db_cleanup script (GoogleCloudPlatform#7593)

## Description Fixes b/223364117 The DAG is now able to perform cleanup operation on airflow databases with larger data in a single table. More context : http://docs/document/d/1eWydWjGXqqQ4lvACOMs5o7qQtnuuyo_0jzwWXJ3_YfY?resourcekey=0-TBb9afCO5cWD5hhW_sRZgQ Note: It's a good idea to open an issue first for discussion. ## Checklist - [x] I have followed [Sample Guidelines from AUTHORING_GUIDE.MD](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md) - [ ] README is updated to include [all relevant information](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md#readme-file) - [ ] **Tests** pass: `nox -s py-3.6` (see [Test Environment Setup](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md#test-environment-setup)) - [x] **Lint** pass: `nox -s lint` (see [Test Environment Setup](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md#test-environment-setup)) - [ ] These samples need a new **API enabled** in testing projects to pass (let us know which ones) - [ ] These samples need a new/updated **env vars** in testing projects set to pass (let us know which ones) - [x] Please **merge** this PR for me once it is approved. - [ ] This sample adds a new sample directory, and I updated the [CODEOWNERS file](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/.github/CODEOWNERS) with the codeowners for this sample

Author: khokharnikunj8

composer/airflow_1_samples/airflow_db_cleanup.py

@@ -186,7 +186,8 @@
  "age_check_column": ImportError.timestamp,
  "keep_last": False,
  "keep_last_filters": None,
- "keep_last_group_by": None
+ "keep_last_group_by": None,
+ "do_not_delete_by_dag_id": True
  })
 
 except Exception as e:
@@ -204,13 +205,15 @@
  "age_check_column": Task.date_done,
  "keep_last": False,
  "keep_last_filters": None,
- "keep_last_group_by": None
+ "keep_last_group_by": None,
+ "do_not_delete_by_dag_id": True
  }, {
  "airflow_db_model": TaskSet,
  "age_check_column": TaskSet.date_done,
  "keep_last": False,
  "keep_last_filters": None,
- "keep_last_group_by": None
+ "keep_last_group_by": None,
+ "do_not_delete_by_dag_id": True
  }))
 
  except Exception as e:
@@ -358,10 +361,15 @@ def cleanup_function(**context):
 
  if ENABLE_DELETE:
  logging.info("Performing Delete...")
- # using bulk delete
- query.delete(synchronize_session=False)
- session.commit()
- logging.info("Finished Performing Delete")
+ if context["params"].get("do_not_delete_by_dag_id"):
+ query.filter(age_check_column <= max_date).delete(synchronize_session=False)
+ session.commit()
+ else:
+ dags = session.query(airflow_db_model.dag_id).distinct()
+ list_dags = [str(list(dag)[0]) for dag in dags]
+ for dag in list_dags:
+ query.filter(age_check_column <= max_date).filter(airflow_db_model.dag_id == dag).delete(synchronize_session=False)
+ session.commit()
  else:
  logging.warn("You've opted to skip deleting the db entries. "
  "Set ENABLE_DELETE to True to delete entries!!!")