restore imports

Author: sycai

pandas_gbq/gbq.py

@@ -11,19 +11,39 @@
 from pandas_gbq.contexts import context
 from pandas_gbq.exceptions import (
     DatasetCreationError,
+    GenericGBQException,
     InvalidColumnOrder,
     InvalidIndexColumn,
+    InvalidPageToken,
+    InvalidSchema,
     NotFoundException,
+    QueryTimeout,
     TableCreationError,
 )
 from pandas_gbq.features import FEATURES
-from pandas_gbq.gbq_connector import GbqConnector
+from pandas_gbq.gbq_connector import GbqConnector, create_user_agent
 import pandas_gbq.schema
 import pandas_gbq.schema.pandas_to_bigquery
 
 logger = logging.getLogger(__name__)
 
 
+__all__ = [
+    # The following items were originally defined in this module
+    # Export them here for backward compatibility.
+    "DatasetCreationError",
+    "InvalidColumnOrder",
+    "InvalidIndexColumn",
+    "InvalidPageToken",
+    "InvalidSchema",
+    "NotFoundException",
+    "TableCreationError",
+    "GenericGBQException",
+    "QueryTimeout",
+    "create_user_agent",
+]
+
+
 def _test_google_api_imports():
     try:
         import packaging  # noqa

tests/system/test_gbq.py

@@ -16,7 +16,7 @@
 import pytest
 import pytz
 
-from pandas_gbq import exceptions, gbq
+from pandas_gbq import gbq
 import pandas_gbq.schema
 
 TABLE_ID = "new_test"
@@ -195,7 +195,7 @@ def test_read_gbq_raises_invalid_column_order(self, project_id):
         col_order = ["string_aaa", "string_1", "string_2"]
 
         # Column string_aaa does not exist. Should raise InvalidColumnOrder
-        with pytest.raises(exceptions.InvalidColumnOrder):
+        with pytest.raises(gbq.InvalidColumnOrder):
             gbq.read_gbq(
                 query,
                 project_id=project_id,
@@ -227,7 +227,7 @@ def test_read_gbq_raises_invalid_index_column(self, project_id):
         col_order = ["string_3", "string_2"]
 
         # Column string_bbb does not exist. Should raise InvalidIndexColumn
-        with pytest.raises(exceptions.InvalidIndexColumn):
+        with pytest.raises(gbq.InvalidIndexColumn):
             gbq.read_gbq(
                 query,
                 project_id=project_id,
@@ -238,7 +238,7 @@ def test_read_gbq_raises_invalid_index_column(self, project_id):
             )
 
     def test_malformed_query(self, project_id):
-        with pytest.raises(exceptions.GenericGBQException):
+        with pytest.raises(gbq.GenericGBQException):
             gbq.read_gbq(
                 "SELCET * FORM [publicdata:samples.shakespeare]",
                 project_id=project_id,
@@ -247,7 +247,7 @@ def test_malformed_query(self, project_id):
             )
 
     def test_bad_project_id(self):
-        with pytest.raises(exceptions.GenericGBQException):
+        with pytest.raises(gbq.GenericGBQException):
             gbq.read_gbq(
                 "SELCET * FROM [publicdata:samples.shakespeare]",
                 project_id="not-my-project",
@@ -256,7 +256,7 @@ def test_bad_project_id(self):
             )
 
     def test_bad_table_name(self, project_id):
-        with pytest.raises(exceptions.GenericGBQException):
+        with pytest.raises(gbq.GenericGBQException):
             gbq.read_gbq(
                 "SELECT * FROM [publicdata:samples.nope]",
                 project_id=project_id,
@@ -316,7 +316,7 @@ def test_legacy_sql(self, project_id):
 
         # Test that a legacy sql statement fails when
         # setting dialect='standard'
-        with pytest.raises(exceptions.GenericGBQException):
+        with pytest.raises(gbq.GenericGBQException):
             gbq.read_gbq(
                 legacy_sql,
                 project_id=project_id,
@@ -341,7 +341,7 @@ def test_standard_sql(self, project_id):
 
         # Test that a standard sql statement fails when using
         # the legacy SQL dialect.
-        with pytest.raises(exceptions.GenericGBQException):
+        with pytest.raises(gbq.GenericGBQException):
             gbq.read_gbq(
                 standard_sql,
                 project_id=project_id,
@@ -485,7 +485,7 @@ def test_timeout_configuration(self, project_id):
             {"query": {"useQueryCache": False}, "jobTimeoutMs": 401},
         ]
         for config in configs:
-            with pytest.raises(exceptions.QueryTimeout):
+            with pytest.raises(gbq.QueryTimeout):
                 gbq.read_gbq(
                     sql_statement,
                     project_id=project_id,
@@ -507,7 +507,7 @@ def test_timeout_configuration(self, project_id):
         ]
 
         for config in configs:
-            with pytest.raises(exceptions.QueryTimeout):
+            with pytest.raises(gbq.QueryTimeout):
                 gbq.read_gbq(
                     sql_statement,
                     project_id=project_id,
@@ -741,7 +741,7 @@ def test_upload_data_if_table_exists_fail(self, project_id):
         )
 
         # Test the default value of if_exists == 'fail'
-        with pytest.raises(exceptions.TableCreationError):
+        with pytest.raises(gbq.TableCreationError):
             gbq.to_gbq(
                 df,
                 self.destination_table + test_id,
@@ -750,7 +750,7 @@ def test_upload_data_if_table_exists_fail(self, project_id):
             )
 
         # Test the if_exists parameter with value 'fail'
-        with pytest.raises(exceptions.TableCreationError):
+        with pytest.raises(gbq.TableCreationError):
             gbq.to_gbq(
                 df,
                 self.destination_table + test_id,
@@ -794,7 +794,7 @@ def test_upload_data_if_table_exists_append(self, project_id):
         assert result["num_rows"][0] == test_size * 2
 
         # Try inserting with a different schema, confirm failure
-        with pytest.raises(exceptions.InvalidSchema):
+        with pytest.raises(gbq.InvalidSchema):
             gbq.to_gbq(
                 df_different_schema,
                 self.destination_table + test_id,
@@ -1029,7 +1029,7 @@ def test_upload_data_with_invalid_user_schema_raises_error(self, project_id):
             {"name": "D", "type": "FLOAT"},
         ]
         destination_table = self.destination_table + test_id
-        with pytest.raises(exceptions.GenericGBQException):
+        with pytest.raises(gbq.GenericGBQException):
             gbq.to_gbq(
                 df,
                 destination_table,
@@ -1047,7 +1047,7 @@ def test_upload_data_with_missing_schema_fields_raises_error(self, project_id):
             {"name": "C", "type": "FLOAT"},
         ]
         destination_table = self.destination_table + test_id
-        with pytest.raises(exceptions.GenericGBQException):
+        with pytest.raises(gbq.GenericGBQException):
             gbq.to_gbq(
                 df,
                 destination_table,
@@ -1193,7 +1193,7 @@ def test_create_dataset(bigquery_client, gbq_dataset, random_dataset_id, project
 
 def test_create_dataset_already_exists(gbq_dataset, random_dataset_id):
     gbq_dataset.create(random_dataset_id)
-    with pytest.raises(exceptions.DatasetCreationError):
+    with pytest.raises(gbq.DatasetCreationError):
         gbq_dataset.create(random_dataset_id)
 
 
@@ -1218,7 +1218,7 @@ def test_create_table(gbq_table):
 def test_create_table_already_exists(gbq_table):
     schema = gbq._generate_bq_schema(make_mixed_dataframe_v1())
     gbq_table.create("test_create_table_exists", schema)
-    with pytest.raises(exceptions.TableCreationError):
+    with pytest.raises(gbq.TableCreationError):
         gbq_table.create("test_create_table_exists", schema)
 
 

tests/unit/test_gbq.py

@@ -19,8 +19,9 @@
 from pandas import DataFrame
 import pytest
 
-from pandas_gbq import exceptions, gbq, gbq_connector
+from pandas_gbq import gbq, gbq_connector
 import pandas_gbq.constants
+import pandas_gbq.exceptions
 import pandas_gbq.features
 from pandas_gbq.features import FEATURES
 
@@ -221,7 +222,7 @@ def test_GbqConnector_process_http_error_transforms_timeout():
     original = google.api_core.exceptions.GoogleAPICallError(
         "Job execution was cancelled: Job timed out after 0s"
     )
-    with pytest.raises(exceptions.QueryTimeout):
+    with pytest.raises(gbq.QueryTimeout):
         gbq.GbqConnector.process_http_error(original)
 
 
@@ -389,7 +390,7 @@ def test_dataset_exists_translates_exception(mock_bigquery_client):
     mock_bigquery_client.get_dataset.side_effect = (
         google.api_core.exceptions.InternalServerError("something went wrong")
     )
-    with pytest.raises(exceptions.GenericGBQException):
+    with pytest.raises(gbq.GenericGBQException):
         connector.exists("not_gonna_work")
 
 
@@ -412,7 +413,7 @@ def test_table_create_translates_exception(mock_bigquery_client):
     mock_bigquery_client.create_table.side_effect = (
         google.api_core.exceptions.InternalServerError("something went wrong")
     )
-    with pytest.raises(exceptions.GenericGBQException):
+    with pytest.raises(gbq.GenericGBQException):
         connector.create(
             "not_gonna_work", {"fields": [{"name": "f", "type": "STRING"}]}
         )
@@ -434,7 +435,7 @@ def test_table_delete_translates_exception(mock_bigquery_client):
     mock_bigquery_client.delete_table.side_effect = (
         google.api_core.exceptions.InternalServerError("something went wrong")
     )
-    with pytest.raises(exceptions.GenericGBQException):
+    with pytest.raises(gbq.GenericGBQException):
         connector.delete("not_gonna_work")
 
 
@@ -460,7 +461,7 @@ def test_table_exists_translates_exception(mock_bigquery_client):
     mock_bigquery_client.get_table.side_effect = (
         google.api_core.exceptions.InternalServerError("something went wrong")
     )
-    with pytest.raises(exceptions.GenericGBQException):
+    with pytest.raises(gbq.GenericGBQException):
         connector.exists("not_gonna_work")
 
 
@@ -919,7 +920,7 @@ def test_read_gbq_with_list_rows_error_translates_exception(
         google.api_core.exceptions.NotFound("table not found"),
     )
 
-    with pytest.raises(exceptions.GenericGBQException, match="table not found"):
+    with pytest.raises(gbq.GenericGBQException, match="table not found"):
         gbq.read_gbq(
             "my-project.my_dataset.read_gbq_table",
             credentials=mock_service_account_credentials,

tests/unit/test_to_gbq.py

@@ -8,7 +8,7 @@
 from pandas import DataFrame
 import pytest
 
-from pandas_gbq import exceptions, gbq
+from pandas_gbq import gbq
 
 
 class FakeDataFrame:
@@ -53,7 +53,7 @@ def test_to_gbq_create_dataset_translates_exception(mock_bigquery_client):
         google.api_core.exceptions.InternalServerError("something went wrong")
     )
 
-    with pytest.raises(exceptions.GenericGBQException):
+    with pytest.raises(gbq.GenericGBQException):
         gbq.to_gbq(DataFrame([[1]]), "my_dataset.my_table", project_id="1234")
 
 
@@ -67,7 +67,7 @@ def test_to_gbq_load_method_translates_exception(
         "error loading data"
     )
 
-    with pytest.raises(exceptions.GenericGBQException):
+    with pytest.raises(gbq.GenericGBQException):
         gbq.to_gbq(
             DataFrame({"int_cole": [1, 2, 3]}),
             "my_dataset.my_table",
@@ -111,11 +111,11 @@ def test_to_gbq_with_if_exists_append_mismatch(mock_bigquery_client):
         "myproj.my_dataset.my_table",
         schema=(SchemaField("col_a", "INTEGER"), SchemaField("col_b", "STRING")),
     )
-    mock_bigquery_client.side_effect = exceptions.InvalidSchema(
+    mock_bigquery_client.side_effect = gbq.InvalidSchema(
         message=r"Provided Schema does not match Table *"
     )
 
-    with pytest.raises(exceptions.InvalidSchema) as exception_block:
+    with pytest.raises((gbq.InvalidSchema)) as exception_block:
         gbq.to_gbq(
             DataFrame({"col_a": [0.25, 1.5, -1.0]}),
             "my_dataset.my_table",
@@ -198,7 +198,7 @@ def test_to_gbq_with_if_exists_unknown():
     ],
 )
 def test_create_user_agent(user_agent, rfc9110_delimiter, expected):
-    from pandas_gbq.gbq_connector import create_user_agent
+    from pandas_gbq.gbq import create_user_agent
 
     result = create_user_agent(user_agent, rfc9110_delimiter)
     assert result == expected