|
|
|
@ -18,9 +18,13 @@ else:
|
|
|
|
|
from common import Reader,Writer
|
|
|
|
|
import json
|
|
|
|
|
from google.oauth2 import service_account
|
|
|
|
|
from google.cloud import bigquery as bq
|
|
|
|
|
from multiprocessing import Lock
|
|
|
|
|
import pandas as pd
|
|
|
|
|
import numpy as np
|
|
|
|
|
import copy
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class SQLRW :
|
|
|
|
|
PROVIDERS = {"postgresql":"5432","redshift":"5432","mysql":"3306","mariadb":"3306"}
|
|
|
|
|
DRIVERS = {"postgresql":pg,"redshift":pg,"mysql":my,"mariadb":my}
|
|
|
|
@ -177,9 +181,32 @@ class SQLWriter(SQLRW,Writer):
|
|
|
|
|
pass
|
|
|
|
|
class BigQuery:
|
|
|
|
|
def __init__(self,**_args):
|
|
|
|
|
path = _args['service_key']
|
|
|
|
|
path = _args['service_key'] if 'service_key' in _args else _args['private_key']
|
|
|
|
|
self.credentials = service_account.Credentials.from_service_account_file(path)
|
|
|
|
|
|
|
|
|
|
self.dataset = _args['dataset'] if 'dataset' in _args else None
|
|
|
|
|
self.path = path
|
|
|
|
|
def meta(self,**_args):
|
|
|
|
|
"""
|
|
|
|
|
This function returns meta data for a given table or query with dataset/table properly formatted
|
|
|
|
|
:param table name of the name WITHOUT including dataset
|
|
|
|
|
:param sql sql query to be pulled,
|
|
|
|
|
"""
|
|
|
|
|
#if 'table' in _args :
|
|
|
|
|
# sql = "SELECT * from :dataset."+ _args['table']" limit 1"
|
|
|
|
|
#else:
|
|
|
|
|
# sql = _args['sql']
|
|
|
|
|
# if 'limit' not in sql.lower() :
|
|
|
|
|
# sql = sql + ' limit 1'
|
|
|
|
|
|
|
|
|
|
#sql = sql.replace(':dataset',self.dataset) if ':dataset' in args else sql
|
|
|
|
|
|
|
|
|
|
#
|
|
|
|
|
# Let us return the schema information now for a given table
|
|
|
|
|
#
|
|
|
|
|
table = _args['table']
|
|
|
|
|
client = bq.Client.from_service_account_json(self.path)
|
|
|
|
|
ref = client.dataset(self.dataset).table(table)
|
|
|
|
|
return client.get_table(ref).schema
|
|
|
|
|
class BQReader(BigQuery,Reader) :
|
|
|
|
|
def __init__(self,**_args):
|
|
|
|
|
|
|
|
|
@ -196,7 +223,8 @@ class BQReader(BigQuery,Reader) :
|
|
|
|
|
SQL = "SELECT * FROM :table ".replace(":table",table)
|
|
|
|
|
if SQL and 'limit' in _args:
|
|
|
|
|
SQL += " LIMIT "+str(_args['limit'])
|
|
|
|
|
|
|
|
|
|
if (':dataset' in SQL or ':DATASET' in SQL) and self.dataset:
|
|
|
|
|
SQL = SQL.replace(':dataset',self.dataset).replace(':DATASET',self.dataset)
|
|
|
|
|
return pd.read_gbq(SQL,credentials=self.credentials,dialect='standard') if SQL else None
|
|
|
|
|
class BQWriter(BigQuery,Writer):
|
|
|
|
|
Lock = Lock()
|
|
|
|
@ -223,7 +251,14 @@ class BQWriter(BigQuery,Writer):
|
|
|
|
|
elif type(_info) == pd.DataFrame :
|
|
|
|
|
_df = _info
|
|
|
|
|
|
|
|
|
|
self.mode['destination_table'] = _args['table'].strip()
|
|
|
|
|
if '.' not in _args['table'] :
|
|
|
|
|
self.mode['destination_table'] = '.'.join([self.dataset,_args['table']])
|
|
|
|
|
else:
|
|
|
|
|
|
|
|
|
|
self.mode['destination_table'] = _args['table'].strip()
|
|
|
|
|
if 'schema' in _args :
|
|
|
|
|
self.mode['table_schema'] = _args['schema']
|
|
|
|
|
_mode = copy.deepcopy(self.mode)
|
|
|
|
|
_df.to_gbq(**self.mode) #if_exists='append',destination_table=partial,credentials=credentials,chunksize=90000)
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
@ -237,4 +272,4 @@ class BQWriter(BigQuery,Writer):
|
|
|
|
|
# # # w = SQLWriter(**_args)
|
|
|
|
|
# # # w.write({"name":"kalara.io","email":"ceo@kalara.io","age":10})
|
|
|
|
|
# r = SQLReader(**_args)
|
|
|
|
|
# print (r.read(filter='age > 0',limit = 20))
|
|
|
|
|
# print (r.read(filter='age > 0',limit = 20))
|
|
|
|
|