Airflow DB Cleanup updates (GoogleCloudPlatform#9124)

kubasieron · engelke · m-strzelczyk · web-flow · commit bd89fb724d94 · 2023-02-17T17:40:40.000+01:00
* Fix and adjust to Airflow 2.3 db cleanup DAG

* Remove cleanup Task and TaskSet tables

* removing duplicated airflow_1 file + some minor fixes

* Add analyze task to the DAG to properly calculate db size

---------

Co-authored-by: Charles Engelke &lt;engelke@google.com&gt;
Co-authored-by: Maciej Strzelczyk &lt;strzelczyk@google.com&gt;
diff --git a/composer/workflows/airflow_db_cleanup.py b/composer/workflows/airflow_db_cleanup.py
@@ -361,6 +361,18 @@ def cleanup_function(**context):
         session.close()
 
 
+def analyze_db():
+    session = settings.Session()
+    session.execute("ANALYZE")
+    session.commit()
+
+
+analyze_op = PythonOperator(
+    task_id="analyze_query",
+    python_callable=analyze_db,
+    provide_context=True,
+    dag=dag)
+
 for db_object in DATABASE_OBJECTS:
 
     cleanup_op = PythonOperator(
@@ -371,5 +383,6 @@ def cleanup_function(**context):
         dag=dag)
 
     print_configuration.set_downstream(cleanup_op)
+    cleanup_op.set_downstream(analyze_op)
 
 # [END composer_metadb_cleanup]
