<commit_before>package org.lumongo.storage.rawfiles;

import java.io.IOException;
import java.io.InputStream;
import java.net.UnknownHostException;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.zip.DeflaterInputStream;
import java.util.zip.InflaterInputStream;

import org.apache.log4j.Logger;
import org.bson.BSON;
import org.lumongo.cluster.message.Lumongo.AssociatedDocument;
import org.lumongo.cluster.message.Lumongo.FetchRequest.FetchType;
import org.lumongo.cluster.message.Lumongo.Metadata;
import org.lumongo.cluster.message.Lumongo.ResultDocument;
import org.lumongo.storage.constants.MongoConstants;
import org.lumongo.util.CommonCompression;
import org.lumongo.util.CommonCompression.CompressionLevel;

import com.google.protobuf.ByteString;
import com.mongodb.BasicDBObject;
import com.mongodb.CommandResult;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBObject;
import com.mongodb.Mongo;
import com.mongodb.MongoException;
import com.mongodb.WriteConcern;
import com.mongodb.gridfs.GridFS;
import com.mongodb.gridfs.GridFSDBFile;
import com.mongodb.gridfs.GridFSFile;

public class MongoDocumentStorage implements DocumentStorage {
	@SuppressWarnings("unused")
	private final static Logger log = Logger.getLogger(MongoDocumentStorage.class);

	public static final String STORAGE_DB_SUFFIX = "_storage";

	private static final Charset UTF_8_CHARSET = Charset.forName("UTF-8");
	private static final String ASSOCIATED_FILES = "associatedFiles";
	private static final String FILES = "files";
	private static final String FILENAME = "filename";
	private static final String VALUE = "value";
	private static final String DOC = "doc";
	private static final String TIMESTAMP = "timestamp";
	private static final String TYPE = "type";
	private static final String COMPRESSED = "compressed";
	private static final String METADATA = "metadata";
	private static final String RESULT_STORAGE_COLLECTION = "resultStorage";
	private static final String UNIQUE_ID_KEY = "unique_id";
	private static final String CHUNKS = "chunks";

	private Mongo pool;
	private String database;

	// TODO: make this configurable?
	private WriteConcern documentWriteConcern = WriteConcern.NORMAL;

	public MongoDocumentStorage(Mongo pool, String database) {
		this(pool, database, false);
	}

	public MongoDocumentStorage(Mongo pool, String database, boolean sharded) {
		this.pool = pool;
		this.database = database;

		if (sharded) {
			DB storageDb = pool.getDB(database);

			DBCollection coll = storageDb.getCollection(ASSOCIATED_FILES + "." + FILES);
			coll.ensureIndex(METADATA + "." + UNIQUE_ID_KEY);

			DB adminDb = pool.getDB(MongoConstants.StandardDBs.ADMIN);
			DBObject enableCommand = new BasicDBObject();
			enableCommand.put(MongoConstants.Commands.ENABLE_SHARDING, database);
			CommandResult cr = adminDb.command(enableCommand);
			if (cr.getErrorMessage() != null) {
				System.err.println("Failed to enable sharding on the database <" + database + "> because <" + cr.getErrorMessage() + ">");
			}

			shardCollection(storageDb, adminDb, RESULT_STORAGE_COLLECTION);
			shardCollection(storageDb, adminDb, ASSOCIATED_FILES + "." + CHUNKS);
		}
	}

	public MongoDocumentStorage(String mongoHost, int mongoPort, String database, boolean sharded) throws UnknownHostException, MongoException {
		this(new Mongo(mongoHost, mongoPort), database, sharded);
	}

	private void shardCollection(DB db, DB adminDb, String collectionName) {
		CommandResult cr;
		DBObject shardCommand = new BasicDBObject();
		DBCollection collection = db.getCollection(collectionName);
		shardCommand.put(MongoConstants.Commands.SHARD_COLLECTION, collection.getFullName());
		shardCommand.put(MongoConstants.Commands.SHARD_KEY, new BasicDBObject(MongoConstants.StandardFields._ID, 1));
		cr = adminDb.command(shardCommand);
		if (cr.getErrorMessage() != null) {
			System.err.println("Failed to shard the collection <" + collectionName + "> because <" + cr.getErrorMessage() + ">");
		}
	}

	private GridFS createGridFSConnection() {
		DB db = pool.getDB(database);
		return new GridFS(db, ASSOCIATED_FILES);
	}

	@Override
	public void storeSourceDocument(String uniqueId, ResultDocument doc) throws Exception {
		DB db = pool.getDB(database);
		DBCollection coll = db.getCollection(RESULT_STORAGE_COLLECTION);
		DBObject document = new BasicDBObject();
		if (!doc.getCompressed() && doc.getType().equals(ResultDocument.Type.BSON)) {
			document.putAll(BSON.decode(doc.getDocument().toByteArray()));
		}
		else if (!doc.getCompressed() && doc.getType().equals(ResultDocument.Type.TEXT)) {
			document.put(DOC, doc.getDocument().toStringUtf8());
		}
		else {
			byte[] bytes = doc.getDocument().toByteArray();
			if (doc.getCompressed()) {
				bytes = CommonCompression.compressZlib(bytes, CompressionLevel.NORMAL);
			}
			document.put(DOC, bytes);
		}

		if (doc.getMetadataCount() > 0) {
			DBObject metadata = new BasicDBObject();
			for (Metadata meta : doc.getMetadataList()) {
				metadata.put(meta.getKey(), meta.getValue());
			}
			document.put(METADATA, metadata);
		}

		document.put(TIMESTAMP, doc.getTimestamp());
		document.put(COMPRESSED, doc.getCompressed());
		document.put(TYPE, doc.getType().toString());
		document.put(MongoConstants.StandardFields._ID, uniqueId);
		coll.save(document, documentWriteConcern);
	}

	@Override
	public ResultDocument getSourceDocument(String uniqueId, FetchType fetchType) throws Exception {
		if (!FetchType.NONE.equals(fetchType)) {
			DB db = pool.getDB(database);
			DBCollection coll = db.getCollection(RESULT_STORAGE_COLLECTION);
			DBObject search = new BasicDBObject(MongoConstants.StandardFields._ID, uniqueId);
			DBObject result = coll.findOne(search);
			if (null != result) {
				boolean compressed = (boolean) result.removeField(COMPRESSED);

				long timestamp = (long) result.removeField(TIMESTAMP);

				ResultDocument.Type type = ResultDocument.Type.valueOf((String) result.removeField(TYPE));
				ResultDocument.Builder dBuilder = ResultDocument.newBuilder();
				dBuilder.setType(type);
				dBuilder.setCompressed(compressed);
				dBuilder.setUniqueId(uniqueId);
				dBuilder.setTimestamp(timestamp);

				if (result.containsField(METADATA)) {
					DBObject metadata = (DBObject) result.removeField(METADATA);
					for (String key : metadata.keySet()) {
						dBuilder.addMetadata(Metadata.newBuilder().setKey(key).setValue((String) metadata.get(VALUE)).build());
					}
				}

				if (FetchType.FULL.equals(fetchType)) {
					ByteString document = null;
					if (!compressed && type.equals(ResultDocument.Type.BSON)) {
						document = ByteString.copyFrom(BSON.encode(result));
					}
					else if (!compressed && type.equals(ResultDocument.Type.TEXT)) {
						document = ByteString.copyFrom(((String) result.get(DOC)).getBytes(UTF_8_CHARSET));
					}
					else {
						byte[] bytes = (byte[]) result.get(DOC);
						if (compressed) {
							bytes = CommonCompression.uncompressZlib(bytes);
						}
						document = ByteString.copyFrom(bytes);
					}
					dBuilder.setDocument(document);
				}


				return dBuilder.build();
			}
		}
		return null;
	}<commit_msg>add _id into the BSON object if stored compressed so that it matches object returned without compression on<commit_after>package org.lumongo.storage.rawfiles;

import java.io.IOException;
import java.io.InputStream;
import java.net.UnknownHostException;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.zip.DeflaterInputStream;
import java.util.zip.InflaterInputStream;

import org.apache.log4j.Logger;
import org.bson.BSON;
import org.bson.BSONObject;
import org.lumongo.cluster.message.Lumongo.AssociatedDocument;
import org.lumongo.cluster.message.Lumongo.FetchRequest.FetchType;
import org.lumongo.cluster.message.Lumongo.Metadata;
import org.lumongo.cluster.message.Lumongo.ResultDocument;
import org.lumongo.storage.constants.MongoConstants;
import org.lumongo.util.CommonCompression;
import org.lumongo.util.CommonCompression.CompressionLevel;

import com.google.protobuf.ByteString;
import com.mongodb.BasicDBObject;
import com.mongodb.CommandResult;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBObject;
import com.mongodb.Mongo;
import com.mongodb.MongoException;
import com.mongodb.WriteConcern;
import com.mongodb.gridfs.GridFS;
import com.mongodb.gridfs.GridFSDBFile;
import com.mongodb.gridfs.GridFSFile;

public class MongoDocumentStorage implements DocumentStorage {
	@SuppressWarnings("unused")
	private final static Logger log = Logger.getLogger(MongoDocumentStorage.class);

	public static final String STORAGE_DB_SUFFIX = "_storage";

	private static final Charset UTF_8_CHARSET = Charset.forName("UTF-8");
	private static final String ASSOCIATED_FILES = "associatedFiles";
	private static final String FILES = "files";
	private static final String FILENAME = "filename";
	private static final String VALUE = "value";
	private static final String DOC = "doc";
	private static final String TIMESTAMP = "timestamp";
	private static final String TYPE = "type";
	private static final String COMPRESSED = "compressed";
	private static final String METADATA = "metadata";
	private static final String RESULT_STORAGE_COLLECTION = "resultStorage";
	private static final String UNIQUE_ID_KEY = "unique_id";
	private static final String CHUNKS = "chunks";

	private Mongo pool;
	private String database;

	// TODO: make this configurable?
	private WriteConcern documentWriteConcern = WriteConcern.NORMAL;

	public MongoDocumentStorage(Mongo pool, String database) {
		this(pool, database, false);
	}

	public MongoDocumentStorage(Mongo pool, String database, boolean sharded) {
		this.pool = pool;
		this.database = database;

		if (sharded) {
			DB storageDb = pool.getDB(database);

			DBCollection coll = storageDb.getCollection(ASSOCIATED_FILES + "." + FILES);
			coll.ensureIndex(METADATA + "." + UNIQUE_ID_KEY);

			DB adminDb = pool.getDB(MongoConstants.StandardDBs.ADMIN);
			DBObject enableCommand = new BasicDBObject();
			enableCommand.put(MongoConstants.Commands.ENABLE_SHARDING, database);
			CommandResult cr = adminDb.command(enableCommand);
			if (cr.getErrorMessage() != null) {
				System.err.println("Failed to enable sharding on the database <" + database + "> because <" + cr.getErrorMessage() + ">");
			}

			shardCollection(storageDb, adminDb, RESULT_STORAGE_COLLECTION);
			shardCollection(storageDb, adminDb, ASSOCIATED_FILES + "." + CHUNKS);
		}
	}

	public MongoDocumentStorage(String mongoHost, int mongoPort, String database, boolean sharded) throws UnknownHostException, MongoException {
		this(new Mongo(mongoHost, mongoPort), database, sharded);
	}

	private void shardCollection(DB db, DB adminDb, String collectionName) {
		CommandResult cr;
		DBObject shardCommand = new BasicDBObject();
		DBCollection collection = db.getCollection(collectionName);
		shardCommand.put(MongoConstants.Commands.SHARD_COLLECTION, collection.getFullName());
		shardCommand.put(MongoConstants.Commands.SHARD_KEY, new BasicDBObject(MongoConstants.StandardFields._ID, 1));
		cr = adminDb.command(shardCommand);
		if (cr.getErrorMessage() != null) {
			System.err.println("Failed to shard the collection <" + collectionName + "> because <" + cr.getErrorMessage() + ">");
		}
	}

	private GridFS createGridFSConnection() {
		DB db = pool.getDB(database);
		return new GridFS(db, ASSOCIATED_FILES);
	}

	@Override
	public void storeSourceDocument(String uniqueId, ResultDocument doc) throws Exception {
		DB db = pool.getDB(database);
		DBCollection coll = db.getCollection(RESULT_STORAGE_COLLECTION);
		DBObject document = new BasicDBObject();
		if (!doc.getCompressed() && doc.getType().equals(ResultDocument.Type.BSON)) {
			document.putAll(BSON.decode(doc.getDocument().toByteArray()));
		}
		else if (!doc.getCompressed() && doc.getType().equals(ResultDocument.Type.TEXT)) {
			document.put(DOC, doc.getDocument().toStringUtf8());
		}
		else {
			byte[] bytes = doc.getDocument().toByteArray();
			if (doc.getCompressed()) {
				bytes = CommonCompression.compressZlib(bytes, CompressionLevel.NORMAL);
			}
			document.put(DOC, bytes);
		}

		if (doc.getMetadataCount() > 0) {
			DBObject metadata = new BasicDBObject();
			for (Metadata meta : doc.getMetadataList()) {
				metadata.put(meta.getKey(), meta.getValue());
			}
			document.put(METADATA, metadata);
		}

		document.put(TIMESTAMP, doc.getTimestamp());
		document.put(COMPRESSED, doc.getCompressed());
		document.put(TYPE, doc.getType().toString());
		document.put(MongoConstants.StandardFields._ID, uniqueId);
		coll.save(document, documentWriteConcern);
	}

	@Override
	public ResultDocument getSourceDocument(String uniqueId, FetchType fetchType) throws Exception {
		if (!FetchType.NONE.equals(fetchType)) {
			DB db = pool.getDB(database);
			DBCollection coll = db.getCollection(RESULT_STORAGE_COLLECTION);
			DBObject search = new BasicDBObject(MongoConstants.StandardFields._ID, uniqueId);
			DBObject result = coll.findOne(search);
			if (null != result) {
				boolean compressed = (boolean) result.removeField(COMPRESSED);

				long timestamp = (long) result.removeField(TIMESTAMP);

				ResultDocument.Type type = ResultDocument.Type.valueOf((String) result.removeField(TYPE));
				ResultDocument.Builder dBuilder = ResultDocument.newBuilder();
				dBuilder.setType(type);
				dBuilder.setCompressed(compressed);
				dBuilder.setUniqueId(uniqueId);
				dBuilder.setTimestamp(timestamp);

				if (result.containsField(METADATA)) {
					DBObject metadata = (DBObject) result.removeField(METADATA);
					for (String key : metadata.keySet()) {
						dBuilder.addMetadata(Metadata.newBuilder().setKey(key).setValue((String) metadata.get(VALUE)).build());
					}
				}

				if (FetchType.FULL.equals(fetchType)) {
					ByteString document = null;
					if (!compressed && type.equals(ResultDocument.Type.BSON)) {
						document = ByteString.copyFrom(BSON.encode(result));
					}
					else if (!compressed && type.equals(ResultDocument.Type.TEXT)) {
						document = ByteString.copyFrom(((String) result.get(DOC)).getBytes(UTF_8_CHARSET));
					}
					else {
						byte[] bytes = (byte[]) result.get(DOC);
						if (compressed) {
							bytes = CommonCompression.uncompressZlib(bytes);
						}
						if (type.equals(ResultDocument.Type.BSON)) {
							BSONObject object = BSON.decode(bytes);
							object.put(MongoConstants.StandardFields._ID, uniqueId);
							document = ByteString.copyFrom(BSON.encode(object));
						}
						else {
							document = ByteString.copyFrom(bytes);
						}
					}
					dBuilder.setDocument(document);
				}


				return dBuilder.build();
			}
		}
		return null;
	}
{'fraction_non_alphanumeric': 0.11555075593952484, 'fraction_numerical': 0.0009449244060475162, 'mean_word_length': 4.981832862333468, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1894818', 'n_tokens_mistral': 5095, 'n_tokens_neox': 4485, 'n_words': 937}