<s> from base64 import b64decode <EOL> import datetime <EOL> import json <EOL> import httplib2 <EOL> import logging <EOL> import sys <EOL> import time <EOL> import requests <EOL> from redash import settings <EOL> from redash . query_runner import * <EOL> from redash . utils import JSONEncoder <EOL> logger = logging . getLogger ( __name__ ) <EOL> try : <EOL> import apiclient . errors <EOL> from apiclient . discovery import build <EOL> from apiclient . errors import HttpError <EOL> from oauth2client . client import SignedJwtAssertionCredentials <EOL> from oauth2client import gce <EOL> enabled = True <EOL> except ImportError : <EOL> enabled = False <EOL> types_map = { <EOL> '<STR_LIT>' : TYPE_INTEGER , <EOL> '<STR_LIT>' : TYPE_FLOAT , <EOL> '<STR_LIT>' : TYPE_BOOLEAN , <EOL> '<STR_LIT>' : TYPE_STRING , <EOL> '<STR_LIT>' : TYPE_DATETIME , <EOL> } <EOL> def transform_row ( row , fields ) : <EOL> column_index = <NUM_LIT:0> <EOL> row_data = { } <EOL> for cell in row [ "<STR_LIT:f>" ] : <EOL> field = fields [ column_index ] <EOL> cell_value = cell [ '<STR_LIT:v>' ] <EOL> if cell_value is None : <EOL> pass <EOL> elif field [ '<STR_LIT:type>' ] == '<STR_LIT>' : <EOL> cell_value = int ( cell_value ) <EOL> elif field [ '<STR_LIT:type>' ] == '<STR_LIT>' : <EOL> cell_value = float ( cell_value ) <EOL> elif field [ '<STR_LIT:type>' ] == '<STR_LIT>' : <EOL> cell_value = cell_value . lower ( ) == "<STR_LIT:true>" <EOL> elif field [ '<STR_LIT:type>' ] == '<STR_LIT>' : <EOL> cell_value = datetime . datetime . fromtimestamp ( float ( cell_value ) ) <EOL> row_data [ field [ "<STR_LIT:name>" ] ] = cell_value <EOL> column_index += <NUM_LIT:1> <EOL> return row_data <EOL> def _load_key ( filename ) : <EOL> f = file ( filename , "<STR_LIT:rb>" ) <EOL> try : <EOL> return f . read ( ) <EOL> finally : <EOL> f . close ( ) <EOL> def _get_query_results ( jobs , project_id , job_id , start_index ) : <EOL> query_reply = jobs . getQueryResults ( projectId = project_id , jobId = job_id , startIndex = start_index ) . execute ( ) <EOL> logging . debug ( '<STR_LIT>' , query_reply ) <EOL> if not query_reply [ '<STR_LIT>' ] : <EOL> time . sleep ( <NUM_LIT:10> ) <EOL> return _get_query_results ( jobs , project_id , job_id , start_index ) <EOL> return query_reply <EOL> class BigQuery ( BaseQueryRunner ) : <EOL> @ classmethod <EOL> def enabled ( cls ) : <EOL> return enabled <EOL> @ classmethod <EOL> def configuration_schema ( cls ) : <EOL> return { <EOL> '<STR_LIT:type>' : '<STR_LIT:object>' , <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT:type>' : '<STR_LIT:string>' , <EOL> '<STR_LIT:title>' : '<STR_LIT>' <EOL> } , <EOL> '<STR_LIT>' : { <EOL> "<STR_LIT:type>" : "<STR_LIT:string>" , <EOL> '<STR_LIT:title>' : '<STR_LIT>' <EOL> } , <EOL> '<STR_LIT>' : { <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> '<STR_LIT:title>' : '<STR_LIT>' <EOL> } <EOL> } , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' ] <EOL> } <EOL> def __init__ ( self , configuration ) : <EOL> super ( BigQuery , self ) . __init__ ( configuration ) <EOL> def _get_bigquery_service ( self ) : <EOL> scope = [ <EOL> "<STR_LIT>" , <EOL> ] <EOL> key = json . loads ( b64decode ( self . configuration [ '<STR_LIT>' ] ) ) <EOL> credentials = SignedJwtAssertionCredentials ( key [ '<STR_LIT>' ] , key [ '<STR_LIT>' ] , scope = scope ) <EOL> http = httplib2 . Http ( timeout = settings . BIGQUERY_HTTP_TIMEOUT ) <EOL> http = credentials . authorize ( http ) <EOL> return build ( "<STR_LIT>" , "<STR_LIT>" , http = http ) <EOL> def _get_project_id ( self ) : <EOL> return self . configuration [ "<STR_LIT>" ] <EOL> def _get_total_bytes_processed ( self , jobs , query ) : <EOL> job_data = { <EOL> "<STR_LIT>" : query , <EOL> "<STR_LIT>" : True , <EOL> } <EOL> response = jobs . query ( projectId = self . _get_project_id ( ) , body = job_data ) . execute ( ) <EOL> return int ( response [ "<STR_LIT>" ] ) <EOL> def _get_query_result ( self , jobs , query ) : <EOL> project_id = self . _get_project_id ( ) <EOL> job_data = { <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : query , <EOL> } <EOL> } <EOL> } <EOL> insert_response = jobs . insert ( projectId = project_id , body = job_data ) . execute ( ) <EOL> current_row = <NUM_LIT:0> <EOL> query_reply = _get_query_results ( jobs , project_id = project_id , <EOL> job_id = insert_response [ '<STR_LIT>' ] [ '<STR_LIT>' ] , start_index = current_row ) <EOL> logger . debug ( "<STR_LIT>" , query_reply ) <EOL> rows = [ ] <EOL> while ( "<STR_LIT>" in query_reply ) and current_row < query_reply [ '<STR_LIT>' ] : <EOL> for row in query_reply [ "<STR_LIT>" ] : <EOL> rows . append ( transform_row ( row , query_reply [ "<STR_LIT>" ] [ "<STR_LIT>" ] ) ) <EOL> current_row += len ( query_reply [ '<STR_LIT>' ] ) <EOL> query_reply = jobs . getQueryResults ( projectId = project_id , jobId = query_reply [ '<STR_LIT>' ] [ '<STR_LIT>' ] , <EOL> startIndex = current_row ) . execute ( ) <EOL> columns = [ { '<STR_LIT:name>' : f [ "<STR_LIT:name>" ] , <EOL> '<STR_LIT>' : f [ "<STR_LIT:name>" ] , <EOL> '<STR_LIT:type>' : types_map . get ( f [ '<STR_LIT:type>' ] , "<STR_LIT:string>" ) } for f in query_reply [ "<STR_LIT>" ] [ "<STR_LIT>" ] ] <EOL> data = { <EOL> "<STR_LIT>" : columns , <EOL> "<STR_LIT>" : rows <EOL> } <EOL> return data <EOL> def run_query ( self , query ) : <EOL> logger . debug ( "<STR_LIT>" , query ) <EOL> bigquery_service = self . _get_bigquery_service ( ) <EOL> jobs = bigquery_service . jobs ( ) <EOL> try : <EOL> if "<STR_LIT>" in self . configuration : <EOL> limitMB = self . configuration [ "<STR_LIT>" ] <EOL> processedMB = self . _get_total_bytes_processed ( jobs , query ) / <NUM_LIT> / <NUM_LIT> <EOL> if limitMB < processedMB : <EOL> return None , "<STR_LIT>" % ( limitMB , processedMB ) <EOL> data = self . _get_query_result ( jobs , query ) <EOL> error = None <EOL> json_data = json . dumps ( data , cls = JSONEncoder ) <EOL> except apiclient . errors . HttpError , e : <EOL> json_data = None <EOL> if e . resp . status == <NUM_LIT> : <EOL> error = json . loads ( e . content ) [ '<STR_LIT:error>' ] [ '<STR_LIT:message>' ] <EOL> else : <EOL> error = e . content <EOL> except KeyboardInterrupt : <EOL> error = "<STR_LIT>" <EOL> json_data = None <EOL> except Exception : <EOL> raise sys . exc_info ( ) [ <NUM_LIT:1> ] , None , sys . exc_info ( ) [ <NUM_LIT:2> ] <EOL> return json_data , error <EOL> class BigQueryGCE ( BigQuery ) : <EOL> @ classmethod <EOL> def type ( cls ) : <EOL> return "<STR_LIT>" <EOL> @ classmethod <EOL> def configuration_schema ( cls ) : <EOL> return { } <EOL> def _get_project_id ( self ) : <EOL> return requests . get ( '<STR_LIT>' , headers = { '<STR_LIT>' : '<STR_LIT>' } ) . content <EOL> def _get_bigquery_service ( self ) : <EOL> credentials = gce . AppAssertionCredentials ( scope = '<STR_LIT>' ) <EOL> http = httplib2 . Http ( ) <EOL> http = credentials . authorize ( http ) <EOL> return build ( "<STR_LIT>" , "<STR_LIT>" , http = http ) <EOL> register ( BigQuery ) <EOL> register ( BigQueryGCE ) </s>
