def read_table(
dataset_id,
table_id,
query_project_id="basedosdados",
billing_project_id=None,
limit=None,
from_file=False,
reauth=False,
):
"""Load data from BigQuery using dataset_id and table_id.
Args:
dataset_id (str): Optional.
Dataset id available in basedosdados. It should always come with table_id.
table_id (str): Optional.
Table id available in basedosdados.dataset_id.
It should always come with dataset_id.
query_project_id (str): Optional.
Which project the table lives. You can change this you want to query different projects.
billing_project_id (str): Optional.
Project that will be billed. Find your Project ID here https://console.cloud.google.com/projectselector2/home/dashboard
reauth (boolean): Optional.
Re-authorize Google Cloud Project in case you need to change user or reset configurations.
limit (int): Optional.
Number of rows to read from table.
Returns:
pd.DataFrame:
Query result
"""
if (dataset_id is not None) and (table_id is not None):
query = f"""
SELECT *
FROM `{query_project_id}.{dataset_id}.{table_id}`"""
if limit is not None:
query += f" LIMIT {limit}"
else:
raise BaseDosDadosException(
"Both table_id and dataset_id should be filled.")
return read_sql(query,
billing_project_id=billing_project_id,
from_file=from_file,
reauth=reauth)
def download(
savepath,
query=None,
dataset_id=None,
table_id=None,
query_project_id="basedosdados",
billing_project_id=None,
limit=None,
from_file=False,
reauth=False,
**pandas_kwargs,
):
"""Download table or query result from basedosdados BigQuery (or other).
* Using a **query**:
`download('select * from `basedosdados.br_suporte.diretorio_municipios` limit 10')`
* Using **dataset_id & table_id**:
`download(dataset_id='br_suporte', table_id='diretorio_municipios')`
You can also add arguments to modify save parameters:
`download(dataset_id='br_suporte', table_id='diretorio_municipios', index=False, sep='|')`
Args:
savepath (str, pathlib.PosixPath):
If savepath is a folder, it saves a file as `savepath / table_id.csv` or
`savepath / query_result.csv` if table_id not available.
If savepath is a file, saves data to file.
query (str): Optional.
Valid SQL Standard Query to basedosdados. If query is available,
dataset_id and table_id are not required.
dataset_id (str): Optional.
Dataset id available in basedosdados. It should always come with table_id.
table_id (str): Optional.
Table id available in basedosdados.dataset_id.
It should always come with dataset_id.
query_project_id (str): Optional.
Which project the table lives. You can change this you want to query different projects.
billing_project_id (str): Optional.
Project that will be billed. Find your Project ID here https://console.cloud.google.com/projectselector2/home/dashboard
limit (int): Optional
Number of rows.
reauth (boolean): Optional.
Re-authorize Google Cloud Project in case you need to change user or reset configurations.
pandas_kwargs ():
Extra arguments accepted by [pandas.to_csv](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.to_csv.html)
Raises:
Exception: If either table_id or dataset_id were are empty.
"""
savepath = Path(savepath)
# make sure that path exists
if savepath.is_dir():
savepath.mkdir(parents=True, exist_ok=True)
else:
savepath.parent.mkdir(parents=True, exist_ok=True)
if (dataset_id is not None) and (table_id is not None):
table = read_table(
dataset_id,
table_id,
query_project_id=query_project_id,
billing_project_id=billing_project_id,
limit=limit,
reauth=reauth,
from_file=from_file,
)
elif query is not None:
query += f" limit {limit}" if limit is not None else ""
table = read_sql(
query,
billing_project_id=billing_project_id,
from_file=from_file,
reauth=reauth,
)
elif query is None:
raise BaseDosDadosException(
"Either table_id, dataset_id or query should be filled.")
if savepath.is_dir():
if table_id is not None:
savepath = savepath / (table_id + ".csv")
else:
savepath = savepath / ("query_result.csv")
pandas_kwargs["index"] = pandas_kwargs.get("index", False)
table.to_csv(savepath, **pandas_kwargs)
def create(
self,
path=None,
job_config_params=None,
force_dataset=True,
if_table_exists="raise",
if_storage_data_exists="raise",
if_table_config_exists="raise",
source_format="csv",
columns_config_url=None,
):
"""Creates BigQuery table at staging dataset.
If you add a path, it automatically saves the data in the storage,
creates a datasets folder and BigQuery location, besides creating the
table and its configuration files.
The new table should be located at `<dataset_id>_staging.<table_id>` in BigQuery.
It looks for data saved in Storage at `<bucket_name>/staging/<dataset_id>/<table_id>/*`
and builds the table.
It currently supports the types:
- Comma Delimited CSV
Data can also be partitioned following the hive partitioning scheme
`<key1>=<value1>/<key2>=<value2>` - for instance,
`year=2012/country=BR`. The partition is automatcally detected
by searching for `partitions` on the `table_config.yaml`.
Args:
path (str or pathlib.PosixPath): Where to find the file that you want to upload to create a table with
job_config_params (dict): Optional.
Job configuration params from bigquery
if_table_exists (str): Optional
What to do if table exists
* 'raise' : Raises Conflict exception
* 'replace' : Replace table
* 'pass' : Do nothing
force_dataset (bool): Creates `<dataset_id>` folder and BigQuery Dataset if it doesn't exists.
if_table_config_exists (str): Optional.
What to do if config files already exist
* 'raise': Raises FileExistError
* 'replace': Replace with blank template
* 'pass'; Do nothing
if_storage_data_exists (str): Optional.
What to do if data already exists on your bucket:
* 'raise' : Raises Conflict exception
* 'replace' : Replace table
* 'pass' : Do nothing
source_format (str): Optional
Data source format. Only 'csv' is supported. Defaults to 'csv'.
columns_config_url (str): google sheets URL.
The URL must be in the format https://docs.google.com/spreadsheets/d/<table_key>/edit#gid=<table_gid>.
The sheet must contain the column name: "coluna" and column description: "descricao"
"""
if path is None:
# Look if table data already exists at Storage
data = self.client["storage_staging"].list_blobs(
self.bucket_name,
prefix=f"staging/{self.dataset_id}/{self.table_id}")
# Raise: Cannot create table without external data
if not data:
raise BaseDosDadosException(
"You must provide a path for uploading data")
# Add data to storage
if isinstance(
path,
(
str,
Path,
),
):
Storage(self.dataset_id, self.table_id,
**self.main_vars).upload(path,
mode="staging",
if_exists=if_storage_data_exists)
# Create Dataset if it doesn't exist
if force_dataset:
dataset_obj = Dataset(self.dataset_id, **self.main_vars)
try:
dataset_obj.init()
except FileExistsError:
pass
dataset_obj.create(if_exists="pass")
self.init(
data_sample_path=path,
if_folder_exists="replace",
if_table_config_exists=if_table_config_exists,
columns_config_url=columns_config_url,
)
table = bigquery.Table(self.table_full_name["staging"])
table.external_data_configuration = Datatype(
self, source_format, "staging",
partitioned=self._is_partitioned()).external_config
# Lookup if table alreay exists
table_ref = None
try:
table_ref = self.client["bigquery_staging"].get_table(
self.table_full_name["staging"])
except google.api_core.exceptions.NotFound:
pass
if isinstance(table_ref, google.cloud.bigquery.table.Table):
if if_table_exists == "pass":
return None
elif if_table_exists == "raise":
raise FileExistsError(
"Table already exists, choose replace if you want to overwrite it"
)
if if_table_exists == "replace":
self.delete(mode="staging")
self.client["bigquery_staging"].create_table(table)
def read_sql(query, billing_project_id=None, from_file=False, reauth=False):
"""Load data from BigQuery using a query. Just a wrapper around pandas.read_gbq
Args:
query (sql):
Valid SQL Standard Query to basedosdados
billing_project_id (str): Optional.
Project that will be billed. Find your Project ID here https://console.cloud.google.com/projectselector2/home/dashboard
reauth (boolean): Optional.
Re-authorize Google Cloud Project in case you need to change user or reset configurations.
Returns:
pd.DataFrame:
Query result
"""
try:
# Set a two hours timeout
bigquery_storage_v1.client.BigQueryReadClient.read_rows = partialmethod(
bigquery_storage_v1.client.BigQueryReadClient.read_rows,
timeout=3600 * 2,
)
return pandas_gbq.read_gbq(
query,
credentials=credentials(from_file=from_file, reauth=reauth),
project_id=billing_project_id,
)
except (OSError, ValueError):
raise BaseDosDadosException(
"\nWe are not sure which Google Cloud project should be billed.\n"
"First, you should make sure that you have a Google Cloud project.\n"
"If you don't have one, set one up following these steps: \n"
"\t1. Go to this link https://console.cloud.google.com/projectselector2/home/dashboard\n"
"\t2. Agree with Terms of Service if asked\n"
"\t3. Click in Create Project\n"
"\t4. Put a cool name in your project\n"
"\t5. Hit create\n"
""
"Copy the Project ID, (notice that it is not the Project Name)\n"
"Now, you have two options:\n"
"1. Add an argument to your function poiting to the billing project id.\n"
" Like `bd.read_table('br_ibge_pib', 'municipios', billing_project_id=<YOUR_PROJECT_ID>)`\n"
"2. You can set a project_id in the environment by running the following command in your terminal: `gcloud config set project <YOUR_PROJECT_ID>`."
" Bear in mind that you need `gcloud` installed.")
except GenericGBQException as e:
if "Reason: 403" in str(e):
raise BaseDosDadosException(
"\nYou still don't have a Google Cloud Project.\n"
"Set one up following these steps: \n"
"1. Go to this link https://console.cloud.google.com/projectselector2/home/dashboard\n"
"2. Agree with Terms of Service if asked\n"
"3. Click in Create Project\n"
"4. Put a cool name in your project\n"
"5. Hit create\n"
"6. Rerun this command with the flag `reauth=True`. \n"
" Like `read_table('br_ibge_pib', 'municipios', reauth=True)`"
)
else:
raise e
def init(
self,
data_sample_path=None,
if_folder_exists="raise",
if_table_config_exists="raise",
source_format="csv",
columns_config_url=None,
):
"""Initialize table folder at metadata_path at `metadata_path/<dataset_id>/<table_id>`.
The folder should contain:
* `table_config.yaml`
* `publish.sql`
You can also point to a sample of the data to auto complete columns names.
Args:
data_sample_path (str, pathlib.PosixPath): Optional.
Data sample path to auto complete columns names
It supports Comma Delimited CSV.
if_folder_exists (str): Optional.
What to do if table folder exists
* 'raise' : Raises FileExistsError
* 'replace' : Replace folder
* 'pass' : Do nothing
if_table_config_exists (str): Optional
What to do if table_config.yaml and publish.sql exists
* 'raise' : Raises FileExistsError
* 'replace' : Replace files with blank template
* 'pass' : Do nothing
source_format (str): Optional
Data source format. Only 'csv' is supported. Defaults to 'csv'.
columns_config_url (str): google sheets URL.
The URL must be in the format https://docs.google.com/spreadsheets/d/<table_key>/edit#gid=<table_gid>.
The sheet must contain the column name: "coluna" and column description: "descricao"
Raises:
FileExistsError: If folder exists and replace is False.
NotImplementedError: If data sample is not in supported type or format.
"""
if not self.dataset_folder.exists():
raise FileExistsError(
f"Dataset folder {self.dataset_folder} folder does not exists. "
"Create a dataset before adding tables.")
try:
self.table_folder.mkdir(exist_ok=(if_folder_exists == "replace"))
except FileExistsError:
if if_folder_exists == "raise":
raise FileExistsError(
f"Table folder already exists for {self.table_id}. ")
elif if_folder_exists == "pass":
return self
if not data_sample_path and if_table_config_exists != "pass":
raise BaseDosDadosException(
"You must provide a path to correctly create config files")
partition_columns = []
if isinstance(
data_sample_path,
(
str,
Path,
),
):
# Check if partitioned and get data sample and partition columns
data_sample_path = Path(data_sample_path)
if data_sample_path.is_dir():
data_sample_path = [
f for f in data_sample_path.glob("**/*")
if f.is_file() and f.suffix == ".csv"
][0]
partition_columns = [
k.split("=")[0]
for k in data_sample_path.as_posix().split("/") if "=" in k
]
columns = Datatype(self, source_format).header(data_sample_path)
else:
columns = ["column_name"]
if if_table_config_exists == "pass":
# Check if config files exists before passing
if (Path(self.table_folder / "table_config.yaml").is_file()
and Path(self.table_folder / "publish.sql").is_file()):
pass
# Raise if no sample to determine columns
elif not data_sample_path:
raise BaseDosDadosException(
"You must provide a path to correctly create config files")
else:
self._make_template(columns, partition_columns)
elif if_table_config_exists == "raise":
# Check if config files already exist
if (Path(self.table_folder / "table_config.yaml").is_file()
and Path(self.table_folder / "publish.sql").is_file()):
raise FileExistsError(
f"table_config.yaml and publish.sql already exists at {self.table_folder}"
)
# if config files don't exist, create them
else:
self._make_template(columns, partition_columns)
else:
# Raise: without a path to data sample, should not replace config files with empty template
self._make_template(columns, partition_columns)
if columns_config_url is not None:
self.update_columns(columns_config_url)
return self