def execute(self, context: 'Context'):
gdrive_hook = GoogleDriveHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
gcs_hook = GCSHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
file_metadata = gdrive_hook.get_file_id(
folder_id=self.folder_id, file_name=self.file_name, drive_id=self.drive_id
)
with gcs_hook.provide_file_and_upload(
bucket_name=self.bucket_name, object_name=self.object_name
) as file:
gdrive_hook.download_file(file_id=file_metadata["id"], file_handle=file)
def serialize_value(value: Any):
if isinstance(value, Dict):
hook = GCSHook()
if value.get("final", False) and "execution_date" in value.keys():
object_name = f"model/{value['execution_date']}/data.pickle"
elif "execution_date" in value.keys():
object_name = (
f"model/{value['execution_date']}/temp/data_{value['model_name']}_"
f"{datetime.now().strftime('%Y-%m-%dT%H:%M:%S')}.pickle")
else:
object_name = f"data/data_{uuid4()}.pickle"
value["location"] = object_name
with hook.provide_file_and_upload(
bucket_name=GCSXComBackend.BUCKET_NAME,
object_name=object_name) as f:
pickle.dump(value, f)
value = f"{GCSXComBackend.PREFIX}{object_name}"
return BaseXCom.serialize_value(value)
