Quick Update GCS Presto

airflow/providers/presto/transfers/gcs_to_presto.py

@@ -18,8 +18,9 @@
 """This module contains Google Cloud Storage to Presto operator."""
 
 import csv
+import json
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Iterable, Optional, Sequence, Union
 
 from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
@@ -33,8 +34,9 @@ class GCSToPrestoOperator(BaseOperator):
     """
     Loads a csv file from Google Cloud Storage into a Presto table.
     Assumptions:
-    1. First row of the csv contains headers
+    1. CSV file should not have headers
     2. Presto table with requisite columns is already created
+    3. Optionally, a separate JSON file with headers or list of headers can be provided
 
     :param source_bucket: Source GCS bucket that contains the csv
     :param source_object: csv file including the path
@@ -55,6 +57,12 @@ class GCSToPrestoOperator(BaseOperator):
         account from the list granting this role to the originating account.
     """
 
+    template_fields: Sequence[str] = (
+        'source_bucket',
+        'source_object',
+        'presto_table',
+    )
+
     def __init__(
         self,
         *,
@@ -63,6 +71,8 @@ def __init__(
         presto_table: str,
         presto_conn_id: str = "presto_default",
         gcp_conn_id: str = "google_cloud_default",
+        schema_fields: Optional[Iterable[str]] = None,
+        schema_object: Optional[str] = None,
         delegate_to: Optional[str] = None,
         impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
         **kwargs,
@@ -73,6 +83,8 @@ def __init__(
         self.presto_table = presto_table
         self.presto_conn_id = presto_conn_id
         self.gcp_conn_id = gcp_conn_id
+        self.schema_fields = schema_fields
+        self.schema_object = schema_object
         self.delegate_to = delegate_to
         self.impersonation_chain = impersonation_chain
 
@@ -93,11 +105,18 @@ def execute(self, context: 'Context') -> None:
                 filename=temp_file.name,
             )
 
-            data = list(csv.reader(temp_file))
-            fields = tuple(data[0])
-            rows = []
-            for row in data[1:]:
-                rows.append(tuple(row))
-
+            data = csv.reader(temp_file)
+            rows = (tuple(row) for row in data)
             self.log.info("Inserting data into %s", self.presto_table)
-            presto_hook.insert_rows(table=self.presto_table, rows=rows, target_fields=fields)
+
+            if self.schema_fields:
+                presto_hook.insert_rows(table=self.presto_table, rows=rows, target_fields=self.schema_fields)
+            elif self.schema_object:
+                blob = gcs_hook.download(
+                    bucket_name=self.source_bucket,
+                    object_name=self.schema_object,
+                )
+                schema_fields = json.loads(blob.decode("utf-8"))
+                presto_hook.insert_rows(table=self.presto_table, rows=rows, target_fields=schema_fields)
+            else:
+                presto_hook.insert_rows(table=self.presto_table, rows=rows)

docs/apache-airflow-providers-presto/operators/transfer/gcs_to_presto.rst

@@ -41,8 +41,9 @@ Load CSV from GCS to Presto Table
 To load a CSV file from Google Cloud Storage to a Presto table you can use the
 :class:`~airflow.providers.presto.transfers.gcs_to_presto.GCSToPrestoOperator`.
 
-This operator assumes that the first row of the CSV contains headers corresponding to the columns in a
-pre-existing presto table.
+This operator assumes that CSV does not have headers and the data is corresponding to the columns in a
+pre-existing presto table. Optionally, you can provide schema as tuple/list of strings or as a path to a
+JSON file in the same bucket as the CSV file.
 
 .. exampleinclude:: /../../airflow/providers/presto/example_dags/example_gcs_to_presto.py
     :language: python

tests/providers/presto/transfers/test_gcs_presto.py

@@ -27,20 +27,19 @@
 PRESTO_CONN_ID = "test_presto"
 PRESTO_TABLE = "test_table"
 TASK_ID = "test_gcs_to_presto"
-VALUES = [("colA", "colB", "colC"), (1, 2, 3)]
+SCHEMA_FIELDS = ["colA", "colB", "colC"]
+SCHEMA_JSON = "path/to/file.json"
 
 
 class TestGCSToPrestoOperator(unittest.TestCase):
-    @mock.patch('airflow.providers.presto.transfers.gcs_to_presto.csv.reader')
     @mock.patch('airflow.providers.presto.transfers.gcs_to_presto.PrestoHook')
     @mock.patch("airflow.providers.presto.transfers.gcs_to_presto.GCSHook")
     @mock.patch("airflow.providers.presto.transfers.gcs_to_presto.NamedTemporaryFile")
-    def test_execute(self, mock_tempfile, mock_gcs_hook, mock_presto_hook, mock_reader):
+    def test_execute_without_schema(self, mock_tempfile, mock_gcs_hook, mock_presto_hook):
         filename = "file://97g23r"
         file_handle = mock.MagicMock()
         mock_tempfile.return_value.__enter__.return_value = file_handle
         mock_tempfile.return_value.__enter__.return_value.name = filename
-        mock_reader.return_value = VALUES
 
         op = GCSToPrestoOperator(
             task_id=TASK_ID,
@@ -67,5 +66,80 @@ def test_execute(self, mock_tempfile, mock_gcs_hook, mock_presto_hook, mock_read
 
         mock_insert = mock_presto_hook.return_value.insert_rows
 
-        fields = VALUES[0]
-        mock_insert.assert_called_once_with(table=PRESTO_TABLE, rows=VALUES[1:], target_fields=fields)
+        mock_insert.assert_called_once()
+
+    @mock.patch('airflow.providers.presto.transfers.gcs_to_presto.PrestoHook')
+    @mock.patch("airflow.providers.presto.transfers.gcs_to_presto.GCSHook")
+    @mock.patch("airflow.providers.presto.transfers.gcs_to_presto.NamedTemporaryFile")
+    def test_execute_schema_fields(self, mock_tempfile, mock_gcs_hook, mock_presto_hook):
+        filename = "file://97g23r"
+        file_handle = mock.MagicMock()
+        mock_tempfile.return_value.__enter__.return_value = file_handle
+        mock_tempfile.return_value.__enter__.return_value.name = filename
+
+        op = GCSToPrestoOperator(
+            task_id=TASK_ID,
+            source_bucket=BUCKET,
+            source_object=PATH,
+            presto_table=PRESTO_TABLE,
+            presto_conn_id=PRESTO_CONN_ID,
+            gcp_conn_id=GCP_CONN_ID,
+            schema_fields=SCHEMA_FIELDS,
+            impersonation_chain=IMPERSONATION_CHAIN,
+        )
+        op.execute(None)
+
+        mock_gcs_hook.assert_called_once_with(
+            gcp_conn_id=GCP_CONN_ID,
+            delegate_to=None,
+            impersonation_chain=IMPERSONATION_CHAIN,
+        )
+
+        mock_presto_hook.assert_called_once_with(presto_conn_id=PRESTO_CONN_ID)
+
+        mock_download = mock_gcs_hook.return_value.download
+
+        mock_download.assert_called_once_with(bucket_name=BUCKET, object_name=PATH, filename=filename)
+
+        mock_insert = mock_presto_hook.return_value.insert_rows
+
+        mock_insert.assert_called_once()
+
+    @mock.patch('airflow.providers.presto.transfers.gcs_to_presto.json.loads')
+    @mock.patch('airflow.providers.presto.transfers.gcs_to_presto.PrestoHook')
+    @mock.patch("airflow.providers.presto.transfers.gcs_to_presto.GCSHook")
+    @mock.patch("airflow.providers.presto.transfers.gcs_to_presto.NamedTemporaryFile")
+    def test_execute_schema_json(self, mock_tempfile, mock_gcs_hook, mock_presto_hook, mock_json_loader):
+        filename = "file://97g23r"
+        file_handle = mock.MagicMock()
+        mock_tempfile.return_value.__enter__.return_value = file_handle
+        mock_tempfile.return_value.__enter__.return_value.name = filename
+        mock_json_loader.return_value = SCHEMA_FIELDS
+
+        op = GCSToPrestoOperator(
+            task_id=TASK_ID,
+            source_bucket=BUCKET,
+            source_object=PATH,
+            presto_table=PRESTO_TABLE,
+            presto_conn_id=PRESTO_CONN_ID,
+            gcp_conn_id=GCP_CONN_ID,
+            schema_object=SCHEMA_JSON,
+            impersonation_chain=IMPERSONATION_CHAIN,
+        )
+        op.execute(None)
+
+        mock_gcs_hook.assert_called_once_with(
+            gcp_conn_id=GCP_CONN_ID,
+            delegate_to=None,
+            impersonation_chain=IMPERSONATION_CHAIN,
+        )
+
+        mock_presto_hook.assert_called_once_with(presto_conn_id=PRESTO_CONN_ID)
+
+        mock_download = mock_gcs_hook.return_value.download
+
+        assert mock_download.call_count == 2
+
+        mock_insert = mock_presto_hook.return_value.insert_rows
+
+        mock_insert.assert_called_once()