[WIP] Substation preloading

network-store-client/src/main/java/com/powsybl/network/store/client/NetworkStoreService.java

@@ -95,6 +95,8 @@ private static NetworkStoreClient createStoreClient(RestClient restClient, Prelo
                 return cachedClient;
             case COLLECTION:
                 return new PreloadingNetworkStoreClient(cachedClient);
+            case SUBSTATION:
+                return new SubstationPreloadingNetworkStoreClient(cachedClient);
             default:
                 throw new IllegalStateException("Unknown preloading strategy: " + preloadingStrategy);
         }

network-store-client/src/main/java/com/powsybl/network/store/client/PreloadingStrategy.java

@@ -11,5 +11,6 @@
  */
 public enum PreloadingStrategy {
     NONE,
-    COLLECTION
+    COLLECTION,
+    SUBSTATION
 }

network-store-client/src/main/java/com/powsybl/network/store/client/RestNetworkStoreClient.java

@@ -730,6 +730,11 @@ public Optional<Resource<IdentifiableAttributes>> getIdentifiable(UUID networkUu
         return get("identifiable", "/networks/{networkUuid}/{variantNum}/identifiables/{id}", networkUuid, variantNum, id);
     }
 
+    @Override
+    public List<Resource<IdentifiableAttributes>> getIdentifiablesWithSameSubstationAs(UUID networkUuid, int variantNum, String id) {
+        return getAll("identifiable", "/networks/{networkUuid}/{variantNum}/identifiables?same_substation_as={id}", networkUuid, variantNum, id);
+    }
+
     @Override
     public void flush() {
         // nothing to do

network-store-client/src/main/java/com/powsybl/network/store/client/SubstationPreloadingNetworkStoreClient.java

@@ -0,0 +1,39 @@
+/**
+ * Copyright (c) 2022, RTE (http://www.rte-france.com)
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ */
+package com.powsybl.network.store.client;
+
+import com.powsybl.network.store.iidm.impl.AbstractForwardingNetworkStoreClient;
+import com.powsybl.network.store.iidm.impl.NetworkStoreClient;
+import com.powsybl.network.store.model.BusbarSectionAttributes;
+import com.powsybl.network.store.model.Resource;
+import com.powsybl.network.store.model.SubstationAttributes;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.UUID;
+
+/**
+ * @author Geoffroy Jamgotchian <geoffroy.jamgotchian at rte-france.com>
+ */
+public class SubstationPreloadingNetworkStoreClient extends AbstractForwardingNetworkStoreClient implements NetworkStoreClient {
+
+    public SubstationPreloadingNetworkStoreClient(NetworkStoreClient delegate) {
+        super(delegate);
+    }
+
+    @Override
+    public List<Resource<BusbarSectionAttributes>> getVoltageLevelBusbarSections(UUID networkUuid, int variantNum, String voltageLevelId) {
+        delegate.getIdentifiablesWithSameSubstationAs(networkUuid, variantNum, voltageLevelId);
+        return super.getVoltageLevelBusbarSections(networkUuid, variantNum, voltageLevelId);
+    }
+
+    @Override
+    public Optional<Resource<SubstationAttributes>> getSubstation(UUID networkUuid, int variantNum, String substationId) {
+        delegate.getIdentifiablesWithSameSubstationAs(networkUuid, variantNum, substationId);
+        return super.getSubstation(networkUuid, variantNum, substationId);
+    }
+}

network-store-iidm-impl/src/main/java/com/powsybl/network/store/iidm/impl/NetworkStoreClient.java

@@ -277,6 +277,8 @@ public interface NetworkStoreClient {
 
     Optional<Resource<IdentifiableAttributes>> getIdentifiable(UUID networkUuid, int variantNum, String id);
 
+    List<Resource<IdentifiableAttributes>> getIdentifiablesWithSameSubstationAs(UUID networkUuid, int variantNum, String id);
+
     void flush();
 
 }

network-store-iidm-impl/src/main/java/com/powsybl/network/store/iidm/impl/OfflineNetworkStoreClient.java

@@ -574,6 +574,11 @@ public Optional<Resource<IdentifiableAttributes>> getIdentifiable(UUID networkUu
         return Optional.empty();
     }
 
+    @Override
+    public List<Resource<IdentifiableAttributes>> getIdentifiablesWithSameSubstationAs(UUID networkUuid, int variantNum, String id) {
+        return Collections.emptyList();
+    }
+
     @Override
     public void flush() {
         // nothing to do

network-store-integration-test/src/test/java/com/powsybl/network/store/integration/NetworkStoreIT.java

@@ -91,11 +91,15 @@ private NetworkStoreService createNetworkStoreService() {
     }
 
     private NetworkStoreService createNetworkStoreService(RestClientMetrics metrics) {
+        return createNetworkStoreService(metrics, PreloadingStrategy.NONE);
+    }
+
+    private NetworkStoreService createNetworkStoreService(RestClientMetrics metrics, PreloadingStrategy preloadingStrategy) {
         RestClient restClient = new RestClientImpl(getBaseUrl());
         if (metrics != null) {
             restClient = new TestRestClient(restClient, metrics);
         }
-        return new NetworkStoreService(restClient, PreloadingStrategy.NONE);
+        return new NetworkStoreService(restClient, preloadingStrategy);
     }
 
     @Test
@@ -5080,4 +5084,20 @@ public void testNanValues() {
             assertEquals(Double.NaN, generator.getTerminal().getQ(), .0001);
         }
     }
+
+    @Test
+    public void testSubstationPreloading() {
+        UUID networkUuid;
+        try (NetworkStoreService service = createNetworkStoreService()) {
+            Network network = EurostagTutorialExample1Factory.create(service.getNetworkFactory());
+            networkUuid = service.getNetworkUuid(network);
+            service.flush(network);
+        }
+
+        RestClientMetrics metrics = new RestClientMetrics();
+        try (NetworkStoreService service = createNetworkStoreService(metrics, PreloadingStrategy.SUBSTATION)) {
+            Network network = service.getNetwork(networkUuid);
+            network.getSubstation("P1");
+        }
+    }
 }

network-store-server/src/main/java/com/powsybl/network/store/server/Mappings.java

@@ -93,6 +93,22 @@ public class Mappings {
                                                    twoWindingsTransformerMappings,
                                                    threeWindingsTransformerMappings);
 
+    private final List<TableMapping> noneContainerTableMapping = List.of(lineMappings,
+                                                                         loadMappings,
+                                                                         generatorMappings,
+                                                                         switchMappings,
+                                                                         batteryMappings,
+                                                                         busbarSectionMappings,
+                                                                         configuredBusMappings,
+                                                                         danglingLineMappings,
+                                                                         shuntCompensatorMappings,
+                                                                         vscConverterStationMappings,
+                                                                         vscConverterStationMappings,
+                                                                         lccConverterStationMappings,
+                                                                         staticVarCompensatorMappings,
+                                                                         twoWindingsTransformerMappings,
+                                                                         threeWindingsTransformerMappings);
+
     private final Map<String, TableMapping> mappingByTable = new LinkedHashMap<>();
 
     private static final String VOLTAGE_LEVEL_ID = "voltageLevelId";
@@ -123,6 +139,10 @@ public class Mappings {
     private static final String REACTIVE_CAPABILITY_CURVE = "reactiveCapabilityCurve";
     private static final String REGULATION_TERMINAL = "regulatingTerminal";
 
+    public List<TableMapping> getNoneContainerTableMapping() {
+        return noneContainerTableMapping;
+    }
+
     public TableMapping getTableMapping(String table) {
         Objects.requireNonNull(table);
         TableMapping tableMapping = mappingByTable.get(table);

network-store-server/src/main/java/com/powsybl/network/store/server/NetworkStoreController.java

@@ -1202,4 +1202,17 @@ public ResponseEntity<TopLevelDocument<IdentifiableAttributes>> getIdentifiable(
                                                                                     @Parameter(description = "Identifiable ID", required = true) @PathVariable("id") String id) {
         return get(() -> repository.getIdentifiable(networkId, variantNum, id));
     }
+
+    @GetMapping(value = "/{networkUuid}/{variantNum}/identifiables", produces = APPLICATION_JSON_VALUE)
+    @Operation(summary = "Get all identifiables with same substation as identifiable with specified id")
+    @ApiResponses(value = {
+            @ApiResponse(responseCode = "200", description = "Successfully get all identifiables with same substation as identifiable with specified id"),
+            @ApiResponse(responseCode = "404", description = "The specified identifiable has not been found")
+    })
+    public ResponseEntity<TopLevelDocument<IdentifiableAttributes>> getIdentifiablesWithSameSubstationAs(@Parameter(description = "Network UUID", required = true) @PathVariable("networkUuid") UUID networkUuid,
+                                                                                                         @Parameter(description = "Variant number", required = true) @PathVariable("variantNum") int variantNum,
+                                                                                                         @Parameter(description = "An identifiable to get the substation", required = true) @RequestParam("same_substation_as") String id) {
+        return ResponseEntity.ok()
+                .body(TopLevelDocument.of(repository.getIdentifiablesWithSameSubstationAs(networkUuid, variantNum, id)));
+    }
 }

network-store-server/src/main/java/com/powsybl/network/store/server/NetworkStoreRepository.java

@@ -508,23 +508,36 @@ private <T extends IdentifiableAttributes> List<Resource<T>> getIdentifiables(UU
     }
 
     private <T extends IdentifiableAttributes> List<Resource<T>> getIdentifiablesInContainer(UUID networkUuid, int variantNum, String containerId,
+                                                                                             String containerColumn,
+                                                                                             TableMapping tableMapping) {
+        return getIdentifiablesInContainer(networkUuid, variantNum, List.of(containerId), Set.of(containerColumn), tableMapping);
+    }
+
+    private <T extends IdentifiableAttributes> List<Resource<T>> getIdentifiablesInContainer(UUID networkUuid, int variantNum, List<String> containerIds,
                                                                                              Set<String> containerColumns,
                                                                                              TableMapping tableMapping) {
         try (var connection = dataSource.getConnection()) {
-            var preparedStmt = connection.prepareStatement(QueryCatalog.buildGetIdentifiablesInContainerQuery(tableMapping.getTable(), tableMapping.getColumnsMapping().keySet(), containerColumns));
+            var preparedStmt = connection.prepareStatement(QueryCatalog.buildGetIdentifiablesInContainerQuery(tableMapping.getTable(), tableMapping.getColumnsMapping().keySet(), containerColumns, containerIds.size()));
             preparedStmt.setObject(1, networkUuid);
             preparedStmt.setInt(2, variantNum);
             for (int i = 0; i < containerColumns.size(); i++) {
-                preparedStmt.setString(3 + i, containerId);
+                for (int j = 0; j < containerIds.size(); j++) {
+                    String containerId = containerIds.get(j);
+                    preparedStmt.setString(3 + i * containerIds.size() + j, containerId);
+                }
             }
             return getIdentifiablesInternal(variantNum, preparedStmt, tableMapping);
         } catch (SQLException e) {
             throw new UncheckedSqlException(e);
         }
     }
 
+    private <T extends IdentifiableAttributes> List<Resource<T>> getIdentifiablesInVoltageLevel(UUID networkUuid, int variantNum, List<String> voltageLevelIds, TableMapping tableMapping) {
+        return getIdentifiablesInContainer(networkUuid, variantNum, voltageLevelIds, tableMapping.getVoltageLevelIdColumns(), tableMapping);
+    }
+
     private <T extends IdentifiableAttributes> List<Resource<T>> getIdentifiablesInVoltageLevel(UUID networkUuid, int variantNum, String voltageLevelId, TableMapping tableMapping) {
-        return getIdentifiablesInContainer(networkUuid, variantNum, voltageLevelId, tableMapping.getVoltageLevelIdColumns(), tableMapping);
+        return getIdentifiablesInContainer(networkUuid, variantNum, List.of(voltageLevelId), tableMapping.getVoltageLevelIdColumns(), tableMapping);
     }
 
     public <T extends IdentifiableAttributes & Contained> void updateIdentifiables(UUID networkUuid, List<Resource<T>> resources,
@@ -628,7 +641,7 @@ public void updateVoltageLevels(UUID networkUuid, List<Resource<VoltageLevelAttr
     }
 
     public List<Resource<VoltageLevelAttributes>> getVoltageLevels(UUID networkUuid, int variantNum, String substationId) {
-        return getIdentifiablesInContainer(networkUuid, variantNum, substationId, Set.of(SUBSTATION_ID), mappings.getVoltageLevelMappings());
+        return getIdentifiablesInContainer(networkUuid, variantNum, substationId, SUBSTATION_ID, mappings.getVoltageLevelMappings());
     }
 
     public Optional<Resource<VoltageLevelAttributes>> getVoltageLevel(UUID networkUuid, int variantNum, String voltageLevelId) {
@@ -1085,4 +1098,76 @@ public Optional<Resource<IdentifiableAttributes>> getIdentifiable(UUID networkUu
         }
         return Optional.empty();
     }
+
+    private List<Resource<VoltageLevelAttributes>> getVoltageLevelsInSubstation(UUID networkUuid, int variantNum, String substationId) {
+        try (var connection = dataSource.getConnection()) {
+            var voltageLevelMapping = mappings.getVoltageLevelMappings();
+            var preparedStmt = connection.prepareStatement(QueryCatalog.buildGetVoltageLevelsInSubstationQuery(voltageLevelMapping.getColumnsMapping().keySet()));
+            preparedStmt.setObject(1, networkUuid);
+            preparedStmt.setInt(2, variantNum);
+            preparedStmt.setString(3, substationId);
+            List<Resource<VoltageLevelAttributes>> resources = new ArrayList<>(3);
+            try (ResultSet resultSet = preparedStmt.executeQuery()) {
+                if (resultSet.next()) {
+                    VoltageLevelAttributes attributes = new VoltageLevelAttributes();
+                    MutableInt columnIndex = new MutableInt(2);
+                    voltageLevelMapping.getColumnsMapping().forEach((columnName, columnMapping) -> {
+                        bindAttributes(resultSet, columnIndex.getValue(), columnMapping, attributes);
+                        columnIndex.increment();
+                    });
+                    resources.add(Resource.voltageLevelBuilder()
+                            .id(resultSet.getString(1))
+                            .variantNum(variantNum)
+                            .attributes(attributes)
+                            .build());
+                }
+            }
+            return resources;
+        } catch (SQLException e) {
+            throw new UncheckedSqlException(e);
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private static <T extends IdentifiableAttributes> Resource<IdentifiableAttributes> downcast(Resource<T> r) {
+        return (Resource<IdentifiableAttributes>) r;
+    }
+
+    public List<Resource<IdentifiableAttributes>> getIdentifiablesWithSameSubstationAs(UUID networkUuid, int variantNum, String id) {
+        Map<String, Resource<IdentifiableAttributes>> resourcesById = new LinkedHashMap<>();
+
+        Resource<IdentifiableAttributes> resource = getIdentifiable(networkUuid, variantNum, id)
+                .orElseThrow(() -> new PowsyblException("Identifiable '" + id + "' not found"));
+        resourcesById.put(resource.getId(), resource);
+
+        // get substation from given identifiable
+        String substationId;
+        if (resource.getType() == ResourceType.SUBSTATION) {
+            substationId = id;
+        } else if (resource.getType() == ResourceType.VOLTAGE_LEVEL) {
+            substationId = ((VoltageLevelAttributes) resource.getAttributes()).getSubstationId();
+        } else if (resource.getAttributes() instanceof InjectionAttributes) {
+            String voltageLevelId = ((InjectionAttributes) resource.getAttributes()).getVoltageLevelId();
+            Resource<VoltageLevelAttributes> voltageLevelResource = getVoltageLevel(networkUuid, variantNum, voltageLevelId).orElseThrow();
+            substationId = voltageLevelResource.getAttributes().getSubstationId();
+        } else {
+            throw new PowsyblException("TODO");
+        }
+
+        // get list of voltage levels contained in this substation
+        var voltageLevelResources = getVoltageLevelsInSubstation(networkUuid, variantNum, substationId);
+        for (var voltageLevelResource : voltageLevelResources) {
+            resourcesById.put(voltageLevelResource.getId(), downcast(voltageLevelResource));
+        }
+
+        // get identifiables contained in these voltage levels
+        List<String> voltageLevelIds = voltageLevelResources.stream().map(Resource::getId).collect(Collectors.toList());
+        for (var tableMapping : mappings.getNoneContainerTableMapping()) {
+            for (Resource<IdentifiableAttributes> identifiableResource : getIdentifiablesInVoltageLevel(networkUuid, variantNum, voltageLevelIds, tableMapping)) {
+                resourcesById.put(identifiableResource.getId(), identifiableResource);
+            }
+        }
+
+        return new ArrayList<>(resourcesById.values());
+    }
 }

network-store-server/src/main/java/com/powsybl/network/store/server/QueryCatalog.java

@@ -9,9 +9,12 @@
 import com.powsybl.network.store.model.Resource;
 
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static com.powsybl.network.store.server.Mappings.*;
+
 import static com.powsybl.network.store.server.Mappings.ELEMENT_TABLES;
 import static com.powsybl.network.store.server.Mappings.NETWORK_TABLE;
 
@@ -22,6 +25,7 @@ public final class QueryCatalog {
 
     static final String VARIANT_ID_COLUMN = "variantId";
     static final String UUID_COLUMN = "uuid";
+    static final String SUBSTATION_ID_COLUMN = "substationId";
     static final String NETWORK_UUID_COLUMN = "networkUuid";
     static final String VARIANT_NUM_COLUMN = "variantNum";
     static final String ID_COLUMN = "id";
@@ -59,7 +63,8 @@ public static String buildGetIdentifiablesQuery(String tableName, Collection<Str
                 " and " + VARIANT_NUM_COLUMN + " = ?";
     }
 
-    public static String buildGetIdentifiablesInContainerQuery(String tableName, Collection<String> columns, Set<String> containerColumns) {
+    public static String buildGetIdentifiablesInContainerQuery(String tableName, Collection<String> columns, Set<String> containerColumns,
+                                                               int containerIdsSize) {
         StringBuilder sql = new StringBuilder()
                 .append("select ").append(ID_COLUMN).append(", ")
                 .append(String.join(", ", columns))
@@ -70,7 +75,12 @@ public static String buildGetIdentifiablesInContainerQuery(String tableName, Col
         var it = containerColumns.iterator();
         while (it.hasNext()) {
             String containerColumn = it.next();
-            sql.append(containerColumn).append(" = ?");
+            sql.append(containerColumn);
+            if (containerIdsSize == 1) {
+                sql.append(" = ?");
+            } else {
+                sql.append(" in (").append(String.join(", ", Collections.nCopies(containerIdsSize, "?"))).append(")");
+            }
             if (it.hasNext()) {
                 sql.append(" or ");
             }
@@ -215,4 +225,13 @@ public static String buildCloneNetworksQuery(Collection<String> columns) {
                 " from network" + " " +
                 "where uuid = ? and variantNum = ?";
     }
+
+    public static String buildGetVoltageLevelsInSubstationQuery(Collection<String> columns) {
+        return "select " + ID_COLUMN + ", " +
+                String.join(", ", columns) +
+                " from " + VOLTAGE_LEVEL_TABLE +
+                " where " + NETWORK_UUID_COLUMN + " = ?" +
+                " and " + VARIANT_NUM_COLUMN + " = ?" +
+                " and " + SUBSTATION_ID_COLUMN + " = ?";
+    }
 }