Make query configuration more general

pandas-dev · Dec 5, 2016 · c199935 · c199935
1 parent 0b365da
commit c199935
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 33 deletions.
diff --git a/doc/source/whatsnew/v0.20.0.txt b/doc/source/whatsnew/v0.20.0.txt
@@ -83,4 +83,4 @@ Performance Improvements
 .. _whatsnew_0200.bug_fixes:
 
 Bug Fixes
-~~~~~~~~~
+~~~~~~~~~
diff --git a/pandas/io/gbq.py b/pandas/io/gbq.py
@@ -395,9 +395,12 @@ def run_query(self, query, **kwargs):
                 }
             }
         }
-        query_config = kwargs.get('query_config')
-        if query_config is not None:
-            job_data['configuration']['query'].update(query_config)
+        configuration = kwargs.get('configuration')
+        if configuration is not None:
+            if 'query' in configuration:
+                job_data['configuration']['query'].update(configuration['query'])
+            else:
+                job_data['configuration'] = configuration
 
         self._start_timer()
         try:
@@ -687,7 +690,7 @@ def read_gbq(query, project_id=None, index_col=None, col_order=None,
         .. versionadded:: 0.19.0
 
     **kwargs: Arbitrary keyword arguments
-        query_config (dict): query configuration parameters for job processing.
+        configuration (dict): query configuration parameters for job processing.
             For more information see `BigQuery SQL Reference
             <https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.query>`
 

diff --git a/pandas/io/tests/test_gbq.py b/pandas/io/tests/test_gbq.py
@@ -713,49 +713,55 @@ def test_invalid_option_for_sql_dialect(self):
 
     def test_query_with_parameters(self):
         sql_statement = "SELECT @param1 + @param2 as VALID_RESULT"
-        query_config = {
-            "useLegacySql": False,
-            "parameterMode": "named",
-            "queryParameters": [
-                {
-                    "name": "param1",
-                    "parameterType": {
-                        "type": "INTEGER"
+        config = {
+            'query': {
+                "useLegacySql": False,
+                "parameterMode": "named",
+                "queryParameters": [
+                    {
+                        "name": "param1",
+                        "parameterType": {
+                            "type": "INTEGER"
+                        },
+                        "parameterValue": {
+                            "value": 1
+                        }
                     },
-                    "parameterValue": {
-                        "value": 1
+                    {
+                        "name": "param2",
+                        "parameterType": {
+                            "type": "INTEGER"
+                        },
+                        "parameterValue": {
+                            "value": 2
+                        }
                     }
-                },
-                {
-                    "name": "param2",
-                    "parameterType": {
-                        "type": "INTEGER"
-                    },
-                    "parameterValue": {
-                        "value": 2
-                    }
-                }
-            ]
+                ]
+            }
         }
-        # Test that an invalid query without query_config
+        # Test that an invalid query without query config
         with tm.assertRaises(ValueError):
             gbq.read_gbq(sql_statement, project_id=_get_project_id(),
                          private_key=_get_private_key_path())
 
         # Test that a correct query with query config
         df = gbq.read_gbq(sql_statement, project_id=_get_project_id(),
                           private_key=_get_private_key_path(),
-                          query_config=query_config)
+                          configuration=config)
         tm.assert_frame_equal(df, DataFrame({'VALID_RESULT': [3]}))
 
-    def test_query_no_cache(self):
+    def test_query_inside_configuration(self):
+        query_no_use = 'SELECT "PI_WRONG" as VALID_STRING'
         query = 'SELECT "PI" as VALID_STRING'
-        query_config = {
-            "useQueryCache": False,
+        config = {
+            'query': {
+                "query": query,
+                "useQueryCache": False,
+            }
         }
-        df = gbq.read_gbq(query, project_id=_get_project_id(),
+        df = gbq.read_gbq(query_no_use, project_id=_get_project_id(),
                           private_key=_get_private_key_path(),
-                          query_config=query_config)
+                          configuration=config)
         tm.assert_frame_equal(df, DataFrame({'VALID_STRING': ['PI']}))