package com.vladmihalcea.book.hpjp.hibernate.logging;

import com.vladmihalcea.book.hpjp.util.AbstractPostgreSQLIntegrationTest;
import com.vladmihalcea.book.hpjp.util.exception.ExceptionUtil;
import org.hibernate.LockMode;
import org.hibernate.LockOptions;
import org.hibernate.Session;
import org.hibernate.jpa.QueryHints;
import org.junit.Test;
import org.slf4j.MDC;

import javax.persistence.*;
import java.util.Properties;

import static org.junit.Assert.assertTrue;

/**
 * @author Vlad Mihalcea
 */
public class MDCLoggingTest extends AbstractPostgreSQLIntegrationTest {

    @Override
    protected Class<?>[] entities() {
        return new Class<?>[] {
            Post.class
        };
    }

    @Override
    protected void additionalProperties(Properties properties) {
        properties.put( "hibernate.jdbc.batch_size", "5" );
    }

    @Override
    protected boolean proxyDataSource() {
        return true;
    }

    @Override
    protected void afterInit() {
        doInJPA(entityManager -> {
            Post post = new Post();
            post.setId( 1L );
            post.setTitle( "Post it!" );

            entityManager.persist(post);
        });
    }

    @Test
    public void testWithoutMDC() {
        doInJPA(entityManager -> {
            Post post = entityManager.createQuery(
                "select p " +
                "from Post p " +
                "where p.id = :id", Post.class)
            .setParameter("id", 1L)
            .setLockMode(LockModeType.PESSIMISTIC_WRITE)
            .getSingleResult();

            try {
                executeSync(() -> {
                    doInJPA(_entityManager -> {
                        Post _post = (Post) _entityManager.createQuery(
                            "select p " +
                            "from Post p " +
                            "where p.id = :id", Post.class)
                        .setParameter("id", 1L)
                        .unwrap(org.hibernate.query.Query.class)
                        .setLockOptions(
                            new LockOptions()
                                .setLockMode(LockMode.PESSIMISTIC_WRITE)
                                .setTimeOut(LockOptions.NO_WAIT)
                        )
                        .getSingleResult();
                    });
                });
            } catch (Exception expected) {
                assertTrue(ExceptionUtil.rootCause(expected).getMessage().contains("could not obtain lock on row in relation"));
            }
        });
    }

    @Test
    public void testWithMDC() {
        try {
            doInJPA(entityManager -> {
                MDC.put("txId", String.format(" TxId: [%s]", transactionId(entityManager)));

                Post post = entityManager.createQuery(
                    "select p " +
                    "from Post p " +
                    "where p.id = :id", Post.class)
                .setParameter("id", 1L)
                .setLockMode(LockModeType.PESSIMISTIC_WRITE)
                .getSingleResult();

                try {
                    executeSync(() -> {
                        try {
                            doInJPA(_entityManager -> {
                                MDC.put("txId", String.format(" TxId: [%s]", transactionId(_entityManager)));

                                Post _post = (Post) _entityManager.createQuery(
                                    "select p " +
                                    "from Post p " +
                                    "where p.id = :id", Post.class)
                                .setParameter("id", 1L)
                                .unwrap(org.hibernate.query.Query.class)
                                .setLockOptions(
                                        new LockOptions()
                                                .setLockMode(LockMode.PESSIMISTIC_WRITE)
                                                .setTimeOut(LockOptions.NO_WAIT)
                                )
                                .getSingleResult();
                            });
                        } finally {
                            MDC.remove("txId");
                        }
                    });
                } catch (Exception expected) {
                    assertTrue(ExceptionUtil.rootCause(expected).getMessage().contains("could not obtain lock on row in relation"));
                }
            });
        } finally {
            MDC.remove("txId");
        }
    }

    @Entity(name = "Post")
    @Table(name = "post")
    public static class Post {

        @Id
        private Long id;

        private String title;

        @Version
        private int version;

        public Long getId() {
            return id;
        }

        public void setId(Long id) {
            this.id = id;
        }

        public String getTitle() {
            return title;
        }

        public void setTitle(String title) {
            this.title = title;
        }
    }
}

@@ -2565,34 +2565,16 @@
 MDC() %7B%0A
-        try %7B%0A    
         
@@ -2612,35 +2612,113 @@
         
-    MDC.put(%22txId%22,
+try(MDC.MDCCloseable closable = MDC.putCloseable(%0A                    %22txId%22,%0A                   
  String.
@@ -2760,34 +2760,53 @@
 (entityManager))
-);
+%0A                )) %7B
 %0A%0A              
@@ -3185,35 +3185,57 @@
                 
-try
+doInJPA(_entityManager -%3E
  %7B%0A             
@@ -3253,83 +3253,138 @@
     
-doInJPA(_entityManager -%3E %7B%0A                                MDC.put(%22txId%22,
+try(MDC.MDCCloseable _closable = MDC.putCloseable(%0A                                    %22txId%22,%0A                                   
  Str
@@ -3439,18 +3439,53 @@
 anager))
-);
+%0A                                )) %7B
 %0A%0A      
@@ -3935,36 +3935,32 @@
                 
-    
 new LockOptions(
@@ -3957,36 +3957,24 @@
 ckOptions()%0A
-            
             
@@ -4034,36 +4034,24 @@
 STIC_WRITE)%0A
-            
             
@@ -4214,34 +4214,32 @@
                %7D
-);
 %0A               
@@ -4252,92 +4252,10 @@
    %7D
- finally %7B%0A                            MDC.remove(%22txId%22);%0A                        %7D
+);
 %0A   
@@ -4357,54 +4357,154 @@
 rue(
-ExceptionUtil.rootCause(expected).getMessage()
+%0A                        ExceptionUtil%0A                        .rootCause(expected)%0A                        .getMessage()%0A                        
 .con
@@ -4544,34 +4544,32 @@
 ow in relation%22)
-);
 %0A               
@@ -4569,30 +4569,31 @@
         
-%7D%0A
     
+);%0A
         
 %7D);%0A    
@@ -4576,36 +4576,32 @@
      );%0A        
-%7D);%0A
         %7D finall
@@ -4597,18 +4597,8 @@
    %7D
- finally %7B
 %0A   
@@ -4610,27 +4610,9 @@
     
-MDC.remove(%22txId%22);
+%7D
 %0A   
@@ -4613,24 +4613,26 @@
  %7D%0A        %7D
+);
 %0A    %7D%0A%0A    

Update example to include MDCClosable