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"]))