def main(
project: Optional[str],
dataset: Optional[str],
module_path: str,
apply: bool,
validate: bool,
) -> None:
client = create_connection()
for local_table in set(find_tables(module_path)):
project = project or local_table.project
assert project, "Project has not been set."
dataset = dataset or local_table.dataset
assert dataset, "Dataset has not been set."
table_identifier = f"{project}.{dataset}.{local_table.full_table_name()}"
print(f"Checking migrations for: {table_identifier}")
try:
remote_table = client.get_table(table_identifier)
except NotFound as not_found:
table_exists_msg = f"Table does not exist in bq: {table_identifier}"
if validate:
raise Exception(table_exists_msg) from not_found
print(table_exists_msg)
if apply:
print("Creating table.")
table = Table(
table_identifier,
schema=local_table.get_schema_fields(),
)
if local_table.time_partitioning:
table.time_partitioning = local_table.time_partitioning
print(client.create_table(table))
else:
new_columns = list(
find_new_columns(local_table.get_schema_fields(),
remote_table.schema))
if new_columns:
new_columns_message = f"Found new columns: {new_columns}"
if validate:
raise Exception(new_columns_message)
print(new_columns_message)
if apply:
print("Applying changes")
remote_table.schema = local_table.get_schema_fields()
print(client.update_table(remote_table, ["schema"]))
def apply_schema_differences(
schema_diffs: _SchemaDiffs,
bigquery_client: BigQueryClient,
) -> None:
print("Applying changes...")
for table_identifier, difference in schema_diffs.items():
if isinstance(difference, MissingTable):
print("Creating table...")
table = Table(
table_identifier,
schema=difference.local_table.get_schema_fields(),
)
if difference.local_table.time_partitioning:
table.time_partitioning = difference.local_table.time_partitioning
remote_table = bigquery_client.create_table(table)
print(remote_table)
elif isinstance(difference, ExistingTable):
difference.remote_table.schema = difference.local_table.get_schema_fields(
)
print(
bigquery_client.update_table(difference.remote_table,
["schema"]))
