Merge pull request #2884 from tseaver/2882-bigquery-copyjob-w-single-source-table

tseaver · web-flow · commit d158cc784075 · 2016-12-28T17:19:50.000-05:00
Accept copy jobs with single 'sourceTable' config.
diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
@@ -742,7 +742,14 @@ def from_api_repr(cls, resource, client):
         dataset = Dataset(dest_config['datasetId'], client)
         destination = Table(dest_config['tableId'], dataset)
         sources = []
-        for source_config in config['sourceTables']:
+        source_configs = config.get('sourceTables')
+        if source_configs is None:
+            single = config.get('sourceTable')
+            if single is None:
+                raise KeyError(
+                    "Resource missing 'sourceTables' / 'sourceTable'")
+            source_configs = [single]
+        for source_config in source_configs:
             dataset = Dataset(source_config['datasetId'], client)
             sources.append(Table(source_config['tableId'], dataset))
         job = cls(name, destination, sources, client=client)
diff --git a/bigquery/unit_tests/test_job.py b/bigquery/unit_tests/test_job.py
@@ -675,7 +675,9 @@ def _verifyResourceProperties(self, job, resource):
         self.assertEqual(job.destination.dataset_name, table_ref['datasetId'])
         self.assertEqual(job.destination.name, table_ref['tableId'])
 
-        sources = config['sourceTables']
+        sources = config.get('sourceTables')
+        if sources is None:
+            sources = [config['sourceTable']]
         self.assertEqual(len(sources), len(job.sources))
         for table_ref, table in zip(sources, job.sources):
             self.assertEqual(table.project, table_ref['projectId'])
@@ -764,6 +766,58 @@ def test_from_api_repr_bare(self):
         self.assertIs(job._client, client)
         self._verifyResourceProperties(job, RESOURCE)
 
+    def test_from_api_repr_w_sourcetable(self):
+        self._setUpConstants()
+        client = _Client(self.PROJECT)
+        RESOURCE = {
+            'id': self.JOB_ID,
+            'jobReference': {
+                'projectId': self.PROJECT,
+                'jobId': self.JOB_NAME,
+            },
+            'configuration': {
+                'copy': {
+                    'sourceTable': {
+                        'projectId': self.PROJECT,
+                        'datasetId': self.DS_NAME,
+                        'tableId': self.SOURCE_TABLE,
+                    },
+                    'destinationTable': {
+                        'projectId': self.PROJECT,
+                        'datasetId': self.DS_NAME,
+                        'tableId': self.DESTINATION_TABLE,
+                    },
+                }
+            },
+        }
+        klass = self._get_target_class()
+        job = klass.from_api_repr(RESOURCE, client=client)
+        self.assertIs(job._client, client)
+        self._verifyResourceProperties(job, RESOURCE)
+
+    def test_from_api_repr_wo_sources(self):
+        self._setUpConstants()
+        client = _Client(self.PROJECT)
+        RESOURCE = {
+            'id': self.JOB_ID,
+            'jobReference': {
+                'projectId': self.PROJECT,
+                'jobId': self.JOB_NAME,
+            },
+            'configuration': {
+                'copy': {
+                    'destinationTable': {
+                        'projectId': self.PROJECT,
+                        'datasetId': self.DS_NAME,
+                        'tableId': self.DESTINATION_TABLE,
+                    },
+                }
+            },
+        }
+        klass = self._get_target_class()
+        with self.assertRaises(KeyError):
+            klass.from_api_repr(RESOURCE, client=client)
+
     def test_from_api_repr_w_properties(self):
         client = _Client(self.PROJECT)
         RESOURCE = self._makeResource()
