<commit_before>package org.drools.persistence.session;

import java.util.Map;

import javax.naming.InitialContext;
import javax.persistence.EntityManager;
import javax.persistence.EntityManagerFactory;
import javax.transaction.RollbackException;
import javax.transaction.Status;
import javax.transaction.Synchronization;
import javax.transaction.SystemException;
import javax.transaction.TransactionManager;
import javax.transaction.UserTransaction;

import org.apache.commons.collections.map.IdentityMap;
import org.drools.KnowledgeBase;
import org.drools.RuleBase;
import org.drools.SessionConfiguration;
import org.drools.StatefulSession;
import org.drools.impl.KnowledgeBaseImpl;
import org.drools.impl.StatefulKnowledgeSessionImpl;
import org.drools.persistence.processinstance.JPAProcessInstanceManager;
import org.drools.persistence.processinstance.JPASignalManager;
import org.drools.persistence.processinstance.JPAWorkItemManager;
import org.drools.process.command.Command;
import org.drools.process.command.CommandService;
import org.drools.reteoo.ReteooStatefulSession;
import org.drools.reteoo.ReteooWorkingMemory;
import org.drools.runtime.Environment;
import org.drools.runtime.EnvironmentName;
import org.drools.runtime.KnowledgeSessionConfiguration;
import org.drools.runtime.StatefulKnowledgeSession;

public class SingleSessionCommandService
    implements
    CommandService {

    private EntityManagerFactory        emf;
    private EntityManager               em;
    private SessionInfo                 sessionInfo;
    private JPASessionMarshallingHelper marshallingHelper;
    private StatefulSession             session;
    private StatefulKnowledgeSession    ksession;
    private Environment                 env;

    public void checkEnvironment(Environment env) {        
        if ( env.get( EnvironmentName.ENTITY_MANAGER_FACTORY ) == null ) {
            throw new IllegalArgumentException( "Environment must have an EntityManagerFactory" );
        }
        
        // @TODO log a warning that all transactions will be locally scoped using the EntityTransaction
//        if ( env.get( EnvironmentName.TRANSACTION_MANAGER ) == null ) {
//            throw new IllegalArgumentException( "Environment must have an EntityManagerFactory" );
//        }        
    }
    
    public SingleSessionCommandService(RuleBase ruleBase,
                                       SessionConfiguration conf,
                                       Environment env) {
        this( new KnowledgeBaseImpl( ruleBase ),
              (SessionConfiguration) conf,
              env );
    }

    public SingleSessionCommandService(int sessionId,
    		                           RuleBase ruleBase,
		                               SessionConfiguration conf,
		                               Environment env) {
		this( sessionId,
		      new KnowledgeBaseImpl( ruleBase ),
		      (SessionConfiguration) conf,
		      env );
	}

    public SingleSessionCommandService(KnowledgeBase kbase,
                                       KnowledgeSessionConfiguration conf,
                                       Environment env) {
        if ( conf == null ) {
            conf = new SessionConfiguration();
        }
        this.env = env;
        this.sessionInfo = new SessionInfo();

        this.session = ((KnowledgeBaseImpl) kbase).ruleBase.newStatefulSession( (SessionConfiguration) conf,
                                                                                this.env );
        
        this.ksession = new StatefulKnowledgeSessionImpl( (ReteooWorkingMemory) session );

        ((JPASignalManager) this.session.getSignalManager()).setCommandService( this );

        this.marshallingHelper = new JPASessionMarshallingHelper( this.ksession,
                                                                  conf );

        this.sessionInfo.setJPASessionMashallingHelper( this.marshallingHelper );

        this.emf = (EntityManagerFactory) env.get( EnvironmentName.ENTITY_MANAGER_FACTORY );
        this.em = emf.createEntityManager(); // how can I ensure this is an extended entity?
        //        System.out.println( ((EntityManagerImpl) this.em).getFlushMode() );
        UserTransaction ut = null;
        try {
            InitialContext ctx = new InitialContext();
            ut = (UserTransaction) ctx.lookup( "java:comp/UserTransaction" );
            ut.begin();
            registerRollbackSync();
            this.em.joinTransaction();

            this.em.persist( this.sessionInfo );

            ut.commit();
            
        } catch ( Throwable t1 ) {
            try {
                if ( ut != null ) {
                    ut.rollback();
                }
            } catch ( Throwable t2 ) {
                throw new RuntimeException( "Could not rollback transaction", t2 );
            }
            throw new RuntimeException( "Could not commit session", t1 );
        }
        
        // update the session id to be the same as the session info id
        ((ReteooStatefulSession) this.session).setId( this.sessionInfo.getId() );

        new Thread( new Runnable() {
            public void run() {
                session.fireUntilHalt();
            }
        } );
    }

    public SingleSessionCommandService(int sessionId,
                                       KnowledgeBase kbase,
                                       KnowledgeSessionConfiguration conf,
                                       Environment env) {
        if ( conf == null ) {
            conf = new SessionConfiguration();
        }

        this.env = env;

        this.emf = (EntityManagerFactory) env.get( EnvironmentName.ENTITY_MANAGER_FACTORY );
        this.em = emf.createEntityManager(); // how can I ensure this is an extended entity?
        //System.out.println(((EntityManagerImpl) this.em).getFlushMode());
        UserTransaction ut = null;
        try {
            InitialContext ctx = new InitialContext();
            ut = (UserTransaction) ctx.lookup( "java:comp/UserTransaction" );
            ut.begin();
            registerRollbackSync();
            this.em.joinTransaction();
            sessionInfo = this.em.find( SessionInfo.class, sessionId );
            ut.commit();
        } catch ( Throwable t1 ) {
            try {
                if ( ut != null ) {
                    ut.rollback();
                }
                throw new RuntimeException( "Could not find session data for id " + sessionId,
                                            t1 );
            } catch ( Throwable t2 ) {
                throw new RuntimeException( "Could not rollback transaction",
                                            t2 );
            }
        }

        this.marshallingHelper = new JPASessionMarshallingHelper( this.sessionInfo,
                                                                  kbase,
                                                                  conf,
                                                                  env );

        this.sessionInfo.setJPASessionMashallingHelper( this.marshallingHelper );        
		this.ksession = this.marshallingHelper.getObject();
		this.session = (StatefulSession) ((StatefulKnowledgeSessionImpl) ksession).session;
        ((JPASignalManager) this.session.getSignalManager()).setCommandService( this );

        new Thread( new Runnable() {
            public void run() {
                session.fireUntilHalt();
            }
        } );
    }

    public StatefulSession getSession() {
        return this.session;
    }

    public synchronized <T> T execute(Command<T> command) {
        session.halt();

        boolean localTransaction = false;
        UserTransaction ut = null;
        try {<commit_msg> - added exception is session id cannot be found<commit_after>package org.drools.persistence.session;

import java.util.Map;

import javax.naming.InitialContext;
import javax.persistence.EntityManager;
import javax.persistence.EntityManagerFactory;
import javax.transaction.RollbackException;
import javax.transaction.Status;
import javax.transaction.Synchronization;
import javax.transaction.SystemException;
import javax.transaction.TransactionManager;
import javax.transaction.UserTransaction;

import org.apache.commons.collections.map.IdentityMap;
import org.drools.KnowledgeBase;
import org.drools.RuleBase;
import org.drools.SessionConfiguration;
import org.drools.StatefulSession;
import org.drools.impl.KnowledgeBaseImpl;
import org.drools.impl.StatefulKnowledgeSessionImpl;
import org.drools.persistence.processinstance.JPAProcessInstanceManager;
import org.drools.persistence.processinstance.JPASignalManager;
import org.drools.persistence.processinstance.JPAWorkItemManager;
import org.drools.process.command.Command;
import org.drools.process.command.CommandService;
import org.drools.reteoo.ReteooStatefulSession;
import org.drools.reteoo.ReteooWorkingMemory;
import org.drools.runtime.Environment;
import org.drools.runtime.EnvironmentName;
import org.drools.runtime.KnowledgeSessionConfiguration;
import org.drools.runtime.StatefulKnowledgeSession;

public class SingleSessionCommandService
    implements
    CommandService {

    private EntityManagerFactory        emf;
    private EntityManager               em;
    private SessionInfo                 sessionInfo;
    private JPASessionMarshallingHelper marshallingHelper;
    private StatefulSession             session;
    private StatefulKnowledgeSession    ksession;
    private Environment                 env;

    public void checkEnvironment(Environment env) {        
        if ( env.get( EnvironmentName.ENTITY_MANAGER_FACTORY ) == null ) {
            throw new IllegalArgumentException( "Environment must have an EntityManagerFactory" );
        }
        
        // @TODO log a warning that all transactions will be locally scoped using the EntityTransaction
//        if ( env.get( EnvironmentName.TRANSACTION_MANAGER ) == null ) {
//            throw new IllegalArgumentException( "Environment must have an EntityManagerFactory" );
//        }        
    }
    
    public SingleSessionCommandService(RuleBase ruleBase,
                                       SessionConfiguration conf,
                                       Environment env) {
        this( new KnowledgeBaseImpl( ruleBase ),
              (SessionConfiguration) conf,
              env );
    }

    public SingleSessionCommandService(int sessionId,
    		                           RuleBase ruleBase,
		                               SessionConfiguration conf,
		                               Environment env) {
		this( sessionId,
		      new KnowledgeBaseImpl( ruleBase ),
		      (SessionConfiguration) conf,
		      env );
	}

    public SingleSessionCommandService(KnowledgeBase kbase,
                                       KnowledgeSessionConfiguration conf,
                                       Environment env) {
        if ( conf == null ) {
            conf = new SessionConfiguration();
        }
        this.env = env;
        this.sessionInfo = new SessionInfo();

        this.session = ((KnowledgeBaseImpl) kbase).ruleBase.newStatefulSession( (SessionConfiguration) conf,
                                                                                this.env );
        
        this.ksession = new StatefulKnowledgeSessionImpl( (ReteooWorkingMemory) session );

        ((JPASignalManager) this.session.getSignalManager()).setCommandService( this );

        this.marshallingHelper = new JPASessionMarshallingHelper( this.ksession,
                                                                  conf );

        this.sessionInfo.setJPASessionMashallingHelper( this.marshallingHelper );

        this.emf = (EntityManagerFactory) env.get( EnvironmentName.ENTITY_MANAGER_FACTORY );
        this.em = emf.createEntityManager(); // how can I ensure this is an extended entity?
        //        System.out.println( ((EntityManagerImpl) this.em).getFlushMode() );
        UserTransaction ut = null;
        try {
            InitialContext ctx = new InitialContext();
            ut = (UserTransaction) ctx.lookup( "java:comp/UserTransaction" );
            ut.begin();
            registerRollbackSync();
            this.em.joinTransaction();

            this.em.persist( this.sessionInfo );

            ut.commit();
            
        } catch ( Throwable t1 ) {
            try {
                if ( ut != null ) {
                    ut.rollback();
                }
            } catch ( Throwable t2 ) {
                throw new RuntimeException( "Could not rollback transaction", t2 );
            }
            throw new RuntimeException( "Could not commit session", t1 );
        }
        
        // update the session id to be the same as the session info id
        ((ReteooStatefulSession) this.session).setId( this.sessionInfo.getId() );

        new Thread( new Runnable() {
            public void run() {
                session.fireUntilHalt();
            }
        } );
    }

    public SingleSessionCommandService(int sessionId,
                                       KnowledgeBase kbase,
                                       KnowledgeSessionConfiguration conf,
                                       Environment env) {
        if ( conf == null ) {
            conf = new SessionConfiguration();
        }

        this.env = env;

        this.emf = (EntityManagerFactory) env.get( EnvironmentName.ENTITY_MANAGER_FACTORY );
        this.em = emf.createEntityManager(); // how can I ensure this is an extended entity?
        //System.out.println(((EntityManagerImpl) this.em).getFlushMode());
        UserTransaction ut = null;
        try {
            InitialContext ctx = new InitialContext();
            ut = (UserTransaction) ctx.lookup( "java:comp/UserTransaction" );
            ut.begin();
            registerRollbackSync();
            this.em.joinTransaction();
            sessionInfo = this.em.find( SessionInfo.class, sessionId );
            ut.commit();
        } catch ( Throwable t1 ) {
            try {
                if ( ut != null ) {
                    ut.rollback();
                }
                throw new RuntimeException( "Could not find session data for id " + sessionId,
                                            t1 );
            } catch ( Throwable t2 ) {
                throw new RuntimeException( "Could not rollback transaction",
                                            t2 );
            }
        }
        
        if (sessionInfo == null) {
        	throw new RuntimeException("Could not find session data for id " + sessionId);
        }

        this.marshallingHelper = new JPASessionMarshallingHelper( this.sessionInfo,
                                                                  kbase,
                                                                  conf,
                                                                  env );

        this.sessionInfo.setJPASessionMashallingHelper( this.marshallingHelper );        
		this.ksession = this.marshallingHelper.getObject();
		this.session = (StatefulSession) ((StatefulKnowledgeSessionImpl) ksession).session;
        ((JPASignalManager) this.session.getSignalManager()).setCommandService( this );

        new Thread( new Runnable() {
            public void run() {
                session.fireUntilHalt();
            }
        } );
    }

    public StatefulSession getSession() {
        return this.session;
    }

    public synchronized <T> T execute(Command<T> command) {
        session.halt();

        boolean localTransaction = false;
        UserTransaction ut = null;
        try {
{'fraction_non_alphanumeric': 0.07893561355369791, 'fraction_numerical': 0.0010209941930955269, 'mean_word_length': 1.8242926653451073, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 44, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1700307', 'n_tokens_mistral': 3766, 'n_tokens_neox': 3625, 'n_words': 848}