<commit_before>package whelk.tools

import se.kb.libris.util.marc.MarcRecord
import se.kb.libris.util.marc.io.Iso2709Deserializer
import whelk.Document
import whelk.converter.MarcJSONConverter
import whelk.converter.marc.MarcFrameConverter
import whelk.importer.MySQLImporter
import whelk.util.LegacyIntegrationTools
import whelk.util.PropertyLoader
import groovy.util.logging.Slf4j as Log

import java.nio.charset.Charset
import java.nio.file.Files
import java.nio.file.Paths
import java.sql.Connection
import java.sql.DriverManager
import java.sql.PreparedStatement
import java.sql.ResultSet

/**
 * Writes documents into a PostgreSQL load-file, which can be efficiently imported into lddb
 */
@Log
class PostgresLoadfileWriter
{
    private static final String JDBC_DRIVER = "com.mysql.jdbc.Driver"

    private final String m_exportFileName;
    private final String m_collection;
    private final Connection m_connection
    private final PreparedStatement m_statement;
    private final ResultSet m_resultSet;
    private final MarcFrameConverter m_marcFrameConverter;
    private final BufferedWriter m_mainTableWriter;
    private final BufferedWriter m_identifiersWriter;

    public PostgresLoadfileWriter(String exportFileName, String collection)
    {
        Class.forName(JDBC_DRIVER);

        Properties props = PropertyLoader.loadProperties("mysql");
        m_connection = DriverManager.getConnection(props.getProperty("mysqlConnectionUrl"));
        m_connection.setAutoCommit(false);

        switch (collection)
        {
            case "auth":
                m_statement = m_connection.prepareStatement("SELECT auth_id, data FROM auth_record WHERE auth_id > ? AND deleted = 0 ORDER BY auth_id", java.sql.ResultSet.TYPE_FORWARD_ONLY, java.sql.ResultSet.CONCUR_READ_ONLY)
                break;
            case "bib":
                m_statement = m_connection.prepareStatement("SELECT bib.bib_id, bib.data, auth.auth_id FROM bib_record bib LEFT JOIN auth_bib auth ON bib.bib_id = auth.bib_id WHERE bib.bib_id > ? AND bib.deleted = 0 ORDER BY bib.bib_id", java.sql.ResultSet.TYPE_FORWARD_ONLY, java.sql.ResultSet.CONCUR_READ_ONLY)
                break;
            case "hold":
                m_statement = m_connection.prepareStatement("SELECT mfhd_id, data, bib_id, shortname FROM mfhd_record WHERE mfhd_id > ? AND deleted = 0 ORDER BY mfhd_id", java.sql.ResultSet.TYPE_FORWARD_ONLY, java.sql.ResultSet.CONCUR_READ_ONLY)
                break;
            default:
                throw new Exception("No valid collection selected");
        }
        m_statement.setFetchSize(Integer.MIN_VALUE)
        m_statement.setInt(1, 0); // start from id 0
        m_resultSet = m_statement.executeQuery();
        m_exportFileName = exportFileName;
        m_collection = collection;
        m_marcFrameConverter = new MarcFrameConverter();
        m_mainTableWriter = Files.newBufferedWriter(Paths.get(exportFileName), Charset.forName("UTF-8"));
        m_identifiersWriter = Files.newBufferedWriter(Paths.get(exportFileName+"_identifiers"), Charset.forName("UTF-8"));
    }

    public void writePostgresLoadFile()
    {
        long startTime = System.currentTimeMillis();
        int savedDocumentsCount = 0;

        // The document we're currently building. May contain several ResultSet rows
        String currentDocumentId = "";
        HashMap documentMap = null;

        while (m_resultSet.next())
        {
            //int recordId = m_resultSet.getInt(1);
            MarcRecord record = Iso2709Deserializer.deserialize(
                    MySQLImporter.normalizeString(new String(m_resultSet.getBytes("data"), "UTF-8")).getBytes("UTF-8"));

            if (record)
            {
                def field001List = record.getControlfields("001");
                if (field001List.size() == 0)
                    continue; // skip document if no 001 control field

                def aList = record.getDatafields("599").collect { it.getSubfields("a").data }.flatten()
                if ("SUPPRESSRECORD" in aList)
                    continue; // skip document if SUPRESSED

                String oldStyleIdentifier = "/"+m_collection+"/"+field001List.get(0).getData()
                String id = LegacyIntegrationTools.generateId(oldStyleIdentifier)

                if (!id.equals(currentDocumentId))
                {
                    // New ID, store current document and begin constructing a new one
                    if (documentMap)
                    {
                        appendToLoadFile(documentMap)

                        if (++savedDocumentsCount % 1000 == 0)
                        {
                            long elapsedSeconds = (System.currentTimeMillis() - startTime) / 1000;
                            if (elapsedSeconds > 0)
                            {
                                long documentsPerSecond = savedDocumentsCount / elapsedSeconds;
                                System.out.println("Working. Currently " + savedDocumentsCount +
                                        " documents saved. Crunching " + documentsPerSecond + " docs / s");
                            }
                        }
                    }
                    currentDocumentId = id;
                    documentMap = new HashMap();
                }

                documentMap.put("record", record)
                documentMap.put("manifest", [(Document.ID_KEY):currentDocumentId,(Document.COLLECTION_KEY):m_collection, (Document.ALTERNATE_ID_KEY): [oldStyleIdentifier]])
            }

            addOaipmhSetSpecs(documentMap, m_resultSet);
        }

        cleanup();
    }

    private void addOaipmhSetSpecs(HashMap documentMap, ResultSet resultSet)
    {
        switch (m_collection)
        {
            case "bib":
                int auth_id = resultSet.getInt("auth_id")
                if (auth_id > 0)
                    documentMap.get("manifest", [:]).get("extraData", [:]).get("oaipmhSetSpecs", []).add("authority:" + auth_id)
                break;
            case "hold":
                int bib_id = resultSet.getInt("bib_id")
                String sigel = resultSet.getString("shortname")
                if (bib_id > 0)
                    documentMap.get("manifest", [:]).get("extraData", [:]).get("oaipmhSetSpecs", []).add("bibid:" + bib_id)
                if (sigel)
                    documentMap.get("manifest", [:]).get("extraData", [:]).get("oaipmhSetSpecs", []).add("location:" + sigel)
                break;<commit_msg>Do not miss the final document in the list.<commit_after>package whelk.tools

import se.kb.libris.util.marc.MarcRecord
import se.kb.libris.util.marc.io.Iso2709Deserializer
import whelk.Document
import whelk.converter.MarcJSONConverter
import whelk.converter.marc.MarcFrameConverter
import whelk.importer.MySQLImporter
import whelk.util.LegacyIntegrationTools
import whelk.util.PropertyLoader
import groovy.util.logging.Slf4j as Log

import java.nio.charset.Charset
import java.nio.file.Files
import java.nio.file.Paths
import java.sql.Connection
import java.sql.DriverManager
import java.sql.PreparedStatement
import java.sql.ResultSet

/**
 * Writes documents into a PostgreSQL load-file, which can be efficiently imported into lddb
 */
@Log
class PostgresLoadfileWriter
{
    private static final String JDBC_DRIVER = "com.mysql.jdbc.Driver"

    private final String m_exportFileName;
    private final String m_collection;
    private final Connection m_connection
    private final PreparedStatement m_statement;
    private final ResultSet m_resultSet;
    private final MarcFrameConverter m_marcFrameConverter;
    private final BufferedWriter m_mainTableWriter;
    private final BufferedWriter m_identifiersWriter;

    public PostgresLoadfileWriter(String exportFileName, String collection)
    {
        Class.forName(JDBC_DRIVER);

        Properties props = PropertyLoader.loadProperties("mysql");
        m_connection = DriverManager.getConnection(props.getProperty("mysqlConnectionUrl"));
        m_connection.setAutoCommit(false);

        switch (collection)
        {
            case "auth":
                m_statement = m_connection.prepareStatement("SELECT auth_id, data FROM auth_record WHERE auth_id > ? AND deleted = 0 ORDER BY auth_id", java.sql.ResultSet.TYPE_FORWARD_ONLY, java.sql.ResultSet.CONCUR_READ_ONLY)
                break;
            case "bib":
                m_statement = m_connection.prepareStatement("SELECT bib.bib_id, bib.data, auth.auth_id FROM bib_record bib LEFT JOIN auth_bib auth ON bib.bib_id = auth.bib_id WHERE bib.bib_id > ? AND bib.deleted = 0 ORDER BY bib.bib_id", java.sql.ResultSet.TYPE_FORWARD_ONLY, java.sql.ResultSet.CONCUR_READ_ONLY)
                break;
            case "hold":
                m_statement = m_connection.prepareStatement("SELECT mfhd_id, data, bib_id, shortname FROM mfhd_record WHERE mfhd_id > ? AND deleted = 0 ORDER BY mfhd_id", java.sql.ResultSet.TYPE_FORWARD_ONLY, java.sql.ResultSet.CONCUR_READ_ONLY)
                break;
            default:
                throw new Exception("No valid collection selected");
        }
        m_statement.setFetchSize(Integer.MIN_VALUE)
        m_statement.setInt(1, 0); // start from id 0
        m_resultSet = m_statement.executeQuery();
        m_exportFileName = exportFileName;
        m_collection = collection;
        m_marcFrameConverter = new MarcFrameConverter();
        m_mainTableWriter = Files.newBufferedWriter(Paths.get(exportFileName), Charset.forName("UTF-8"));
        m_identifiersWriter = Files.newBufferedWriter(Paths.get(exportFileName+"_identifiers"), Charset.forName("UTF-8"));
    }

    public void writePostgresLoadFile()
    {
        long startTime = System.currentTimeMillis();
        int savedDocumentsCount = 0;

        // The document we're currently building. May contain several ResultSet rows
        String currentDocumentId = "";
        HashMap documentMap = null;

        while (m_resultSet.next())
        {
            //int recordId = m_resultSet.getInt(1);
            MarcRecord record = Iso2709Deserializer.deserialize(
                    MySQLImporter.normalizeString(new String(m_resultSet.getBytes("data"), "UTF-8")).getBytes("UTF-8"));

            if (record)
            {
                def field001List = record.getControlfields("001");
                if (field001List.size() == 0)
                    continue; // skip document if no 001 control field

                def aList = record.getDatafields("599").collect { it.getSubfields("a").data }.flatten()
                if ("SUPPRESSRECORD" in aList)
                    continue; // skip document if SUPRESSED

                String oldStyleIdentifier = "/"+m_collection+"/"+field001List.get(0).getData()
                String id = LegacyIntegrationTools.generateId(oldStyleIdentifier)

                if (!id.equals(currentDocumentId))
                {
                    // New ID, store current document and begin constructing a new one
                    if (documentMap)
                    {
                        appendToLoadFile(documentMap)

                        if (++savedDocumentsCount % 1000 == 0)
                        {
                            long elapsedSeconds = (System.currentTimeMillis() - startTime) / 1000;
                            if (elapsedSeconds > 0)
                            {
                                long documentsPerSecond = savedDocumentsCount / elapsedSeconds;
                                System.out.println("Working. Currently " + savedDocumentsCount +
                                        " documents saved. Crunching " + documentsPerSecond + " docs / s");
                            }
                        }
                    }
                    currentDocumentId = id;
                    documentMap = new HashMap();
                }

                documentMap.put("record", record)
                documentMap.put("manifest", [(Document.ID_KEY):currentDocumentId,(Document.COLLECTION_KEY):m_collection, (Document.ALTERNATE_ID_KEY): [oldStyleIdentifier]])
            }

            addOaipmhSetSpecs(documentMap, m_resultSet);
        }

        // Don't forget about the last document being constructed. Must be written as well.
        if (!new HashMap().equals(documentMap))
            appendToLoadFile(documentMap)

        cleanup();
    }

    private void addOaipmhSetSpecs(HashMap documentMap, ResultSet resultSet)
    {
        switch (m_collection)
        {
            case "bib":
                int auth_id = resultSet.getInt("auth_id")
                if (auth_id > 0)
                    documentMap.get("manifest", [:]).get("extraData", [:]).get("oaipmhSetSpecs", []).add("authority:" + auth_id)
                break;
            case "hold":
                int bib_id = resultSet.getInt("bib_id")
                String sigel = resultSet.getString("shortname")
                if (bib_id > 0)
                    documentMap.get("manifest", [:]).get("extraData", [:]).get("oaipmhSetSpecs", []).add("bibid:" + bib_id)
                if (sigel)
                    documentMap.get("manifest", [:]).get("extraData", [:]).get("oaipmhSetSpecs", []).add("location:" + sigel)
                break;
{'fraction_non_alphanumeric': 0.09380411051720244, 'fraction_numerical': 0.007980124971768426, 'mean_word_length': 2.219583131362094, 'pattern_counts': {'":': 10, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1890645', 'n_tokens_mistral': 3782, 'n_tokens_neox': 3480, 'n_words': 770}