jroachgolf84 commented on code in PR #42900:
URL: https://github.com/apache/airflow/pull/42900#discussion_r1799619239


##########
providers/src/airflow/providers/amazon/aws/operators/redshift_data.py:
##########
@@ -197,16 +196,38 @@ def execute_complete(
             raise AirflowException("statement_id should not be empty.")
 
         self.log.info("%s completed successfully.", self.task_id)
-        if self.return_sql_result:
-            result = self.hook.conn.get_statement_result(Id=statement_id)
-            self.log.debug("Statement result: %s", result)
-            return result
 
-        return statement_id
+        # Use the get_sql_results method to return the results of the SQL 
query, or the statement_ids,
+        # depending on the value of self.return_sql_result
+        return self.get_sql_results(return_sql_result=self.return_sql_result)
+
+    def get_sql_results(self, return_sql_result: bool) -> 
list[GetStatementResultResponseTypeDef] | list[str]:
+        """
+        Retrieve either the result of the SQL query, or the statement ID(s).
+
+        :param return_sql_result:
+        """
+        # ISSUE-40427: Pull the statement, and check to see if there are 
sub-statements. If that is the
+        # case, pull each of the sub-statement ID's, and grab the results. 
Otherwise, just use
+        # self.statement_id
+        statement: DescribeStatementResponseTypeDef = 
self.hook.conn.describe_statement(Id=self.statement_id)

Review Comment:
   Updated, should take care of the issue!



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to