Sanitize filenames in MySQLHook (#33328)

apache · Jan 3, 2024 · 3598a52 · 3598a52
1 parent bd1d2f4
commit 3598a52
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 37 deletions.
diff --git a/airflow/providers/mysql/hooks/mysql.py b/airflow/providers/mysql/hooks/mysql.py
@@ -214,10 +214,8 @@ def bulk_load(self, table: str, tmp_file: str) -> None:
         conn = self.get_conn()
         cur = conn.cursor()
         cur.execute(
-            f"""
-            LOAD DATA LOCAL INFILE '{tmp_file}'
-            INTO TABLE {table}
-            """
+            f"LOAD DATA LOCAL INFILE %s INTO TABLE {table}",
+            (tmp_file,),
         )
         conn.commit()
         conn.close()  # type: ignore[misc]
@@ -227,10 +225,8 @@ def bulk_dump(self, table: str, tmp_file: str) -> None:
         conn = self.get_conn()
         cur = conn.cursor()
         cur.execute(
-            f"""
-            SELECT * INTO OUTFILE '{tmp_file}'
-            FROM {table}
-            """
+            f"SELECT * INTO OUTFILE %s FROM {table}",
+            (tmp_file,),
         )
         conn.commit()
         conn.close()  # type: ignore[misc]
@@ -294,12 +290,8 @@ def bulk_load_custom(
         cursor = conn.cursor()
 
         cursor.execute(
-            f"""
-            LOAD DATA LOCAL INFILE '{tmp_file}'
-            {duplicate_key_handling}
-            INTO TABLE {table}
-            {extra_options}
-            """
+            f"LOAD DATA LOCAL INFILE %s %s INTO TABLE {table} %s",
+            (tmp_file, duplicate_key_handling, extra_options),
         )
 
         cursor.close()

diff --git a/tests/providers/mysql/hooks/test_mysql.py b/tests/providers/mysql/hooks/test_mysql.py
@@ -281,21 +281,11 @@ def test_run_multi_queries(self):
 
     def test_bulk_load(self):
         self.db_hook.bulk_load("table", "/tmp/file")
-        self.cur.execute.assert_called_once_with(
-            """
-            LOAD DATA LOCAL INFILE '/tmp/file'
-            INTO TABLE table
-            """
-        )
+        self.cur.execute.assert_called_once_with("LOAD DATA LOCAL INFILE %s INTO TABLE table", ("/tmp/file",))
 
     def test_bulk_dump(self):
         self.db_hook.bulk_dump("table", "/tmp/file")
-        self.cur.execute.assert_called_once_with(
-            """
-            SELECT * INTO OUTFILE '/tmp/file'
-            FROM table
-            """
-        )
+        self.cur.execute.assert_called_once_with("SELECT * INTO OUTFILE %s FROM table", ("/tmp/file",))
 
     def test_serialize_cell(self):
         assert "foo" == self.db_hook._serialize_cell("foo", None)
@@ -310,14 +300,14 @@ def test_bulk_load_custom(self):
             IGNORE 1 LINES""",
         )
         self.cur.execute.assert_called_once_with(
-            """
-            LOAD DATA LOCAL INFILE '/tmp/file'
-            IGNORE
-            INTO TABLE table
-            FIELDS TERMINATED BY ';'
+            "LOAD DATA LOCAL INFILE %s %s INTO TABLE table %s",
+            (
+                "/tmp/file",
+                "IGNORE",
+                """FIELDS TERMINATED BY ';'
             OPTIONALLY ENCLOSED BY '"'
-            IGNORE 1 LINES
-            """
+            IGNORE 1 LINES""",
+            ),
         )
 
 
@@ -412,8 +402,5 @@ def test_mysql_hook_test_bulk_dump_mock(self, mock_get_conn, client):
             hook.bulk_dump(table, tmp_file)
 
             assert mock_execute.call_count == 1
-            query = f"""
-                SELECT * INTO OUTFILE '{tmp_file}'
-                FROM {table}
-            """
+            query = f"SELECT * INTO OUTFILE %s FROM {table}"
             assert_equal_ignore_multiple_spaces(mock_execute.call_args.args[0], query)