def process(self, data: Dict[str, Any], context):
"""Check all the running jobs
Arguments:
event {Dict[str, Any]} -- data sent from the PubSub message
context {Dict[str, Any]} -- context data. unused
"""
firestore = Firestore(in_cloud=True, email=None, project=None)
documents = firestore.get_all_jobs()
for document in documents:
for T in [t for t in Type if not t.name.startswith('_')]:
config = firestore.get_report_config(T, document.id)
if config:
if config.get('dest_project'):
# authenticate against supplied project with supplied key
project = config.get('dest_project') or os.environ.get(
'GCP_PROJECT')
client_key = json.loads(
Cloud_Storage.fetch_file(
bucket=
f"{os.environ.get('GCP_PROJECT') or 'galvanic-card-234919'}-report2bq-tokens",
file=f"{config['email']}_user_token.json"))
server_key = json.loads(
Cloud_Storage.fetch_file(
bucket=
f"{os.environ.get('GCP_PROJECT') or 'galvanic-card-234919'}-report2bq-tokens",
file='client_secrets.json'))
client_key['client_id'] = (
server_key.get('web')
or server_key.get('installed')).get('client_id')
client_key['client_secret'] = (
server_key.get('web') or
server_key.get('installed')).get('client_secret')
logging.info(client_key)
creds = Credentials.from_authorized_user_info(
client_key)
bq = bigquery.Client(project=project,
credentials=creds)
else:
bq = bigquery.Client()
api_repr = document.get().to_dict()
if api_repr:
try:
job = LoadJob.from_api_repr(api_repr, bq)
job.reload()
if job.state == 'DONE':
if job.error_result:
logging.error(job.errors)
self._handle_finished(job=job,
id=document.id,
config=config,
report_type=T)
firestore.mark_import_job_complete(
document.id, job)
except Exception as e:
logging.error(
f"""Error loading job {document.id} for monitoring."""
)
break
class SA360(object):
def __init__(self,
email: str,
project: str,
append: bool = False,
infer_schema: bool = False):
self.email = email
self.project = project
self.creds = Credentials(email=email, project=project)
self.credentials = storage.Client()._credentials
self.transport = AuthorizedSession(credentials=self.credentials)
self.append = append
self.infer_schema = infer_schema
self.firestore = Firestore(email=email, project=project)
self.chunk_multiplier = int(os.environ.get('CHUNK_MULTIPLIER', 64))
self.bucket = f'{self.project}-report2bq-upload'
def _soupify(self, data: BytesIO) -> BeautifulSoup:
return BeautifulSoup(data, 'lxml')
def process(self, bucket: str, report_details: Dict[str, Any]) -> None:
input_buffer = BytesIO()
repeater = self._stream_processor(bucket=bucket,
report_details=report_details,
repeatable=False)
# old_id, report_details['id'] = report_details['id'], f'{report_details["id"]}-repeat'
# self.upload_report(bucket=bucket, report_details=report_details, input_buffer=repeater)
# report_details['id'] = old_id
def handle_offline_report(self, run_config: Dict[str, Any]) -> bool:
sa360_service = DiscoverService.get_service(Service.SA360, self.creds)
request = sa360_service.reports().get(reportId=run_config['file_id'])
try:
report = request.execute()
if report['isReportReady']:
report_config = self.firestore.get_report_config(
type=Type.SA360_RPT, id=run_config['report_id'])
csv_header, csv_types = self.read_header(report)
schema = CSVHelpers.create_table_schema(
csv_header, csv_types if self.infer_schema else None)
report_config['schema'] = schema
report_config['files'] = report['files']
if 'dest_project' in run_config:
report_config['dest_project'] = run_config['dest_project']
if 'dest_dataset' in run_config:
report_config['dest_dataset'] = run_config['dest_dataset']
if 'notify_topic' in run_config:
report_config['notifier'] = {
'topic': run_config['notify_topic'],
}
if 'notify_message' in run_config:
report_config['notifier']['message'] = run_config[
'notify_message']
# update the report details please...
self.firestore.update_document(Type.SA360_RPT,
run_config['report_id'],
report_config)
# ... then stream the file to GCS a la DV360/CM
self._stream_report_to_gcs(report_details=report_config,
run_config=run_config)
return report['isReportReady']
except Exception as e:
logging.error(
f'Report fetch error: Run {run_config["file_id"]} for report {run_config["report_id"]}'
)
return False
def read_header(self, report_config: dict) -> list:
r = urllib.request.Request(report_config['files'][0]['url'])
for header in self.creds.get_auth_headers():
r.add_header(header, self.creds.get_auth_headers()[header])
with closing(urlopen(r)) as report:
data = report.read(self.chunk_multiplier * 1024 * 1024)
bytes_io = BytesIO(data)
return CSVHelpers.get_column_types(bytes_io)
@measure_memory
def _stream_report_to_gcs(self, report_details: Dict[str, Any],
run_config: Dict[str, Any]) -> None:
"""Multi-threaded stream to GCS
Arguments:
bucket {str} -- GCS Bucket
report_details {dict} -- Report definition
"""
queue = Queue()
report_id = run_config['report_id']
chunk_size = self.chunk_multiplier * 1024 * 1024
out_file = BytesIO()
streamer = ThreadedGCSObjectStreamUpload(client=Cloud_Storage.client(),
bucket_name=self.bucket,
blob_name=f'{report_id}.csv',
chunk_size=chunk_size,
queue=queue)
streamer.start()
r = urllib.request.Request(report_details['files'][0]['url'])
for header in self.creds.get_auth_headers():
r.add_header(header, self.creds.get_auth_headers()[header])
with closing(urlopen(r)) as _report:
_downloaded = 0
chunk_id = 1
_report_size = int(_report.headers['content-length'])
while _downloaded < _report_size:
chunk = _report.read(chunk_size)
_downloaded += len(chunk)
queue.put((chunk_id, chunk))
chunk_id += 1
queue.join()
streamer.stop()
@timeit
def _fetch_data(self, report_details: Dict[str, Any],
buffer: BytesIO) -> int:
try:
report_url = report_details['url']
request = requests.Download(report_url, stream=buffer)
request.consume(transport=self.transport)
return self._stream_size(buffer)
except Exception as e:
logging.error(e)
return -1
@timeit
def _stream_size(self, buffer: BytesIO) -> int:
pos = buffer.tell()
buffer.seek(0, SEEK_END)
size = buffer.tell()
buffer.seek(pos)
return size
def _extract_keys(self, buffer: BytesIO, key: str) -> Tuple[str, BytesIO]:
b = buffer.getvalue()
start_pos = b.find((f'<{key}>').encode('utf-8'))
if start_pos == -1:
buffer.seek(0)
extract = None
new_stream = None
else:
end_pos = b.find((f'</{key}>').encode('utf-8'), start_pos)
buffer.seek(start_pos)
content = buffer.read(end_pos + len(f'</{key}>') - start_pos)
extract = content.decode('utf-8')
new_stream = BytesIO(buffer.read())
return extract, new_stream
@timeit
def _get_connection(self, report_url: str):
auth_headers = self.creds.get_auth_headers()
conn = req.get(report_url, stream=True, headers=auth_headers)
return conn
def _find_fieldnames(self, buffer: BytesIO) -> Tuple[str, BytesIO]:
header, buffer = self._extract_keys(buffer=buffer, key='thead')
if header:
fieldnames = [
CSVHelpers.sanitize_string(field)
for field in re.findall(r'\<th[^>]*\>([^<]*)\<\/th\>', header)
]
# logging.info(f'Fields: {fieldnames}')
del header
else:
fieldnames = None
return fieldnames, buffer
def _next_chunk(self,
stream,
html_chunk_size: int = None) -> Tuple[bytes, bool]:
_buffer = BytesIO()
last_chunk = False
while len(_buffer.getvalue()) < html_chunk_size and not last_chunk:
try:
_block = stream.__next__()
if _block: _buffer.write(_block)
except StopIteration:
last_chunk = True
return _buffer.getvalue(), last_chunk
@measure_memory
def _stream_processor(self,
bucket: str,
report_details: Dict[str, Any],
repeatable: bool = False) -> BytesIO:
repeater = BytesIO()
report_url = report_details['url']
remainder = b''
queue = Queue()
output_buffer = StringIO()
html_chunk_size = 2048 * 1024
chunk_size = 1024 * 1024
streamer = ThreadedGCSObjectStreamUpload(
client=Cloud_Storage.client(credentials=self.creds),
bucket_name=bucket,
blob_name='{id}.csv'.format(id=report_details['id']),
chunk_size=chunk_size,
queue=queue)
streamer.daemon = True
streamer.start()
try:
chunk_id = 0
conn = self._get_connection(report_url)
_stream = conn.iter_content(chunk_size=html_chunk_size)
source_size = 0
done = False
fieldnames = None
while not done:
# logging.info(f'Processing chunk {chunk_id}')
# logging.info(f'Processing chunk {chunk_id}, remainder {remainder.decode("utf-8")}')
chunk = BytesIO()
chunk.write(remainder)
remainder = b''
block, done = self._next_chunk(_stream, html_chunk_size)
source_size += len(block)
# logging.info(f'{len(block):,}, begins {block[0:80]} : ends {block[-80:].decode("utf-8")}')
if repeatable: repeater.write(block)
chunk.write(block)
if len(chunk.getvalue()) < html_chunk_size and not done:
continue
# logging.info(f'Chunk size {len(chunk.getvalue()):,} bytes')
chunk.seek(0)
if chunk_id == 0:
fieldnames, chunk = self._find_fieldnames(buffer=chunk)
# find last </tr> on any section but the last, chop off the last portion and store
last_tr_pos = chunk.getvalue().rfind(b'</tr>')
if last_tr_pos == -1:
# logging.debug(f'HALP! {chunk.getvalue()}')
remainder = chunk.getvalue()
continue
else:
last_tr_pos += 5
chunk.seek(last_tr_pos)
remainder = chunk.read()
# logging.debug(f'Remainder: {remainder}')
chunk.truncate(last_tr_pos)
rows = []
while True:
tr, chunk = self._extract_keys(chunk, 'tr')
if chunk:
rows.append([
unescape(field) for field in re.findall(
r'\<td[^>]*\>([^<]*)\<\/td\>', tr)
])
else:
break
# queue for upload
report_data = []
for row in rows:
report_data.append(dict(zip(fieldnames, row)))
writer = csv.DictWriter(output_buffer, fieldnames=fieldnames)
if chunk_id == 0: writer.writeheader()
[writer.writerow(row) for row in report_data]
output_buffer.seek(0)
# logging.info(f'Sending chunk {chunk_id} size {len(output_buffer.getvalue())}')
queue.put((chunk_id, output_buffer.getvalue().encode('utf-8')))
chunk_id += 1
chunk = BytesIO()
output_buffer.seek(0)
output_buffer.truncate(0)
logging.info(f'SA360 report length: {source_size:,} bytes')
queue.join()
streamer.stop()
report_details['schema'] = CSVHelpers.create_table_schema(
fieldnames)
return repeater
except Exception as e:
logging.error(e)
@measure_memory
def upload_report(self,
bucket: str,
report_details: Dict[str, Any],
input_buffer: BytesIO = None):
output_buffer = StringIO() #BytesIO()
try:
if not input_buffer:
input_buffer = BytesIO()
request = requests.Download(report_details['url'],
stream=input_buffer)
request.consume(transport=self.transport)
logging.info('Report data size: {bytes}'.format(bytes=0))
input_buffer.seek(0)
soup = self._soupify(input_buffer)
# del input_buffer
headers = soup.find('thead').find_all('th')
fieldnames = []
for header in headers:
fieldnames.append(CSVHelpers.sanitize_string(header.string))
rows = soup.find('tbody').find_all('tr')
report_data = []
for row in rows:
data = []
for col in row.contents:
data.append(col.string)
report_data.append(dict(zip(fieldnames, data)))
writer = csv.DictWriter(output_buffer, fieldnames=fieldnames)
writer.writeheader()
for row in report_data:
writer.writerow(row)
output_buffer.seek(0)
Cloud_Storage.write_file(bucket=bucket,
file=f"{report_details['id']}.csv",
data=output_buffer.getvalue())
report_details['schema'] = CSVHelpers.create_table_schema(
fieldnames)
except Exception as e:
logging.error(e)
class Report2BQ(object):
def __init__(self,
product: Type,
email=None,
project=None,
report_id=None,
profile=None,
sa360_url=None,
force: bool = False,
append: bool = False,
infer_schema: bool = False,
dest_project: str = None,
dest_dataset: str = 'report2bq',
notify_topic: str = None,
notify_message: str = None):
self.product = product
self.force = force
self.email = email
self.append = append
self.infer_schema = infer_schema
self.report_id = report_id
self.sa360_url = unquote(sa360_url) if sa360_url else None
self.cm_profile = profile
self.project = project
self.dest_project = dest_project
self.dest_dataset = dest_dataset
self.notify_topic = notify_topic
self.notify_message = notify_message
self.firestore = Firestore(email=email, project=project)
def handle_report_fetcher(self, fetcher: ReportFetcher):
# Get Latest Report
report_object = fetcher.get_latest_report_file(self.report_id)
# Normalize Report Details
report_data = fetcher.fetch_report_config(report_object=report_object,
report_id=self.report_id)
last_report = self.firestore.get_report_config(fetcher.report_type,
self.report_id)
if last_report:
if report_data['last_updated'] == last_report[
'last_updated'] and not self.force:
logging.info('No change: ignoring.')
return
report_data = fetcher.normalize_report_details(
report_object=report_object, report_id=self.report_id)
report_data['email'] = self.email
report_data['append'] = self.append
if self.dest_project: report_data['dest_project'] = self.dest_project
if self.dest_dataset: report_data['dest_dataset'] = self.dest_dataset
if self.notify_topic:
report_data['notifier'] = {
'topic': self.notify_topic,
}
if self.notify_message:
report_data['notifier']['message'] = self.notify_message
if report_object:
csv_header, csv_types = fetcher.read_header(report_data)
if csv_header:
schema = CSVHelpers.create_table_schema(
csv_header, csv_types if self.infer_schema else None)
report_data['schema'] = schema
fetcher.stream_to_gcs(f'{self.project}-report2bq-upload',
report_data)
self.firestore.store_report_config(fetcher.report_type, self.report_id,
report_data)
def handle_sa360(self):
sa360 = SA360(project=self.project,
email=self.email,
infer_schema=self.infer_schema,
append=self.append)
logging.info(self.sa360_url)
id = re.match(r'^.*rid=([0-9]+).*$', self.sa360_url).group(1)
report_data = self.firestore.get_report_config(Type.SA360, id)
if not report_data:
# Create new report details structure
report_data = {'id': id, 'url': self.sa360_url}
report_data['table_name'] = 'SA360_{id}'.format(id=id)
report_data['email'] = self.email
if self.dest_project: report_data['dest_project'] = self.dest_project
if self.dest_dataset: report_data['dest_dataset'] = self.dest_dataset
if self.notify_topic:
report_data['notifier'] = {
'topic': self.notify_topic,
}
if self.notify_message:
report_data['notifier']['message'] = self.notify_message
sa360.process(
bucket='{project}-report2bq-upload'.format(project=self.project),
report_details=report_data)
self.firestore.store_report_config(Type.SA360, id, report_data)
def run(self):
logging.info(f'Product: {self.product}')
if self.product in [Type.DV360, Type.CM]:
fetcher = FetcherFactory.create_fetcher(self.product,
email=self.email,
project=self.project,
profile=self.cm_profile)
self.handle_report_fetcher(fetcher=fetcher)
elif self.product == Type.SA360:
self.handle_sa360()