Fix Query.data_sources method to return dict instead of list of tuples

google/datalab/bigquery/_query.py

@@ -197,7 +197,7 @@ def subqueries(self):
   @property
   def data_sources(self):
     """ Get a dictionary of external data sources referenced by the query."""
-    return self._data_sources
+    return dict(self._data_sources)
 
   def dry_run(self, context=None, query_params=None):
     """Dry run a query, to check the validity of the query and return some useful statistics.
@@ -218,7 +218,7 @@ def dry_run(self, context=None, query_params=None):
     api = _api.Api(context)
     try:
       query_result = api.jobs_insert_query(self.sql, dry_run=True,
-                                           table_definitions=self._data_sources,
+                                           table_definitions=self.data_sources,
                                            query_params=query_params)
     except Exception as e:
       raise e
@@ -263,7 +263,7 @@ def execute_async(self, output_options=None, sampling=None, context=None, query_
                                            append=append, overwrite=overwrite, batch=batch,
                                            use_cache=output_options.use_cache,
                                            allow_large_results=output_options.allow_large_results,
-                                           table_definitions=self._data_sources,
+                                           table_definitions=self.data_sources,
                                            query_params=query_params)
     except Exception as e:
       raise e

tests/bigquery/query_tests.py

@@ -43,6 +43,15 @@ def test_parameter_validation(self):
     self.assertEqual(q.udfs, {'testudf': udf})
     self.assertEqual(q._sql, sql)
 
+    with self.assertRaises(Exception) as error:
+      q = TestCases._create_query(sql, data_sources=['test_datasource'])
+    test_datasource = TestCases._create_data_source('gs://test/path')
+    env = {'test_datasource': test_datasource}
+    q = TestCases._create_query(sql, env=env, data_sources=['test_datasource'])
+    self.assertIsNotNone(q)
+    self.assertEqual(q.data_sources, {'test_datasource': test_datasource})
+    self.assertEqual(q._sql, sql)
+
   @mock.patch('google.datalab.bigquery._api.Api.tabledata_list')
   @mock.patch('google.datalab.bigquery._api.Api.jobs_insert_query')
   @mock.patch('google.datalab.bigquery._api.Api.jobs_query_results')
@@ -200,6 +209,10 @@ def _create_query(sql='SELECT * ...', name=None, env=None, udfs=None, data_sourc
   def _create_udf(name, code, return_type):
     return google.datalab.bigquery.UDF(name, code, return_type)
 
+  @staticmethod
+  def _create_data_source(source):
+    return google.datalab.bigquery.ExternalDataSource(source=source)
+
   @staticmethod
   def _create_context():
     project_id = 'test'