def run(spark):
"""Indexes Protobuf files by updating the S3Prefixes and VehPosPb tables
Args:
spark: Spark Session object
"""
log = utils.get_logger()
with DBConnCommonQueries() as conn:
dbtables.create_if_not_exists(conn, dbtables.S3Prefixes)
dbtables.create_if_not_exists(conn, dbtables.VehPosPb)
pfxs = explore_s3prefixes()
s3Mgr = s3.S3Mgr()
for pfx in pfxs:
fullPfx = '/'.join(("pb", "VehiclePos", pfx))
keys = s3Mgr.fetch_keys(fullPfx)
if len(keys) > 0:
log.info("PROCESSING %d KEYS FOR %s", len(keys), pfx)
file_list = spark.sparkContext.parallelize(keys)
file_list \
.map(dbtables.VehPosPb.build_tuple_from_protobuf) \
.foreachPartition(push_vehpospb_dbtpls)
log.info("PROCESSED %d KEYS FOR %s", len(keys), pfx)
tpl = (pfx, len(keys))
with DBConn() as conn:
dbtables.S3Prefixes.insert_values(conn, pfx, len(keys))
conn.commit()
log.info("PUSHED S3Prefix %s", str(tpl))
def download_feed(dirName, url, *args):
"""Downloads a real-time vehicle positions file to the local storage
and then uploads it to S3 and removes from the local storage
Args:
dirName: the local directory where the file will be saved initially
url: the URL to download the file from
*args: placeholder for any additional arguments, unused
"""
fName = datetime.now().strftime("%Y%m%d-%H%M%S.pb")
r = requests.get(url)
fPath = os.path.join(Settings.ProjPath, "pb", dirName, fName)
with open(fPath, "wb") as handle:
handle.write(r.content)
try:
# always use '/' as path separator in S3
objKey = '/'.join(["pb", dirName, fName.replace('-', '/')])
s3Mgr = s3.S3Mgr()
s3Mgr.upload_file(fPath, objKey)
os.remove(fPath)
except Exception: # pylint: disable=broad-except
log = utils.get_logger()
log.warning("Error while saving the file %s to S3 and/or DB", fPath)
log.warning(traceback.format_exc())
pass # do not interfere with other threads that might succeed
def read_feed_descs():
"""Retrieves a list of GTFS feed descriptions available on S3
"""
objKey = '/'.join(["GTFS", "MBTA_archived_feeds.txt"])
s3Mgr = s3.S3Mgr()
content = s3Mgr.fetch_object_body(objKey)
return gtfs.read_feed_descs(content)
def _process_pqdate(pqDate):
mxdstr = '0' if Settings.MaxAbsDelay <= 0 else str(Settings.MaxAbsDelay)
s3_prefix = f"HlyDelays{mxdstr}/{pqDate.strftime('%Y%m%d')}.pq"
rexpr = re.compile(r'.*part-.*\.parquet$')
s3Mgr = s3.S3Mgr()
for key in s3Mgr.fetch_keys(s3_prefix):
if rexpr.match(key):
key = f's3://{Settings.S3BucketName}/{key}'
df = pd.read_parquet(key)
_process_df(df, pqDate)
logger.info(f'Processed entries from {key}')
def update_s3(self, dfHlyDelays, pqDate):
"""Saves a delays dataframe to the S3 in parquet
"""
s3Mgr = s3.S3Mgr()
mxdstr = '0' if Settings.MaxAbsDelay <= 0 else str(
Settings.MaxAbsDelay)
pfx = f"HlyDelays{mxdstr}/{pqDate.strftime('%Y%m%d.pq')}"
if s3Mgr.prefix_exists(pfx):
s3Mgr.delete_prefix(pfx)
time.sleep(5)
dfHlyDelays = dfHlyDelays \
.withColumn(
'route_stop',
F.concat(
dfHlyDelays.RouteId, F.lit(':::'),
F.lit('['), dfHlyDelays.StopName, F.lit(']')
)
)
dfHlyDelays = dfHlyDelays \
.groupBy(dfHlyDelays.route_stop) \
.agg(
F.collect_list(
F.struct(
dfHlyDelays.DateEST, dfHlyDelays.HourEST,
dfHlyDelays.AvgDelay, dfHlyDelays.AvgDist, dfHlyDelays.Cnt
)
).alias('vals_unsorted')
)
udf_ret_type = ArrayType(
StructType([
StructField("DateEST", DateType(), False),
StructField("HourEST", IntegerType(), False),
StructField("AvgDelay", DoubleType(), False),
StructField("AvgDist", DoubleType(), False),
StructField("Cnt", IntegerType(), False)
]))
udf_sort_vals = F.udf(
lambda vals: list(
sorted(vals, key=lambda r: (r.DateEST, r.HourEST))),
udf_ret_type)
dfHlyDelays = dfHlyDelays \
.withColumn('vals', udf_sort_vals(dfHlyDelays.vals_unsorted)) \
.drop('vals_unsorted')
dfHlyDelays.printSchema()
s3_path = "s3a://%s/%s" % (Settings.S3BucketName, pfx)
dfHlyDelays.write.mode("overwrite").parquet(s3_path)
def upload_zips(feedDescs):
"""Checks whether all feed descriptions are available in S3, then
retrieves and downloads the missing ones
Args:
feedDescs: a list of MBTA_ArchivedFeedDesc objects
"""
s3Mgr = s3.S3Mgr()
for feedDesc in feedDescs:
s3Key = '/'.join(["GTFS", feedDesc.s3Key])
archS3Key = '/'.join(["GTFS_Archived", feedDesc.s3Key])
if not s3Mgr.prefix_exists(s3Key) and not s3Mgr.prefix_exists(archS3Key):
upload_zip(s3Mgr, feedDesc, s3Key)
def update_archive_txt(): """Retrieves a text file describing the currently available archives from MBTA and stores it in S3 if it can be successfully parsed by the app """ url = "https://cdn.mbta.com/archive/archived_feeds.txt" r = requests.get(url) # parse first to check if we can work with this feedDescs = gtfs.read_feed_descs(r.content) objKey = '/'.join(["GTFS", "MBTA_archived_feeds.txt"]) s3Mgr = s3.S3Mgr() s3Mgr.put_object_body(objKey, r.content) return feedDescs
def main():
"""Compresses text files from obsolete feeds in S3, uploads archives to S3,
and removes the text files
Uses tar.bz2 format
"""
s3Mgr = s3.S3Mgr()
objKey = '/'.join(["GTFS", "MBTA_archived_feeds.txt"])
content = s3Mgr.fetch_object_body(objKey)
feedDescs = gtfs.read_feed_descs(content)
with DBConn() as conn:
dtNow = dbtables.PqDates.select_latest_processed(conn)
if not dtNow:
return
logger.info('Latest processed parquet date is %s' % str(dtNow))
for fd in feedDescs:
daysDiff = (dtNow - fd.endDate).total_seconds() / (24 * 3600)
if daysDiff > Settings.GTFS_ObsoleteAfterDays:
archive_gtfs_files(s3Mgr, fd)
def _rollback_2020():
s3Mgr = s3.S3Mgr()
seasons = ['Winter', 'Spring', 'Summer', 'Fall']
pfxs20 = [x + ' 2020' for x in seasons]
pfx20_lens = [len(x) for x in pfxs20]
for objKey in s3Mgr.fetch_keys('GTFS_Archived'):
m = re.match(r'GTFS_Archived/(?P<feedname>.*)\.tar.bz2', objKey)
if not m:
continue
feedname = m['feedname']
for pfx, pfx_len in zip(pfxs20, pfx20_lens):
if feedname[0:pfx_len] == pfx:
logger.info(f'Proceeding with {objKey}')
data = s3Mgr.fetch_object_body(objKey)
buffer = BytesIO(data)
tbz2 = tarfile.open(mode="r:bz2", fileobj=buffer)
for member in tbz2.getmembers():
extr_key = '/'.join(['GTFS', feedname, member.name])
if s3Mgr.prefix_exists(extr_key):
logger.warning(f'Prefix {extr_key} already exists')
continue
s3Mgr.put_object_body(extr_key,
tbz2.extractfile(member).read())
s3Mgr.delete_key(objKey)