Add exportContext.offload flag to CLOUD_SQL_EXPORT_VALIDATION. (apach…

…e#23614)
okirialbert · May 10, 2022 · 74557e4 · 74557e4
1 parent 1bd75dd
commit 74557e4
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/airflow/providers/google/cloud/example_dags/example_cloud_sql.py b/airflow/providers/google/cloud/example_dags/example_cloud_sql.py
@@ -123,7 +123,12 @@
 # [END howto_operator_cloudsql_patch_body]
 # [START howto_operator_cloudsql_export_body]
 export_body = {
-    "exportContext": {"fileType": "sql", "uri": EXPORT_URI, "sqlExportOptions": {"schemaOnly": False}}
+    "exportContext": {
+        "fileType": "sql",
+        "uri": EXPORT_URI,
+        "sqlExportOptions": {"schemaOnly": False},
+        "offload": True,
+    }
 }
 # [END howto_operator_cloudsql_export_body]
 # [START howto_operator_cloudsql_import_body]

diff --git a/airflow/providers/google/cloud/operators/cloud_sql.py b/airflow/providers/google/cloud/operators/cloud_sql.py
@@ -154,6 +154,7 @@
                 ],
             ),
             dict(name="csvExportOptions", type="dict", optional=True, fields=[dict(name="selectQuery")]),
+            dict(name="offload", optional=True),
         ],
     )
 ]

diff --git a/tests/providers/google/cloud/operators/test_cloud_sql.py b/tests/providers/google/cloud/operators/test_cloud_sql.py
@@ -123,6 +123,7 @@
         "databases": [],
         "sqlExportOptions": {"tables": ["table1", "table2"], "schemaOnly": False},
         "csvExportOptions": {"selectQuery": "SELECT * FROM TABLE"},
+        "offload": True,
     }
 }
 IMPORT_BODY = {
-Original file line number
+Diff line change
@@ Expand Up / @@ -154,6 +154,7 @@ @@
                     ],
                 ),
                 dict(name="csvExportOptions", type="dict", optional=True, fields=[dict(name="selectQuery")]),
+                dict(name="offload", optional=True),
             ],
         )
     ]
@@ Expand Down @@