How to delete records from MongoDB using Glue Job

Often AWS engineers find themselves in a situation where they need to achieve something, but AWS does not have the required feature. We recently found ourselves in one. We were already using a glue job to insert a large number of records in MongoDB. Now we wanted to delete these records via the same job if a delete flag is enabled.

We knew that Spark has a feature called repartition to scale.

partitioned_df = df.repartition(partition_count, col("salt_column”))

It writes parallelly to the DB via write_dynamic_frame.

glueContext.write_dynamic_frame.from_options(partitioned_df,
connection_type='documentdb',
connection_options=get_documentdb_options(collection))

But there are no delete operations like write_dynamic_frame. So we wrote to the support team at AWS support and they replied that such a feature was not implemented yet.

As engineers, we set out to find alternate solutions, and we discovered that we can achieve this using mapPartitions options.

ids = partitioned_df.rdd.mapPartitions(delete_from_db).collect()

mapPartitions return is a new RDD by applying a function to each partition of this RDD. In each partition, we can delete the records manually.

def delete_from_db(partitioned_df):
ids = []
collection = ''
for row in partitioned_df:
ids.append(row['id'])
collection = row['collection']
if(len(ids) > 0) :
return delete_records(collection, ids)
else :
return ids

def delete_records(collection, ids):
mongo_uri ="mongodb://{}:{}@{}/?ssl={}&ssl_cert_reqs=CERT_NONE".format(username, password), uri.split('//')[1],'true')
mongo_client = pymongo.MongoClient(mongo_uri)
db = mongo_client[database_name]
collection = db[collection]
chunked = chunks(ids, 1024)
for chunked_ids in chunked:
query = {"_id" : { '$in':chunked_ids}}
collection.delete_many(query)
return ids