Merge pull request #2705 from spotify/claire/fix_dataflow_bq_return_type

Fix BigQueryTarget parsing in beam_dataflow module

Author: honnix

luigi/contrib/beam_dataflow.py

@@ -219,6 +219,7 @@ class BeamDataflowJobTask(MixinNaiveBulkComplete, luigi.Task):
     def __init__(self):
         if not isinstance(self.dataflow_params, DataflowParamKeys):
             raise ValueError("dataflow_params must be of type DataflowParamKeys")
+        super(BeamDataflowJobTask, self).__init__()
 
     @abstractmethod
     def dataflow_executable(self):
@@ -471,9 +472,13 @@ def _format_output_args(self):
 
     @staticmethod
     def get_target_path(target):
+        """
+            Given a luigi Target, determine a stringly typed path to pass as a
+            Dataflow job argument.
+        """
         if isinstance(target, luigi.LocalTarget) or isinstance(target, gcs.GCSTarget):
             return target.path
         elif isinstance(target, bigquery.BigQueryTarget):
-            "{}:{}.{}".format(target.project_id, target.dataset_id, target.table_id)
+            return "{}:{}.{}".format(target.table.project_id, target.table.dataset_id, target.table.table_id)
         else:
-            raise ValueError("Target not supported")
+            raise ValueError("Target %s not supported" % target)

test/contrib/beam_dataflow_test.py

@@ -17,7 +17,7 @@
 
 import json
 import luigi
-from luigi.contrib import beam_dataflow
+from luigi.contrib import beam_dataflow, bigquery, gcs
 from luigi import local_target
 import mock
 from mock import MagicMock, patch
@@ -282,6 +282,20 @@ def output(self):
 
         self.assertEqual(TestTaskDictOfMixedCompleteOutput().complete(), False)
 
+    def test_get_target_path(self):
+        bq_target = bigquery.BigQueryTarget("p", "d", "t", client="fake_client")
+        self.assertEqual(
+            SimpleTestTask.get_target_path(bq_target),
+            "p:d.t")
+
+        gcs_target = gcs.GCSTarget("gs://foo/bar.txt", client="fake_client")
+        self.assertEqual(
+            SimpleTestTask.get_target_path(gcs_target),
+            "gs://foo/bar.txt")
+
+        with self.assertRaises(ValueError):
+            SimpleTestTask.get_target_path("not_a_target")
+
     def test_dataflow_runner_resolution(self):
         task = SimpleTestTask()
         # Test that supported runners are passed through