| Current File : //usr/local/lib64/python3.6/site-packages/pandas/tests/io/test_gbq.py |
from contextlib import ExitStack as does_not_raise
from datetime import datetime
import os
import platform
import random
import string
import numpy as np
import pytest
import pytz
import pandas as pd
from pandas import DataFrame
api_exceptions = pytest.importorskip("google.api_core.exceptions")
bigquery = pytest.importorskip("google.cloud.bigquery")
service_account = pytest.importorskip("google.oauth2.service_account")
pandas_gbq = pytest.importorskip("pandas_gbq")
PROJECT_ID = None
PRIVATE_KEY_JSON_PATH = None
PRIVATE_KEY_JSON_CONTENTS = None
VERSION = platform.python_version()
def _skip_if_no_project_id():
if not _get_project_id():
pytest.skip("Cannot run integration tests without a project id")
def _skip_if_no_private_key_path():
if not _get_private_key_path():
pytest.skip("Cannot run integration tests without a private key json file path")
def _in_travis_environment():
return "TRAVIS_BUILD_DIR" in os.environ and "GBQ_PROJECT_ID" in os.environ
def _get_project_id():
if _in_travis_environment():
return os.environ.get("GBQ_PROJECT_ID")
return PROJECT_ID or os.environ.get("GBQ_PROJECT_ID")
def _get_private_key_path():
if _in_travis_environment():
return os.path.join(
*[os.environ.get("TRAVIS_BUILD_DIR"), "ci", "travis_gbq.json"]
)
private_key_path = PRIVATE_KEY_JSON_PATH
if not private_key_path:
private_key_path = os.environ.get("GBQ_GOOGLE_APPLICATION_CREDENTIALS")
return private_key_path
def _get_credentials():
private_key_path = _get_private_key_path()
if private_key_path:
return service_account.Credentials.from_service_account_file(private_key_path)
def _get_client():
project_id = _get_project_id()
credentials = _get_credentials()
return bigquery.Client(project=project_id, credentials=credentials)
def generate_rand_str(length: int = 10) -> str:
return "".join(random.choices(string.ascii_lowercase, k=length))
def make_mixed_dataframe_v2(test_size):
# create df to test for all BQ datatypes except RECORD
bools = np.random.randint(2, size=(1, test_size)).astype(bool)
flts = np.random.randn(1, test_size)
ints = np.random.randint(1, 10, size=(1, test_size))
strs = np.random.randint(1, 10, size=(1, test_size)).astype(str)
times = [datetime.now(pytz.timezone("US/Arizona")) for t in range(test_size)]
return DataFrame(
{
"bools": bools[0],
"flts": flts[0],
"ints": ints[0],
"strs": strs[0],
"times": times[0],
},
index=range(test_size),
)
def test_read_gbq_without_deprecated_kwargs(monkeypatch):
captured_kwargs = {}
def mock_read_gbq(sql, **kwargs):
captured_kwargs.update(kwargs)
return DataFrame([[1.0]])
monkeypatch.setattr("pandas_gbq.read_gbq", mock_read_gbq)
pd.read_gbq("SELECT 1")
assert "verbose" not in captured_kwargs
assert "private_key" not in captured_kwargs
def test_read_gbq_with_new_kwargs(monkeypatch):
captured_kwargs = {}
def mock_read_gbq(sql, **kwargs):
captured_kwargs.update(kwargs)
return DataFrame([[1.0]])
monkeypatch.setattr("pandas_gbq.read_gbq", mock_read_gbq)
pd.read_gbq("SELECT 1", use_bqstorage_api=True, max_results=1)
assert captured_kwargs["use_bqstorage_api"]
assert captured_kwargs["max_results"]
def test_read_gbq_without_new_kwargs(monkeypatch):
captured_kwargs = {}
def mock_read_gbq(sql, **kwargs):
captured_kwargs.update(kwargs)
return DataFrame([[1.0]])
monkeypatch.setattr("pandas_gbq.read_gbq", mock_read_gbq)
pd.read_gbq("SELECT 1")
assert "use_bqstorage_api" not in captured_kwargs
assert "max_results" not in captured_kwargs
@pytest.mark.parametrize("progress_bar", [None, "foo"])
def test_read_gbq_progress_bar_type_kwarg(monkeypatch, progress_bar):
# GH 29857
captured_kwargs = {}
def mock_read_gbq(sql, **kwargs):
captured_kwargs.update(kwargs)
return DataFrame([[1.0]])
monkeypatch.setattr("pandas_gbq.read_gbq", mock_read_gbq)
pd.read_gbq("SELECT 1", progress_bar_type=progress_bar)
assert "progress_bar_type" in captured_kwargs
@pytest.mark.single
class TestToGBQIntegrationWithServiceAccountKeyPath:
@pytest.fixture()
def gbq_dataset(self):
# Setup Dataset
_skip_if_no_project_id()
_skip_if_no_private_key_path()
dataset_id = "pydata_pandas_bq_testing_" + generate_rand_str()
self.client = _get_client()
self.dataset = self.client.dataset(dataset_id)
# Create the dataset
self.client.create_dataset(bigquery.Dataset(self.dataset))
table_name = generate_rand_str()
destination_table = f"{dataset_id}.{table_name}"
yield destination_table
# Teardown Dataset
self.client.delete_dataset(self.dataset, delete_contents=True)
def test_roundtrip(self, gbq_dataset):
destination_table = gbq_dataset
test_size = 20001
df = make_mixed_dataframe_v2(test_size)
df.to_gbq(
destination_table,
_get_project_id(),
chunksize=None,
credentials=_get_credentials(),
)
result = pd.read_gbq(
f"SELECT COUNT(*) AS num_rows FROM {destination_table}",
project_id=_get_project_id(),
credentials=_get_credentials(),
dialect="standard",
)
assert result["num_rows"][0] == test_size
@pytest.mark.parametrize(
"if_exists, expected_num_rows, expectation",
[
("append", 300, does_not_raise()),
("fail", 200, pytest.raises(pandas_gbq.gbq.TableCreationError)),
("replace", 100, does_not_raise()),
],
)
def test_gbq_if_exists(
self, if_exists, expected_num_rows, expectation, gbq_dataset
):
# GH 29598
destination_table = gbq_dataset
test_size = 200
df = make_mixed_dataframe_v2(test_size)
df.to_gbq(
destination_table,
_get_project_id(),
chunksize=None,
credentials=_get_credentials(),
)
with expectation:
df.iloc[:100].to_gbq(
destination_table,
_get_project_id(),
if_exists=if_exists,
chunksize=None,
credentials=_get_credentials(),
)
result = pd.read_gbq(
f"SELECT COUNT(*) AS num_rows FROM {destination_table}",
project_id=_get_project_id(),
credentials=_get_credentials(),
dialect="standard",
)
assert result["num_rows"][0] == expected_num_rows