Sree Kuchibhotla | 559e45b | 2016-02-19 03:02:16 -0800 | [diff] [blame] | 1 | #!/usr/bin/env python2.7 |
Craig Tiller | 6169d5f | 2016-03-31 07:46:18 -0700 | [diff] [blame] | 2 | # Copyright 2015, Google Inc. |
Sree Kuchibhotla | 559e45b | 2016-02-19 03:02:16 -0800 | [diff] [blame] | 3 | # All rights reserved. |
| 4 | # |
| 5 | # Redistribution and use in source and binary forms, with or without |
| 6 | # modification, are permitted provided that the following conditions are |
| 7 | # met: |
| 8 | # |
| 9 | # * Redistributions of source code must retain the above copyright |
| 10 | # notice, this list of conditions and the following disclaimer. |
| 11 | # * Redistributions in binary form must reproduce the above |
| 12 | # copyright notice, this list of conditions and the following disclaimer |
| 13 | # in the documentation and/or other materials provided with the |
| 14 | # distribution. |
| 15 | # * Neither the name of Google Inc. nor the names of its |
| 16 | # contributors may be used to endorse or promote products derived from |
| 17 | # this software without specific prior written permission. |
| 18 | # |
| 19 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS |
| 20 | # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT |
| 21 | # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR |
| 22 | # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT |
| 23 | # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, |
| 24 | # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT |
| 25 | # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, |
| 26 | # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY |
| 27 | # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| 28 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| 29 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| 30 | |
| 31 | import argparse |
| 32 | import json |
| 33 | import uuid |
| 34 | import httplib2 |
| 35 | |
| 36 | from apiclient import discovery |
| 37 | from apiclient.errors import HttpError |
| 38 | from oauth2client.client import GoogleCredentials |
| 39 | |
| 40 | NUM_RETRIES = 3 |
| 41 | |
| 42 | |
| 43 | def create_big_query(): |
| 44 | """Authenticates with cloud platform and gets a BiqQuery service object |
| 45 | """ |
| 46 | creds = GoogleCredentials.get_application_default() |
| 47 | return discovery.build('bigquery', 'v2', credentials=creds) |
| 48 | |
| 49 | |
| 50 | def create_dataset(biq_query, project_id, dataset_id): |
| 51 | is_success = True |
| 52 | body = { |
| 53 | 'datasetReference': { |
| 54 | 'projectId': project_id, |
| 55 | 'datasetId': dataset_id |
| 56 | } |
| 57 | } |
| 58 | |
| 59 | try: |
| 60 | dataset_req = biq_query.datasets().insert(projectId=project_id, body=body) |
| 61 | dataset_req.execute(num_retries=NUM_RETRIES) |
| 62 | except HttpError as http_error: |
| 63 | if http_error.resp.status == 409: |
| 64 | print 'Warning: The dataset %s already exists' % dataset_id |
| 65 | else: |
| 66 | # Note: For more debugging info, print "http_error.content" |
| 67 | print 'Error in creating dataset: %s. Err: %s' % (dataset_id, http_error) |
| 68 | is_success = False |
| 69 | return is_success |
| 70 | |
| 71 | |
| 72 | def create_table(big_query, project_id, dataset_id, table_id, table_schema, |
| 73 | description): |
Jan Tattermusch | 7d54db8 | 2016-04-14 16:57:45 -0700 | [diff] [blame] | 74 | fields = [{'name': field_name, |
| 75 | 'type': field_type, |
| 76 | 'description': field_description |
| 77 | } for (field_name, field_type, field_description) in table_schema] |
| 78 | return create_table2(big_query, project_id, dataset_id, table_id, |
| 79 | fields, description) |
| 80 | |
| 81 | |
| 82 | def create_table2(big_query, project_id, dataset_id, table_id, fields_schema, |
| 83 | description): |
Sree Kuchibhotla | 559e45b | 2016-02-19 03:02:16 -0800 | [diff] [blame] | 84 | is_success = True |
| 85 | |
| 86 | body = { |
| 87 | 'description': description, |
| 88 | 'schema': { |
Jan Tattermusch | 7d54db8 | 2016-04-14 16:57:45 -0700 | [diff] [blame] | 89 | 'fields': fields_schema |
Sree Kuchibhotla | 559e45b | 2016-02-19 03:02:16 -0800 | [diff] [blame] | 90 | }, |
| 91 | 'tableReference': { |
| 92 | 'datasetId': dataset_id, |
| 93 | 'projectId': project_id, |
| 94 | 'tableId': table_id |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | try: |
| 99 | table_req = big_query.tables().insert(projectId=project_id, |
| 100 | datasetId=dataset_id, |
| 101 | body=body) |
| 102 | res = table_req.execute(num_retries=NUM_RETRIES) |
| 103 | print 'Successfully created %s "%s"' % (res['kind'], res['id']) |
| 104 | except HttpError as http_error: |
| 105 | if http_error.resp.status == 409: |
| 106 | print 'Warning: Table %s already exists' % table_id |
| 107 | else: |
| 108 | print 'Error in creating table: %s. Err: %s' % (table_id, http_error) |
| 109 | is_success = False |
| 110 | return is_success |
| 111 | |
| 112 | |
| 113 | def insert_rows(big_query, project_id, dataset_id, table_id, rows_list): |
| 114 | is_success = True |
| 115 | body = {'rows': rows_list} |
| 116 | try: |
| 117 | insert_req = big_query.tabledata().insertAll(projectId=project_id, |
| 118 | datasetId=dataset_id, |
| 119 | tableId=table_id, |
| 120 | body=body) |
Sree Kuchibhotla | 559e45b | 2016-02-19 03:02:16 -0800 | [diff] [blame] | 121 | res = insert_req.execute(num_retries=NUM_RETRIES) |
Jan Tattermusch | ac4251a | 2016-04-15 14:44:59 -0700 | [diff] [blame] | 122 | if res.get('insertErrors', None): |
| 123 | print 'Error inserting rows! Response: %s' % res |
| 124 | is_success = False |
Sree Kuchibhotla | 559e45b | 2016-02-19 03:02:16 -0800 | [diff] [blame] | 125 | except HttpError as http_error: |
Jan Tattermusch | ac4251a | 2016-04-15 14:44:59 -0700 | [diff] [blame] | 126 | print 'Error inserting rows to the table %s' % table_id |
Sree Kuchibhotla | 559e45b | 2016-02-19 03:02:16 -0800 | [diff] [blame] | 127 | is_success = False |
Jan Tattermusch | ac4251a | 2016-04-15 14:44:59 -0700 | [diff] [blame] | 128 | |
Sree Kuchibhotla | 559e45b | 2016-02-19 03:02:16 -0800 | [diff] [blame] | 129 | return is_success |
| 130 | |
| 131 | |
| 132 | def sync_query_job(big_query, project_id, query, timeout=5000): |
| 133 | query_data = {'query': query, 'timeoutMs': timeout} |
| 134 | query_job = None |
| 135 | try: |
| 136 | query_job = big_query.jobs().query( |
| 137 | projectId=project_id, |
| 138 | body=query_data).execute(num_retries=NUM_RETRIES) |
| 139 | except HttpError as http_error: |
| 140 | print 'Query execute job failed with error: %s' % http_error |
| 141 | print http_error.content |
| 142 | return query_job |
| 143 | |
| 144 | # List of (column name, column type, description) tuples |
| 145 | def make_row(unique_row_id, row_values_dict): |
Sree Kuchibhotla | 2715a39 | 2016-02-24 12:01:52 -0800 | [diff] [blame] | 146 | """row_values_dict is a dictionary of column name and column value. |
Sree Kuchibhotla | 559e45b | 2016-02-19 03:02:16 -0800 | [diff] [blame] | 147 | """ |
| 148 | return {'insertId': unique_row_id, 'json': row_values_dict} |