Repository: apache/airflow
Pull request: Feature qubole hook support headers
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.

Feel free to ping committers for the review!

In case of existing issue, reference it using one of the following:

How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->

closes: #15598
related: #15598

---
Added support to include column headers in the result from Qubole.

Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines)** for more information.
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).


Problem statement:
Qubole Hook Does Not Support 'include_headers'


**Description**

Qubole Hook and Operator do not support `include_header` param for getting results with headers
Add Support for `include_header` get_results(... arguments=[True])
**Use case / motivation**

It's very hard to work with CSV results from db without headers. 
This is super important when using Qubole's databases.

**Are you willing to submit a PR?**

Not sure yet, I can give it a try

**Related Issues**

Hints from issue discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

@levyitay assigned to you

Commit message: Added missing issue number

Patch to review:
diff --git a/airflow/providers/qubole/CHANGELOG.rst b/airflow/providers/qubole/CHANGELOG.rst
index eac8d8e171110..3512f3d873380 100644
--- a/airflow/providers/qubole/CHANGELOG.rst
+++ b/airflow/providers/qubole/CHANGELOG.rst
@@ -19,6 +19,15 @@
 Changelog
 ---------
 
+1.0.3
+.....
+
+Features
+~~~~~~~~
+
+*  ``Feature add support for include_headers in get_results (#15598)``
+
+
 1.0.2
 .....
 
diff --git a/airflow/providers/qubole/hooks/qubole.py b/airflow/providers/qubole/hooks/qubole.py
index 99e6027408709..579441eff9e36 100644
--- a/airflow/providers/qubole/hooks/qubole.py
+++ b/airflow/providers/qubole/hooks/qubole.py
@@ -203,7 +203,7 @@ def kill(self, ti):
             self.cmd.cancel()
 
     # pylint: disable=consider-using-with
-    def get_results(self, ti=None, fp=None, inline: bool = True, delim=None, fetch: bool = True) -> str:
+    def get_results(self, ti=None, fp=None, inline: bool = True, delim=None, fetch: bool = True, include_headers: bool = False) -> str:
         """
         Get results (or just s3 locations) of a command from Qubole and save into a file
 
@@ -225,7 +225,7 @@ def get_results(self, ti=None, fp=None, inline: bool = True, delim=None, fetch:
             cmd_id = ti.xcom_pull(key="qbol_cmd_id", task_ids=self.task_id)
             self.cmd = self.cls.find(cmd_id)
 
-        self.cmd.get_results(fp, inline, delim, fetch)  # type: ignore[attr-defined]
+        self.cmd.get_results(fp, inline, delim, fetch, arguments=[include_headers])  # type: ignore[attr-defined]
         fp.flush()
         fp.close()
         return fp.name
diff --git a/airflow/providers/qubole/hooks/qubole_check.py b/airflow/providers/qubole/hooks/qubole_check.py
index 987ea59ae2021..c15f47f54ff1c 100644
--- a/airflow/providers/qubole/hooks/qubole_check.py
+++ b/airflow/providers/qubole/hooks/qubole_check.py
@@ -112,7 +112,7 @@ def get_query_results(self) -> Optional[str]:
             cmd_id = self.cmd.id
             self.log.info("command id: %d", cmd_id)
             query_result_buffer = StringIO()
-            self.cmd.get_results(fp=query_result_buffer, inline=True, delim=COL_DELIM)
+            self.cmd.get_results(fp=query_result_buffer, inline=True, delim=COL_DELIM, arguments=[True])
             query_result = query_result_buffer.getvalue()
             query_result_buffer.close()
             return query_result
diff --git a/airflow/providers/qubole/operators/qubole.py b/airflow/providers/qubole/operators/qubole.py
index b090af6dc6ebe..35024153d1ed8 100644
--- a/airflow/providers/qubole/operators/qubole.py
+++ b/airflow/providers/qubole/operators/qubole.py
@@ -250,9 +250,9 @@ def on_kill(self, ti=None) -> None:
         else:
             self.get_hook().kill(ti)
 
-    def get_results(self, ti=None, fp=None, inline: bool = True, delim=None, fetch: bool = True) -> str:
+    def get_results(self, ti=None, fp=None, inline: bool = True, delim=None, fetch: bool = True, include_headers: bool = False) -> str:
         """get_results from Qubole"""
-        return self.get_hook().get_results(ti, fp, inline, delim, fetch)
+        return self.get_hook().get_results(ti, fp, inline, delim, fetch,include_headers)
 
     def get_log(self, ti) -> None:
         """get_log from Qubole"""
diff --git a/airflow/providers/qubole/provider.yaml b/airflow/providers/qubole/provider.yaml
index bd4cd5827ee08..dad3364cd8ded 100644
--- a/airflow/providers/qubole/provider.yaml
+++ b/airflow/providers/qubole/provider.yaml
@@ -22,6 +22,7 @@ description: |
     `Qubole <https://www.qubole.com/>`__
 
 versions:
+  - 1.0.3
   - 1.0.2
   - 1.0.1
   - 1.0.0


Code review comments:
diff_hunk: @@ -250,9 +250,9 @@ def on_kill(self, ti=None) -> None:
         else:
             self.get_hook().kill(ti)
 
-    def get_results(self, ti=None, fp=None, inline: bool = True, delim=None, fetch: bool = True) -> str:
+    def get_results(self, ti=None, fp=None, inline: bool = True, delim=None, fetch: bool = True, include_headers: bool = False) -> str:
         """get_results from Qubole"""
-        return self.get_hook().get_results(ti, fp, inline, delim, fetch)
+        return self.get_hook().get_results(ti, fp, inline, delim, fetch,include_headers)
path: airflow/providers/qubole/operators/qubole.py
text: @user1:
```suggestion
        return self.get_hook().get_results(ti, fp, inline, delim, fetch, include_headers)
```

Before pylint wastes a CI run to catch it…

@author:
Thanks, not sure why I missed it. 
Fixed along with some other warnings