python-bigquery-dataframes/samples/snippets/conftest.py at e90ac46847e54bab2f7061b84c3e8225b47da5d9 · googleapis/python-bigquery-dataframes · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Generator, Iterator

from google.cloud import bigquery, storage
import pytest
import test_utils.prefixer

import bigframes.pandas as bpd

prefixer = test_utils.prefixer.Prefixer(
    "python-bigquery-dataframes", "samples/snippets"
)

routine_prefixer = test_utils.prefixer.Prefixer("bigframes", "")


@pytest.fixture(scope="session", autouse=True)
def cleanup_datasets(bigquery_client: bigquery.Client) -> None:
    for dataset in bigquery_client.list_datasets():
        if prefixer.should_cleanup(dataset.dataset_id):
            bigquery_client.delete_dataset(
                dataset, delete_contents=True, not_found_ok=True
            )


@pytest.fixture(scope="session")
def bigquery_client() -> bigquery.Client:
    bigquery_client = bigquery.Client()
    return bigquery_client


@pytest.fixture(scope="session")
def storage_client(project_id: str) -> storage.Client:
    return storage.Client(project=project_id)


@pytest.fixture(scope="session")
def project_id(bigquery_client: bigquery.Client) -> str:
    return bigquery_client.project


@pytest.fixture(scope="session")
def gcs_bucket(storage_client: storage.Client) -> Generator[str, None, None]:
    bucket_name = "bigframes_blob_test_with_data_wipeout"

    yield bucket_name

    bucket = storage_client.get_bucket(bucket_name)
    for blob in bucket.list_blobs():
        blob.delete()


@pytest.fixture(scope="session")
def gcs_bucket_snippets(storage_client: storage.Client) -> Generator[str, None, None]:
    bucket_name = "bigframes_blob_test_snippet_with_data_wipeout"

    yield bucket_name

    bucket = storage_client.get_bucket(bucket_name)
    for blob in bucket.list_blobs():
        blob.delete()


@pytest.fixture(autouse=True)
def reset_session() -> None:
    """An autouse fixture ensuring each sample runs in a fresh session.

    This allows us to have samples that query data in different locations.
    """
    bpd.reset_session()
    bpd.options.bigquery.location = None


@pytest.fixture(scope="session")
def dataset_id(bigquery_client: bigquery.Client, project_id: str) -> Iterator[str]:
    dataset_id = prefixer.create_prefix()
    full_dataset_id = f"{project_id}.{dataset_id}"
    dataset = bigquery.Dataset(full_dataset_id)
    bigquery_client.create_dataset(dataset)
    yield dataset_id
    bigquery_client.delete_dataset(dataset, delete_contents=True, not_found_ok=True)


@pytest.fixture(scope="session")
def dataset_id_eu(bigquery_client: bigquery.Client, project_id: str) -> Iterator[str]:
    dataset_id = prefixer.create_prefix()
    full_dataset_id = f"{project_id}.{dataset_id}"
    dataset = bigquery.Dataset(full_dataset_id)
    dataset.location = "EU"
    bigquery_client.create_dataset(dataset)
    yield dataset_id
    bigquery_client.delete_dataset(dataset, delete_contents=True, not_found_ok=True)


@pytest.fixture
def random_model_id(
    bigquery_client: bigquery.Client, project_id: str, dataset_id: str
) -> Iterator[str]:
    """Create a new table ID each time, so random_model_id can be used as
    target for load jobs.
    """
    random_model_id = prefixer.create_prefix()
    full_model_id = f"{project_id}.{dataset_id}.{random_model_id}"
    yield full_model_id
    bigquery_client.delete_model(full_model_id, not_found_ok=True)


@pytest.fixture
def random_model_id_eu(
    bigquery_client: bigquery.Client, project_id: str, dataset_id_eu: str
) -> Iterator[str]:
    """
    Create a new table ID each time, so random_model_id_eu can be used
    as a target for load jobs.
    """
    random_model_id_eu = prefixer.create_prefix()
    full_model_id = f"{project_id}.{dataset_id_eu}.{random_model_id_eu}"
    yield full_model_id
    bigquery_client.delete_model(full_model_id, not_found_ok=True)


@pytest.fixture
def routine_id() -> Iterator[str]:
    """Create a new BQ routine ID each time, so random_routine_id can be used as
    target for udf creation.
    """
    random_routine_id = routine_prefixer.create_prefix()
    yield random_routine_id