org.drools.persistence.SessionMarshallingHelper Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of drools-persistence-jpa Show documentation
JPA implementation for Drools
The newest version!
/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.drools.persistence;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

import org.drools.serialization.protobuf.marshalling.InternalMarshaller;
import org.drools.serialization.protobuf.marshalling.KieSessionInitializer;
import org.kie.api.KieBase;
import org.kie.api.marshalling.Marshaller;
import org.kie.api.marshalling.ObjectMarshallingStrategy;
import org.kie.api.runtime.Environment;
import org.kie.api.runtime.EnvironmentName;
import org.kie.api.runtime.KieSession;
import org.kie.api.runtime.KieSessionConfiguration;
import org.kie.internal.marshalling.MarshallerFactory;

public class SessionMarshallingHelper {

    private KieBase kbase;
    private KieSessionConfiguration       conf;
    private KieSession      			  ksession;
    private InternalMarshaller            marshaller;
    private Environment                   env;

    /**
     * Exist Info, so load session from here
     */
    public SessionMarshallingHelper( KieBase kbase,
                                     KieSessionConfiguration conf,
                                     Environment env) {
        this.kbase = kbase;
        this.conf = conf;
        this.env = env;
        ObjectMarshallingStrategy[] strategies = (ObjectMarshallingStrategy[]) env.get( EnvironmentName.OBJECT_MARSHALLING_STRATEGIES );

        // use strategies if provided in the environment
        this.marshaller = strategies != null ?
                          (InternalMarshaller) MarshallerFactory.newMarshaller( kbase, strategies ) :
                          (InternalMarshaller) MarshallerFactory.newMarshaller( kbase ) ;
    }

    /** 
     * new session, don't write now as info will request it on update callback
     */
    public SessionMarshallingHelper( KieSession ksession,
                                     KieSessionConfiguration conf) {
        this.ksession = ksession;
        this.kbase = ksession.getKieBase();
        this.conf = conf;
        this.env = ksession.getEnvironment();
        ObjectMarshallingStrategy[] strategies = (ObjectMarshallingStrategy[]) this.env.get( EnvironmentName.OBJECT_MARSHALLING_STRATEGIES );

        // use strategies if provided in the environment
        this.marshaller = strategies != null ?
                          (InternalMarshaller) MarshallerFactory.newMarshaller( kbase, strategies ) :
                          (InternalMarshaller) MarshallerFactory.newMarshaller( kbase ) ;
    }

    public byte[] getSnapshot() {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        try {
            marshaller.marshall( baos,
                                 ksession );
        } catch ( IOException e ) {
            throw new RuntimeException( "Unable to get session snapshot",
                                        e );
        }

        return baos.toByteArray();
    }

    public KieSession loadSnapshot( byte[] bytes,
                                    KieSession ksession,
                                    KieSessionInitializer initializer ) {
        this.ksession = ksession;
        ByteArrayInputStream bais = new ByteArrayInputStream( bytes );
        try {
            this.marshaller.setInitializer( initializer );
            if ( this.ksession != null ) {
                this.marshaller.unmarshall( bais,
                                            this.ksession );
            } else {
                this.ksession = this.marshaller.unmarshall( bais,
                                                            this.conf,
                                                            this.env );
            }
        } catch ( Exception e ) {
            throw new RuntimeException( "Unable to load session snapshot",
                                        e );
        } finally {
            this.marshaller.setInitializer( null );
        }
        return this.ksession;
    }


    public KieSession getObject() {
        return ksession;
    }

    public KieBase getKbase() {
        return kbase;
    }

    public KieSessionConfiguration getConf() {
        return conf;
    }
    
    public Marshaller getMarshaller() {
    	return marshaller;
    }

    
}