[WIP] Apply best practices for SQL queries

airflow/migrations/utils.py

@@ -35,15 +35,15 @@ def get_mssql_table_constraints(conn, table_name) -> dict[str, dict[str, list[st
     :return: a dictionary of ((constraint name, constraint type), column name) of table
     """
     query = text(
-        f"""SELECT tc.CONSTRAINT_NAME , tc.CONSTRAINT_TYPE, ccu.COLUMN_NAME
+        """SELECT tc.CONSTRAINT_NAME , tc.CONSTRAINT_TYPE, ccu.COLUMN_NAME
      FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc
      JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE AS ccu ON ccu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME
-     WHERE tc.TABLE_NAME = '{table_name}' AND
+     WHERE tc.TABLE_NAME = :table_name AND
      (tc.CONSTRAINT_TYPE = 'PRIMARY KEY' or UPPER(tc.CONSTRAINT_TYPE) = 'UNIQUE'
      or UPPER(tc.CONSTRAINT_TYPE) = 'FOREIGN KEY')
     """
     )
-    result = conn.execute(query).fetchall()
+    result = conn.execute(query, {"table_name": table_name}).fetchall()
     constraint_dict = defaultdict(lambda: defaultdict(list))
     for constraint, constraint_type, col_name in result:
         constraint_dict[constraint_type][constraint].append(col_name)

airflow/migrations/versions/0046_1_10_5_change_datetime_to_datetime2_6_on_mssql_.py

@@ -260,14 +260,14 @@ def get_table_constraints(conn, table_name) -> dict[tuple[str, str], list[str]]:
     :return: a dictionary of ((constraint name, constraint type), column name) of table
     """
     query = text(
-        f"""SELECT tc.CONSTRAINT_NAME , tc.CONSTRAINT_TYPE, ccu.COLUMN_NAME
+        """SELECT tc.CONSTRAINT_NAME , tc.CONSTRAINT_TYPE, ccu.COLUMN_NAME
      FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc
      JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE AS ccu ON ccu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME
-     WHERE tc.TABLE_NAME = '{table_name}' AND
+     WHERE tc.TABLE_NAME = :table_name AND
      (tc.CONSTRAINT_TYPE = 'PRIMARY KEY' or UPPER(tc.CONSTRAINT_TYPE) = 'UNIQUE')
     """
     )
-    result = conn.execute(query).fetchall()
+    result = conn.execute(query, {"table_name": table_name}).fetchall()
     constraint_dict = defaultdict(list)
     for constraint, constraint_type, column in result:
         constraint_dict[(constraint, constraint_type)].append(column)

airflow/migrations/versions/0060_2_0_0_remove_id_column_from_xcom.py

@@ -50,14 +50,14 @@ def get_table_constraints(conn, table_name) -> dict[tuple[str, str], list[str]]:
     :return: a dictionary of ((constraint name, constraint type), column name) of table
     """
     query = text(
-        f"""SELECT tc.CONSTRAINT_NAME , tc.CONSTRAINT_TYPE, ccu.COLUMN_NAME
+        """SELECT tc.CONSTRAINT_NAME , tc.CONSTRAINT_TYPE, ccu.COLUMN_NAME
      FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc
      JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE AS ccu ON ccu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME
-     WHERE tc.TABLE_NAME = '{table_name}' AND
+     WHERE tc.TABLE_NAME = :table_name AND
      (tc.CONSTRAINT_TYPE = 'PRIMARY KEY' or UPPER(tc.CONSTRAINT_TYPE) = 'UNIQUE')
     """
     )
-    result = conn.execute(query).fetchall()
+    result = conn.execute(query, {"table_name": table_name}).fetchall()
     constraint_dict = defaultdict(list)
     for constraint, constraint_type, column in result:
         constraint_dict[(constraint, constraint_type)].append(column)

airflow/migrations/versions/0069_2_0_0_add_scheduling_decision_to_dagrun_and_.py

@@ -69,15 +69,20 @@ def upgrade():
 
     # Set it to true here as it makes us take the slow/more complete path, and when it's next parsed by the
     # DagParser it will get set to correct value.
-
-    op.execute(
-        f"""
+    query = sa.text(
+        """
         UPDATE dag SET
-            concurrency={concurrency},
-            has_task_concurrency_limits={1 if is_sqlite or is_mssql else sa.true()}
+            concurrency=:concurrency,
+            has_task_concurrency_limits=:has_task_concurrency_limits
         where concurrency IS NULL
         """
+    ).bindparams(
+        {
+            "concurrency": concurrency,
+            "has_task_concurrency_limits": (1 if is_sqlite or is_mssql else sa.true()),
+        }
     )
+    op.execute(query)
 
     with op.batch_alter_table("dag", schema=None) as batch_op:
         batch_op.alter_column("concurrency", type_=sa.Integer(), nullable=False)

airflow/migrations/versions/0088_2_2_0_improve_mssql_compatibility.py

@@ -49,7 +49,7 @@ def is_table_empty(conn, table_name):
     :param table_name: table name
     :return: Booelan indicating if the table is present
     """
-    return conn.execute(text(f"select TOP 1 * from {table_name}")).first() is None
+    return conn.execute(text("select TOP 1 * from :table_name"), {"table_name": table_name}).first() is None
 
 
 def get_table_constraints(conn, table_name) -> dict[tuple[str, str], list[str]]:
@@ -68,11 +68,11 @@ def get_table_constraints(conn, table_name) -> dict[tuple[str, str], list[str]]:
         f"""SELECT tc.CONSTRAINT_NAME , tc.CONSTRAINT_TYPE, ccu.COLUMN_NAME
      FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc
      JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE AS ccu ON ccu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME
-     WHERE tc.TABLE_NAME = '{table_name}' AND
+     WHERE tc.TABLE_NAME = :table_name AND
      (tc.CONSTRAINT_TYPE = 'PRIMARY KEY' or UPPER(tc.CONSTRAINT_TYPE) = 'UNIQUE')
     """
     )
-    result = conn.execute(query).fetchall()
+    result = conn.execute(query, {"table_name": table_name}).fetchall()
     constraint_dict = defaultdict(list)
     for constraint, constraint_type, column in result:
         constraint_dict[(constraint, constraint_type)].append(column)
@@ -111,15 +111,15 @@ def create_constraints(operator, column_name, constraint_dict):
 
 def _is_timestamp(conn, table_name, column_name):
     query = text(
-        f"""SELECT
+        """SELECT
     TYPE_NAME(C.USER_TYPE_ID) AS DATA_TYPE
     FROM SYS.COLUMNS C
     JOIN SYS.TYPES T
     ON C.USER_TYPE_ID=T.USER_TYPE_ID
-    WHERE C.OBJECT_ID=OBJECT_ID('{table_name}') and C.NAME='{column_name}';
+    WHERE C.OBJECT_ID=OBJECT_ID(:table_name) and C.NAME=:column_name';
     """
     )
-    column_type = conn.execute(query).fetchone()[0]
+    column_type = conn.execute(query, {"table_name": table_name, "column_name": column_name}).fetchone()[0]
     return column_type == "timestamp"
 
 

airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py

@@ -41,7 +41,11 @@ def upgrade():
     with op.batch_alter_table("slot_pool") as batch_op:
         batch_op.add_column(sa.Column("include_deferred", sa.Boolean))
     # Different databases support different literal for FALSE. This is fine.
-    op.execute(sa.text(f"UPDATE slot_pool SET include_deferred = {sa.false().compile(op.get_bind())}"))
+    op.execute(
+        sa.text(
+            "UPDATE slot_pool SET include_deferred = :include_deferred"
+        ).bindparams({"include_deferred": (sa.false().compile(op.get_bind()))})
+    )
     with op.batch_alter_table("slot_pool") as batch_op:
         batch_op.alter_column("include_deferred", existing_type=sa.Boolean, nullable=False)
 

airflow/providers/amazon/aws/hooks/redshift_data.py

@@ -162,23 +162,28 @@ def get_table_primary_key(
 
         :return: Primary key columns list
         """
-        sql = f"""
+        sql = """
             select kcu.column_name
             from information_schema.table_constraints tco
                     join information_schema.key_column_usage kcu
                         on kcu.constraint_name = tco.constraint_name
                             and kcu.constraint_schema = tco.constraint_schema
                             and kcu.constraint_name = tco.constraint_name
             where tco.constraint_type = 'PRIMARY KEY'
-            and kcu.table_schema = {schema}
-            and kcu.table_name = {table}
+            and kcu.table_schema = :schema
+            and kcu.table_name = :table
         """
+        parameters = [
+            {"name": "schema", "value": {"stringValue": schema}},
+            {"name": "table", "value": {"stringValue": table}},
+        ]
         stmt_id = self.execute_query(
             sql=sql,
             database=database,
             cluster_identifier=cluster_identifier,
             workgroup_name=workgroup_name,
             db_user=db_user,
+            parameters=parameters,
             secret_arn=secret_arn,
             statement_name=statement_name,
             with_event=with_event,

airflow/providers/amazon/aws/transfers/redshift_to_s3.py

@@ -124,7 +124,11 @@ def __init__(
         if select_query:
             self.select_query = select_query
         elif self.schema and self.table:
-            self.select_query = f"SELECT * FROM {self.schema}.{self.table}"
+            self.select_query = "SELECT * FROM :select_query_schema.:select_query_table"
+            self.parameters = list(self.parameters) + [
+                {"name": "select_query_schema", "value": {"stringValue": schema}},
+                {"name": "select_query_table", "value": {"stringValue": table}},
+            ]
         else:
             raise ValueError(
                 "Please provide both `schema` and `table` params or `select_query` to fetch the data."
@@ -141,12 +145,16 @@ def __init__(
     def _build_unload_query(
         self, credentials_block: str, select_query: str, s3_key: str, unload_options: str
     ) -> str:
+        self.parameters = list(self.parameters) + [
+            {"name": "unload_query_credentials_block", "value": {"stringValue": credentials_block}},
+            {"name": "unload_query_options", "value": {"stringValue": unload_options}},
+        ]
         return f"""
                     UNLOAD ('{select_query}')
                     TO 's3://{self.s3_bucket}/{s3_key}'
                     credentials
-                    '{credentials_block}'
-                    {unload_options};
+                    :unload_query_credentials_block
+                    :unload_query_options;
         """
 
     def execute(self, context: Context) -> None:

airflow/providers/amazon/aws/transfers/s3_to_redshift.py

@@ -161,9 +161,11 @@ def execute(self, context: Context) -> None:
         )
 
         sql: str | Iterable[str]
+        parameters: list[dict] = []
 
         if self.method == "REPLACE":
-            sql = ["BEGIN;", f"DELETE FROM {destination};", copy_statement, "COMMIT"]
+            sql = ["BEGIN;", "DELETE FROM :destination;", copy_statement, "COMMIT"]
+            parameters = [{"name": "destination", "value": {"stringValue": destination}}]
         elif self.method == "UPSERT":
             if isinstance(redshift_hook, RedshiftDataHook):
                 keys = self.upsert_keys or redshift_hook.get_table_primary_key(
@@ -178,20 +180,25 @@ def execute(self, context: Context) -> None:
             where_statement = " AND ".join([f"{self.table}.{k} = {copy_destination}.{k}" for k in keys])
 
             sql = [
-                f"CREATE TABLE {copy_destination} (LIKE {destination} INCLUDING DEFAULTS);",
+                "CREATE TABLE :copy_destination (LIKE :destination INCLUDING DEFAULTS);",
                 copy_statement,
                 "BEGIN;",
-                f"DELETE FROM {destination} USING {copy_destination} WHERE {where_statement};",
-                f"INSERT INTO {destination} SELECT * FROM {copy_destination};",
+                "DELETE FROM :destination USING :copy_destination WHERE :where_statement;",
+                "INSERT INTO :destination SELECT * FROM :copy_destination;",
                 "COMMIT",
             ]
+            parameters = [
+                {"name": "copy_destination", "value": {"stringValue": copy_destination}},
+                {"name": "destination", "value": {"stringValue": destination}},
+                {"name": "where_statement", "value": {"stringValue": where_statement}},
+            ]
 
         else:
             sql = copy_statement
 
         self.log.info("Executing COPY command...")
         if isinstance(redshift_hook, RedshiftDataHook):
-            redshift_hook.execute_query(sql=sql, **self.redshift_data_api_kwargs)
+            redshift_hook.execute_query(sql=sql, parameters=parameters, **self.redshift_data_api_kwargs)
         else:
-            redshift_hook.run(sql, autocommit=self.autocommit)
+            redshift_hook.run(sql, parameters=parameters, autocommit=self.autocommit)
         self.log.info("COPY command complete...")

airflow/providers/apache/cassandra/hooks/cassandra.py

@@ -200,9 +200,11 @@ def record_exists(self, table: str, keys: dict[str, str]) -> bool:
         if "." in table:
             keyspace, table = table.split(".", 1)
         ks_str = " AND ".join(f"{key}=%({key})s" for key in keys)
-        query = f"SELECT * FROM {keyspace}.{table} WHERE {ks_str}"
+        query = f"SELECT * FROM %(query_keyspace).%(query_table) WHERE {ks_str}"
         try:
-            result = self.get_conn().execute(query, keys)
+            result = self.get_conn().execute(
+                query, {**keys, "query_keyspace": keyspace, "query_table": table}
+            )
             return result.one() is not None
         except Exception:
             return False

docs/apache-airflow/img/airflow_erd.sha256

@@ -1 +1 @@
-8229a936107bee851d6a39c791b842b11f295ffa308b18106e45298a50871493
+e90d46d6a9661d874ac97f189900d05dbef80cac7f73029b44d6f0b9fe12c35c