EGxW9CtM

1from google.cloud import bigquery
2
3"""
4Use this to convert an online BigQuery dataset's table schema into a PostgreSQL (eg. on CloudSQL) CREATE TABLE command.
5"""
6
7
8def get_table_schema(project_name, dataset_name, table_name):
9    """Return 'project_name.dataset_name.table_name''s BigQuery table schema object."""
10    bq_client = bigquery.Client(project=project_name)
11
12    dataset = bq_client.dataset(dataset_name)
13    table_ref = dataset.table(table_name)
14
15    table = bq_client.get_table(table_ref)
16    return table.schema
17
18
19def bq_to_pg_field_type(field_type):
20    """Return an equivalent variable type from BigQuery to PostgreSQL
21    TODO: support more types.
22    """
23    if field_type == 'STRING':
24        return 'TEXT'
25    if field_type == 'INTEGER':
26        return 'INTEGER'
27    if field_type == 'FLOAT':
28        return 'REAL'
29    else:
30        raise TypeError(
31            'UNKNOWN TYPE: {}. Unable to convert that BigQuery field type to PostgreSQL. Giving up.'.format(field_type))
32
33
34def bq_schema_to_pg_query(bq_schema, destination_table_name, primary_keys, add_if_not_exists=True):
35    """Return a PostgreSQL CREATE TABLE query as string from a given BigQuery schema object."""
36    pg_query = "CREATE TABLE {} {} ({}, PRIMARY KEY({}));"
37    pg_fields_definition = []
38
39    for field in bq_schema:
40        pg_fields_definition.append(field.name + ' ' + bq_to_pg_field_type(field.field_type))
41
42    return pg_query.format("IF NOT EXISTS" if add_if_not_exists else "", destination_table_name,
43                           ", ".join(pg_fields_definition), ", ".join(primary_keys))
44
45
46def get_pg_schema_from_bq_table(project_name, dataset_name, table_name, destination_table_name=None, primary_keys=[],
47                                add_if_not_exists=True):
48    """Return a PostgreSQL CREATE Table query string, from a given BigQuery dataset's table attributes.
49    If destination_table_name is omitted, defaults to table_name.
50    """
51    if destination_table_name is None:
52        destination_table_name = table_name
53    if not isinstance(primary_keys, (list, tuple)):
54        raise TypeError('primary_keys argument must be iterable')
55    schema = get_table_schema(project_name, dataset_name, table_name)
56    return bq_schema_to_pg_query(schema, destination_table_name, primary_keys, add_if_not_exists)
57
58
59if __name__ == '__main__':
60    """
61    Example:
62    
63    print get_pg_schema_from_bq_table('some_project', 'dataset', 'table', 'dest_table', ['my_primary_key'], True)
64    
65    Yields:
66    
67    CREATE TABLE IF NOT EXISTS dest_table 
68    (my_primary_key INTEGER, some_field TEXT, another_field REAL, PRIMARY KEY(my_primary_key));
69    """
70    pass