ai.platon.pulsar.persist.gora.GoraStorage Maven / Gradle / Ivy

Go to download
package ai.platon.pulsar.persist.gora;

import ai.platon.pulsar.common.config.ImmutableConfig;
import ai.platon.pulsar.common.config.Params;
import ai.platon.pulsar.persist.HadoopUtils;
import ai.platon.pulsar.persist.gora.generated.GWebPage;
import org.apache.gora.persistency.Persistent;
import org.apache.gora.store.DataStore;
import org.apache.gora.store.DataStoreFactory;
import org.apache.gora.util.GoraException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

import static ai.platon.pulsar.common.config.AppConstants.MONGO_STORE_CLASS;
import static ai.platon.pulsar.common.config.AppConstants.WEBPAGE_SCHEMA;
import static ai.platon.pulsar.common.config.CapabilityTypes.*;

public class GoraStorage {
    public static final Logger logger = LoggerFactory.getLogger(GoraStorage.class);

    // load properties from gora.properties
    public static Properties properties = DataStoreFactory.createProps();
    private static Map dataStores = new HashMap<>();

    @SuppressWarnings("unchecked")
    public synchronized static  DataStore
    createDataStore(ImmutableConfig conf, Class keyClass, Class persistentClass)
            throws GoraException, ClassNotFoundException {
        String className = conf.get(STORAGE_DATA_STORE_CLASS, MONGO_STORE_CLASS);
        Class> dataStoreClass = (Class>)Class.forName(className);
        return createDataStore(conf, keyClass, persistentClass, dataStoreClass);
    }

    @SuppressWarnings("unchecked")
    public synchronized static  DataStore
    createDataStore(ImmutableConfig conf,
                    Class keyClass, Class persistentClass, Class> dataStoreClass
    ) throws GoraException {
        String crawlId = conf.get(STORAGE_CRAWL_ID, "");
        String schemaPrefix = "";
        if (!crawlId.isEmpty()) {
            schemaPrefix = crawlId + "_";
        }

        String schema;
        if (GWebPage.class.equals(persistentClass)) {
            schema = conf.get(STORAGE_SCHEMA_WEBPAGE, WEBPAGE_SCHEMA);
        } else {
            throw new UnsupportedOperationException("Unable to create storage for class " + persistentClass);
        }

        Object o = dataStores.get(schema);
        if (o == null) {
            org.apache.hadoop.conf.Configuration hadoopConf = HadoopUtils.INSTANCE.toHadoopConfiguration(conf);
            String realSchema = schemaPrefix + schema;
            hadoopConf.set(STORAGE_PREFERRED_SCHEMA_NAME, realSchema);
            DataStore dataStore = DataStoreFactory.createDataStore(dataStoreClass,
                    keyClass, persistentClass, hadoopConf, properties, schema);

            dataStores.put(realSchema, dataStore);

            Params.of(
                    "Backend data store", dataStore.getClass().getSimpleName(),
                    "realSchema", dataStore.getSchemaName()
            ).withLogger(logger).info(true);

            return dataStore;
        }

        return (DataStore) o;
    }

    public synchronized static void close() {
        dataStores.forEach((schema, store) -> {
            if (store instanceof DataStore) {
                logger.info("Closing data store <{}>", schema);
                ((DataStore) store).close();
            }
        });
        dataStores.clear();
    }
}