Files
Szpitale-graph/szpitale-graph/src/main/java/com/developx/szpitale/load/GraphLoader.java
T

357 lines
16 KiB
Java

package com.developx.szpitale.load;
import com.developx.szpitale.model.*;
import org.neo4j.driver.Driver;
import org.neo4j.driver.Session;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.*;
import java.util.stream.Collectors;
public class GraphLoader {
private static final Logger log = LoggerFactory.getLogger(GraphLoader.class);
private static final int BATCH_SIZE = 1000;
private final Driver driver;
public GraphLoader(Driver driver) {
this.driver = driver;
}
public LoadStats load(List<CanonicalDataset> datasets) {
LoadStats stats = new LoadStats();
try (Session session = driver.session()) {
for (CanonicalDataset dataset : datasets) {
stats.addVoivodeship(dataset.voivodeship());
loadHospitals(session, dataset.hospitals(), stats);
loadPeople(session, dataset.people(), stats);
loadRoles(session, dataset.roles(), stats);
loadAffiliations(session, dataset.affiliations(), stats);
loadMandates(session, dataset.mandates(), stats);
loadOrganizations(session, dataset.organizations(), stats);
loadCompanies(session, dataset.companies(), stats);
loadPersonOrganizationLinks(session, dataset.personOrganizationLinks(), stats);
loadOrganizationLinks(session, dataset.organizationLinks(), stats);
loadPersonCompanyLinks(session, dataset.personCompanyLinks(), stats);
}
}
return stats;
}
private void loadHospitals(Session session, List<Hospital> hospitals, LoadStats stats) {
List<Map<String, Object>> batch = new ArrayList<>();
for (Hospital h : hospitals) {
batch.add(Map.of(
"id", h.id(),
"name", h.name(),
"shortName", h.shortName(),
"city", h.city(),
"voivodeship", h.voivodeship(),
"legalForm", h.legalForm() != null ? h.legalForm().name() : null,
"foundingBody", h.foundingBody(),
"supervisoryBodyType", h.supervisoryBodyType() != null ? h.supervisoryBodyType().name() : null,
"website", h.website(),
"sourceUrls", h.sourceUrls()
));
}
for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
String cypher = "UNWIND $batch AS h " +
"MERGE (hospital:Hospital {id: h.id}) " +
"SET hospital += {name:h.name, shortName:h.shortName, city:h.city, " +
" voivodeship:h.voivodeship, legalForm:h.legalForm, " +
" foundingBody:h.foundingBody, supervisoryBodyType:h.supervisoryBodyType, " +
" website:h.website} " +
"MERGE (v:Voivodeship {slug: h.voivodeship}) " +
"MERGE (hospital)-[:W_WOJEWODZTWIE]->(v)";
session.run(cypher, Map.of("batch", subList));
stats.addHospitals(subList.size());
}
}
private void loadPeople(Session session, List<Person> people, LoadStats stats) {
List<Map<String, Object>> batch = new ArrayList<>();
for (Person p : people) {
batch.add(Map.of(
"id", p.id(),
"fullName", p.fullName(),
"displayName", p.displayName(),
"titles", p.titles() != null ? p.titles() : Collections.emptyList(),
"sourceUrls", p.sourceUrls()
));
}
for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
String cypher = "UNWIND $batch AS p " +
"MERGE (person:Person {id: p.id}) " +
"SET person += {fullName:p.fullName, displayName:p.displayName, " +
" titles:p.titles, sourceUrls:p.sourceUrls}";
session.run(cypher, Map.of("batch", subList));
stats.addPeople(subList.size());
}
}
private void loadRoles(Session session, List<Role> roles, LoadStats stats) {
Map<String, List<Role>> byRelType = new LinkedHashMap<>();
for (Role r : roles) {
String relType = roleTypeToRelName(r);
byRelType.computeIfAbsent(relType, k -> new ArrayList<>()).add(r);
}
for (Map.Entry<String, List<Role>> entry : byRelType.entrySet()) {
String relType = entry.getKey();
List<Role> items = entry.getValue();
for (int i = 0; i < items.size(); i += BATCH_SIZE) {
List<Role> subList = items.subList(i, Math.min(i + BATCH_SIZE, items.size()));
String cypher = "UNWIND $batch AS r " +
"MATCH (h:Hospital {id: r.hospitalId}) " +
"MATCH (p:Person {id: r.personId}) " +
"MERGE (h)-[rrel:`" + escapeBacktick(relType) + "`]->(p) " +
"SET rrel += {roleType: r.roleType, roleLabel: r.roleLabel, " +
" organ: r.organ, status: r.status, " +
" validFrom: r.validFrom, validTo: r.validTo, note: r.note}";
session.run(cypher, Map.of("batch", subList.stream()
.map(r -> Map.of(
"hospitalId", r.hospitalId(),
"personId", r.personId(),
"roleType", r.roleType().name(),
"roleLabel", r.roleLabel(),
"organ", r.organ().name(),
"status", r.status().name(),
"validFrom", r.validFrom(),
"validTo", r.validTo(),
"note", r.note()
))
.collect(Collectors.toList())));
stats.addRoles(subList.size());
}
}
}
private void loadAffiliations(Session session, List<Affiliation> affiliations, LoadStats stats) {
for (int i = 0; i < affiliations.size(); i += BATCH_SIZE) {
List<Affiliation> subList = affiliations.subList(
i, Math.min(i + BATCH_SIZE, affiliations.size()));
String cypher = "UNWIND $batch AS a " +
"MATCH (p:Person {id: a.personId}) " +
"MERGE (party:Party {name: a.name}) SET party.type = a.type " +
"MERGE (p)-[m:CZLONEK_PARTII]->(party) " +
"SET m += {confidence: a.confidence, note: a.note, sourceUrls: a.sourceUrls}";
session.run(cypher, Map.of("batch", subList.stream()
.map(a -> Map.of(
"personId", a.personId(),
"name", a.name(),
"type", a.type().name(),
"confidence", a.confidence().name(),
"note", a.note(),
"sourceUrls", a.sourceUrls()
))
.collect(Collectors.toList())));
stats.addAffiliations(subList.size());
}
}
private void loadMandates(Session session, List<Mandate> mandates, LoadStats stats) {
if (mandates.isEmpty()) return;
for (int i = 0; i < mandates.size(); i += BATCH_SIZE) {
List<Mandate> subList = mandates.subList(
i, Math.min(i + BATCH_SIZE, mandates.size()));
String cypher = "UNWIND $batch AS m " +
"MATCH (p:Person {id: m.personId}) " +
"MERGE (g:GovBody {name: m.body}) " +
" SET g.type = m.mandateType, g.voivodeship = m.voivodeship, g.term = m.term " +
"MERGE (p)-[:PELNI_MANDAT {mandateType: m.mandateType, term: m.term}]->(g)";
session.run(cypher, Map.of("batch", subList.stream()
.map(m -> Map.of(
"personId", m.personId(),
"body", m.body(),
"mandateType", m.mandateType().name(),
"term", m.term(),
"voivodeship", m.voivodeship(),
"sourceUrls", m.sourceUrls()
))
.collect(Collectors.toList())));
stats.addMandates(subList.size());
}
}
private void loadOrganizations(Session session, List<Organization> organizations, LoadStats stats) {
if (organizations.isEmpty()) return;
List<Map<String, Object>> batch = new ArrayList<>();
for (Organization o : organizations) {
batch.add(Map.of(
"id", o.id(),
"name", o.name(),
"krs", o.krs(),
"nip", o.nip(),
"sourceUrls", o.sourceUrls()
));
}
for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
String cypher = "UNWIND $batch AS o " +
"MERGE (org:Organization {id: o.id}) " +
"SET org += {name: o.name, krs: o.krs, nip: o.nip, sourceUrls: o.sourceUrls}";
session.run(cypher, Map.of("batch", subList));
stats.addOrganizations(subList.size());
}
}
private void loadCompanies(Session session, List<Company> companies, LoadStats stats) {
if (companies.isEmpty()) return;
List<Map<String, Object>> batch = new ArrayList<>();
for (Company c : companies) {
batch.add(Map.of(
"id", c.id(),
"name", c.name(),
"krs", c.krs(),
"nip", c.nip(),
"ceidgId", c.ceidgId(),
"sourceUrls", c.sourceUrls()
));
}
for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
String cypher = "UNWIND $batch AS c " +
"MERGE (comp:Company {id: c.id}) " +
"SET comp += {name: c.name, krs: c.krs, nip: c.nip, ceidgId: c.ceidgId, sourceUrls: c.sourceUrls}";
session.run(cypher, Map.of("batch", subList));
stats.addCompanies(subList.size());
}
}
private void loadPersonOrganizationLinks(Session session, List<PersonOrganizationLink> links, LoadStats stats) {
if (links.isEmpty()) return;
List<Map<String, Object>> batch = new ArrayList<>();
for (PersonOrganizationLink l : links) {
batch.add(Map.of(
"personId", l.personId(),
"organizationId", l.organizationId(),
"roleLabel", l.roleLabel(),
"basis", l.basis().name(),
"sourceUrls", l.sourceUrls()
));
}
for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
String cypher = "UNWIND $batch AS l " +
"MATCH (p:Person {id: l.personId}), (org:Organization {id: l.organizationId}) " +
"MERGE (p)-[m:CZLONEK_ORGANIZACJI]->(org) " +
"SET m += {roleLabel: l.roleLabel, basis: l.basis, sourceUrls: l.sourceUrls}";
session.run(cypher, Map.of("batch", subList));
stats.addPersonOrganizationLinks(subList.size());
}
}
private void loadOrganizationLinks(Session session, List<OrganizationLink> links, LoadStats stats) {
if (links.isEmpty()) return;
List<Map<String, Object>> batch = new ArrayList<>();
for (OrganizationLink l : links) {
batch.add(Map.of(
"fromId", l.fromOrganizationId(),
"toId", l.toOrganizationId(),
"basis", l.basis().name(),
"note", l.note(),
"sourceUrls", l.sourceUrls()
));
}
for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
String cypher = "UNWIND $batch AS l " +
"MATCH (a:Organization {id: l.fromId}), (b:Organization {id: l.toId}) " +
"MERGE (a)-[m:POWIAZANA_Z]->(b) " +
"SET m += {basis: l.basis, note: l.note, sourceUrls: l.sourceUrls}";
session.run(cypher, Map.of("batch", subList));
stats.addOrganizationLinks(subList.size());
}
}
private void loadPersonCompanyLinks(Session session, List<PersonCompanyLink> links, LoadStats stats) {
if (links.isEmpty()) return;
List<Map<String, Object>> batch = new ArrayList<>();
for (PersonCompanyLink l : links) {
batch.add(Map.of(
"personId", l.personId(),
"companyId", l.companyId(),
"roleLabel", l.roleLabel(),
"basis", l.basis().name(),
"sourceUrls", l.sourceUrls()
));
}
for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
String cypher = "UNWIND $batch AS l " +
"MATCH (p:Person {id: l.personId}), (comp:Company {id: l.companyId}) " +
"MERGE (p)-[m:POWIAZANY_Z_FIRMA]->(comp) " +
"SET m += {roleLabel: l.roleLabel, basis: l.basis, sourceUrls: l.sourceUrls}";
session.run(cypher, Map.of("batch", subList));
stats.addPersonCompanyLinks(subList.size());
}
}
private String roleTypeToRelName(Role r) {
String organ = r.organ().name();
String role = r.roleType().name();
if ("DYREKCJA".equals(organ) && "DYREKTOR".equals(role)) return "DYREKTOR";
if ("DYREKCJA".equals(organ) && ("ZASTEPCA_DYREKTORA".equals(role) || "GLOWNY_KSIEGOWY".equals(role) || "PRZELOZONA_PIELEGNIARKOW".equals(role))) return "ZASTEPCA_DYREKTORA";
if ("PRZEWODNICZACY_ORGANU".equals(role)) return "PRZEWODNICZY_ORGANOWI";
return "CZLONEK_ORGANU";
}
private String escapeBacktick(String s) {
return s.replace("`", "\\`");
}
public static class LoadStats {
public int hospitalCount;
public int personCount;
public int roleCount;
public int affiliationCount;
public int mandateCount;
public int organizationCount;
public int companyCount;
public int personOrganizationLinkCount;
public int organizationLinkCount;
public int personCompanyLinkCount;
public final List<String> voivodeships = new ArrayList<>();
public void addVoivodeship(String v) { voivodeships.add(v); }
public void addHospitals(int count) { hospitalCount += count; }
public void addPeople(int count) { personCount += count; }
public void addRoles(int count) { roleCount += count; }
public void addAffiliations(int count) { affiliationCount += count; }
public void addMandates(int count) { mandateCount += count; }
public void addOrganizations(int count) { organizationCount += count; }
public void addCompanies(int count) { companyCount += count; }
public void addPersonOrganizationLinks(int count) { personOrganizationLinkCount += count; }
public void addOrganizationLinks(int count) { organizationLinkCount += count; }
public void addPersonCompanyLinks(int count) { personCompanyLinkCount += count; }
@Override
public String toString() {
return String.format(
"Load stats - voivodeships: %s, hospitals: %d, people: %d, roles: %d, affiliations: %d, mandates: %d, organizations: %d, companies: %d, personOrgLinks: %d, orgLinks: %d, personCompanyLinks: %d",
String.join(", ", voivodeships), hospitalCount, personCount, roleCount, affiliationCount, mandateCount,
organizationCount, companyCount, personOrganizationLinkCount, organizationLinkCount, personCompanyLinkCount
);
}
}
}