Fix airflow_db_cleanup script (#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!!!")