forked from bellwether/minerva
117 lines
3.4 KiB
Python
117 lines
3.4 KiB
Python
import boto3
|
|
import os
|
|
import random
|
|
import time
|
|
import pyarrow as pa
|
|
import pyarrow.dataset
|
|
import pprint
|
|
import datetime
|
|
import json
|
|
from minerva import parallel_map
|
|
|
|
pp = pprint.PrettyPrinter(indent=4)
|
|
|
|
class Redshift:
|
|
def __init__(self, profile, output, db=None, cluster=None):
|
|
self.session = boto3.session.Session(profile_name=profile)
|
|
self.redshift = self.session.client("redshift-data")
|
|
self.output = output
|
|
self.database = db
|
|
self.cluster = cluster
|
|
|
|
def query(self, sql):
|
|
q = Query(self, sql)
|
|
q.run()
|
|
return q
|
|
|
|
def execute(self, sql):
|
|
e = Execute(self, sql)
|
|
e.run()
|
|
return e
|
|
|
|
def download(self, s3):
|
|
bucket = s3.split("/")[2]
|
|
file = os.path.join(*s3.split("/")[3:])
|
|
tmp = f"/tmp/{random.random()}.bin"
|
|
self.session.client('s3').download_file(bucket, file, tmp)
|
|
|
|
return tmp
|
|
|
|
class Execute:
|
|
"""
|
|
Execute will not return results, but will execute the SQL and return the final state.
|
|
Execute is meant to be used for DML statements such as CREATE DATABASE/TABLE
|
|
"""
|
|
def __init__(self, handler, sql):
|
|
self.handler = handler
|
|
self.redshift = handler.redshift
|
|
self.sql = sql
|
|
self.info_cache = None
|
|
|
|
def query(self):
|
|
return self.sql
|
|
|
|
def run(self):
|
|
resp = self.redshift.execute_statement(Sql=self.query(),
|
|
Database=self.handler.database,
|
|
ClusterIdentifier=self.handler.cluster)
|
|
self.query_id = resp['Id']
|
|
return resp
|
|
|
|
def status(self):
|
|
return self.info()['Status']
|
|
|
|
def info(self):
|
|
res = self.redshift.describe_statement(Id=self.query_id)
|
|
self.info_cache = res
|
|
return self.info_cache
|
|
|
|
def finish(self):
|
|
stat = self.status()
|
|
while stat in ['SUBMITTED', 'PICKED', 'STARTED']:
|
|
time.sleep(5)
|
|
stat = self.status()
|
|
|
|
return stat # finalized state
|
|
|
|
|
|
class Query(Execute):
|
|
DATA_STYLE = 'parquet'
|
|
|
|
def query(self):
|
|
self.out = os.path.join(self.handler.output,
|
|
str(random.random()))
|
|
query = f"unload ({repr(self.sql)}) to {repr(self.out)} " + \
|
|
f"iam_role default " + \
|
|
f"format as {self.DATA_STYLE} " + \
|
|
f"manifest"
|
|
return query
|
|
|
|
def manifest_files(self):
|
|
status = self.finish()
|
|
|
|
if status == "FINISHED":
|
|
# Track the runtime
|
|
self.runtime = self.info_cache['UpdatedAt'] - self.info_cache['CreatedAt']
|
|
|
|
# Because we're using `UNLOAD`, we get a manifest of the files
|
|
# that make up our data.
|
|
manif = self.out + "manifest"
|
|
tmp = self.handler.download(manif)
|
|
with open(tmp, "r") as f:
|
|
js = json.load(f)
|
|
|
|
# Filter empty strings
|
|
files = [e['url'].strip() for e in js['entries'] if e['url'].strip()]
|
|
|
|
return files
|
|
else:
|
|
return status # canceled or error
|
|
|
|
# TODO parallelize this
|
|
def results(self):
|
|
#local = [self.handler.download(f) for f in self.manifest_files()]
|
|
local = parallel_map(self.handler.download, self.manifest_files())
|
|
self.ds = pa.dataset.dataset(local)
|
|
return self.ds
|
|
|