upload.d

/+ dub.json:{
"name": "kafka2s3",
"dependencies": {
	"asdf": "~>0.1.0-alpha3",
	"vibe-s3": {"path": "vibe-s3"},
	"vibe-d": "~>0.7.29"
},
"versions": ["VibeCustomMain"]
}+/

import vibe.d;
import vibe.aws.s3;
import vibe.aws.credentials;
import vibe.aws.aws;

import std.format;
import std.array;
import std.exception;
import std.algorithm;
import std.process : environment;

import asdf;

// For ASDF optimization
struct InputStreamChain
{
	InputStream stream;
	ubyte[] buffer;
	size_t length;

	this(InputStream stream, size_t bufferLength = 4096 * 2)
	{
		this.stream = stream;
		buffer = new ubyte[bufferLength];
		popFront();
	}

	bool empty()
	{
		return length == 0;
	}

	ubyte[] front()
	{
		return buffer[0 .. length];
	}

	void popFront()
	{
		for(length = 0; length < buffer.length; )
		{
			auto len = stream.leastSize;
			if(len == 0)
				break;
			auto newLength = length + len;
			newLength = min(newLength, buffer.length);
			stream.read(buffer[length .. newLength]);
			length = newLength;
		}
	}
}

int main(string[] args)
{
	if (args.length != 2 || !args[1].startsWith("s3://"))
	{
		logError("Usage:\n\tvibe-lambda s3://bucketB/absolutePathTo");
		return 1;
	}

	setLogLevel(LogLevel.debug_);
	auto creds = new EnvAWSCredentials;

	auto region = environment.get("S3_REGION").enforce("S3_REGION environment variable should be set.");
    auto cfg = ClientConfiguration();
    cfg.maxErrorRetry = 1;

	auto arg = args[1][5 .. $];
	auto path = arg.find("/");
	auto bucket = arg[0 .. $ - path.length];
	auto s3 = new S3(bucket, region, creds, cfg);

	// Zipped input ?
	bool gunzip = true;

	runTask({
			// Yet another vibe.d bug: StdinStream is 100-1000 times slower!
			// import vibe.stream.stdio;
			// InputStream stream = new StdinStream();
			InputStream stream = openFile("../data.gz");
			import vibe.stream.taskpipe;
			
			if(stream.empty)
			{
				logInfo("Input stream is empty.");
				return;
			}
			logInfo("Download ... ");
			
			TaskPipe pipe;
			OutputStream output;
			void delegate(const(char)[]) del;
			Task uploadTask;
			long counter; // TODO: init from Kafka
			long counterStart;
			Date date;

			void initializeOutput()
			{
				counterStart = counter;
				pipe = new TaskPipe;
				output = new GzipOutputStream(pipe);
				del = cast(typeof(del)) &output.write;
				uploadTask = runTask({
					logInfo("Multipart Upload ... ");
					s3.multipartUpload(format("%s-%s-incremental.jsonl.gz", path, date.toISOExtString), pipe);
					logInfo("Multipart Upload Finished.");
				});
			}
			
			void finalizeOutput()
			{
				logInfo("output.finalize");
				output.finalize;
				logInfo("finalize.finalize");
				pipe.finalize;
				logInfo("Wait uploading");
				uploadTask.join;
				logInfo("finalized");
			}
			
			foreach(asdf; InputStreamChain(gunzip ? new GzipInputStream(stream) : stream).parseJsonByLine(4096))
			{
				counter++;
				if (asdf.data.length == 0)
				{
					logInfo("msg is brocken");
					continue;
				}
				auto tsAsdf = asdf["ts"];
				if (tsAsdf.data.length == 0)
				{
					logInfo("ts not found");
					continue;
				}
				long ts;
				try ts = cast(long) tsAsdf;
				catch(Exception e)
				{
					logInfo("Ts format is brocken: %s", e.msg);
					continue;
				}
				// default timezone is LocalTime. UTC is used instead.
				auto currDate = cast(Date) SysTime.fromUnixTime(ts, UTC());
				if (currDate > date)
				{
					if (date != date.init)
					{
						finalizeOutput();
						initializeOutput();
						date = currDate;
					}
					else
					{
						date = currDate;
						initializeOutput();
					}
				}
				//tsAsdf.toString(del); // fast to execute
				asdf.toString(del);
				del("\n");
				if(counter % 1000 == 0)
					logInfo("counter = %s, ts = %s, asdf-len = %s", counter, ts, asdf.data.length);
			}

			finalizeOutput();
			logInfo("Exit event loop.");
			exitEventLoop();
		});

	lowerPrivileges();

	logDiagnostic("Running event loop...");
	int status;
	version (VibeDebugCatchAll) {
		try {
			status = runEventLoop();
		} catch( Throwable th ){
			logError("Unhandled exception in event loop: %s", th.msg);
			logDiagnostic("Full exception: %s", th.toString().sanitize());
			return 1;
		}
	} else {
		status = runEventLoop();
	}

	logDiagnostic("Event loop exited with status %d.", status);
	return status;
}