/* * ToroDB * Copyright © 2014 8Kdata Technology (www.8kdata.com) * * This program is free software: you can redistribute it and/or modify * it under the terms of the GNU Affero General Public License as published by * the Free Software Foundation, either version 3 of the License, or * (at your option) any later version. * * This program is distributed in the hope that it will be useful, * but WITHOUT ANY WARRANTY; without even the implied warranty of * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the * GNU Affero General Public License for more details. * * You should have received a copy of the GNU Affero General Public License * along with this program. If not, see <http://www.gnu.org/licenses/>. */ package com.torodb.torod.pipeline.impl; import com.google.common.base.Preconditions; import com.google.common.collect.Iterators; import com.google.inject.assistedinject.Assisted; import com.torodb.core.backend.WriteBackendTransaction; import com.torodb.core.d2r.CollectionData; import com.torodb.core.d2r.D2RTranslatorFactory; import com.torodb.core.dsl.backend.BackendTransactionJob; import com.torodb.core.dsl.backend.BackendTransactionJobFactory; import com.torodb.core.exceptions.user.UserException; import com.torodb.core.transaction.RollbackException; import com.torodb.core.transaction.metainf.MetaDatabase; import com.torodb.core.transaction.metainf.MutableMetaCollection; import com.torodb.kvdocument.values.KvDocument; import com.torodb.torod.pipeline.D2RTranslationBatchFunction; import com.torodb.torod.pipeline.DefaultToBackendFunction; import com.torodb.torod.pipeline.InsertPipeline; import org.jooq.lambda.Unchecked; import org.jooq.lambda.UncheckedException; import java.util.stream.Stream; import javax.inject.Inject; /** * */ public class SameThreadInsertPipeline implements InsertPipeline { private final D2RTranslatorFactory translatorFactory; private final MetaDatabase metaDb; private final MutableMetaCollection mutableMetaCollection; private final WriteBackendTransaction backendConnection; private final BackendTransactionJobFactory jobFactory; private int docBatchSize = 100; @Inject public SameThreadInsertPipeline(@Assisted D2RTranslatorFactory translatorFactory, @Assisted MetaDatabase metaDb, @Assisted MutableMetaCollection mutableMetaCollection, @Assisted WriteBackendTransaction backendConnection, BackendTransactionJobFactory factory) { this.translatorFactory = translatorFactory; this.metaDb = metaDb; this.mutableMetaCollection = mutableMetaCollection; this.backendConnection = backendConnection; this.jobFactory = factory; } @Override public void insert(Stream<KvDocument> docs) throws RollbackException, UserException { D2RTranslationBatchFunction d2rFun = new D2RTranslationBatchFunction(translatorFactory, metaDb, mutableMetaCollection); DefaultToBackendFunction r2BackendFun = new DefaultToBackendFunction(jobFactory, metaDb, mutableMetaCollection); try { Iterators.partition(docs.iterator(), docBatchSize) .forEachRemaining(list -> { CollectionData collData = d2rFun.apply(list); Iterable<BackendTransactionJob> jobs = r2BackendFun.apply(collData); jobs.forEach(Unchecked.consumer(job -> job.execute(backendConnection))); }); } catch (UncheckedException ex) { Throwable cause = ex.getCause(); if (cause != null && cause instanceof UserException) { throw (UserException) cause; } throw ex; } } @Override public int getDocsBatchLength() { return docBatchSize; } @Override public void setDocsBatchLength(int newBatchLength) { Preconditions.checkArgument(newBatchLength > 0, "The new batch size must be higher than 0, but %s was recived", newBatchLength); this.docBatchSize = newBatchLength; } public static interface Factory { SameThreadInsertPipeline create( D2RTranslatorFactory translatorFactory, MetaDatabase metaDb, MutableMetaCollection mutableMetaCollection, WriteBackendTransaction backendConnection); } }