Merge pull request #115 from nymta/transformer-extensions

Fix #113 - enable implementations of GtfsTransformStrategy to use GTFS extensions
OneBusAway · Jul 14, 2020 · cd9a0bf · cd9a0bf
2 parents 8cb24f9 + 9f3d2d2
commit cd9a0bf
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 31 deletions.
diff --git a/...saway-gtfs-transformer/src/main/java/org/onebusaway/gtfs_transformer/GtfsTransformer.java b/...saway-gtfs-transformer/src/main/java/org/onebusaway/gtfs_transformer/GtfsTransformer.java
@@ -39,7 +39,6 @@
 import org.onebusaway.gtfs_transformer.factory.TransformFactory;
 import org.onebusaway.gtfs_transformer.services.GtfsEntityTransformStrategy;
 import org.onebusaway.gtfs_transformer.services.GtfsTransformStrategy;
-import org.onebusaway.gtfs_transformer.services.SchemaUpdateStrategy;
 import org.onebusaway.gtfs_transformer.services.TransformContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -62,8 +61,6 @@ public class GtfsTransformer {
 
   private List<GtfsEntityTransformStrategy> _entityTransformStrategies = new ArrayList<GtfsEntityTransformStrategy>();
 
-  private List<SchemaUpdateStrategy> _outputSchemaUpdates = new ArrayList<SchemaUpdateStrategy>();
-
   private TransformContext _context = new TransformContext();
 
   private GtfsReader _reader = new GtfsReader();
@@ -114,10 +111,6 @@ public void addEntityTransform(GtfsEntityTransformStrategy entityTransform) {
     _entityTransformStrategies.add(entityTransform);
   }
 
-  public void addOutputSchemaUpdate(SchemaUpdateStrategy outputSchemaUpdate) {
-    _outputSchemaUpdates.add(outputSchemaUpdate);
-  }
-
   public void addParameter(String key, Object value) {
     _parameters.put(key, value);
   }
@@ -186,6 +179,14 @@ private void readGtfs() throws IOException {
     if (!_entityTransformStrategies.isEmpty())
       dao = new DaoInterceptor(_dao);
 
+
+    DefaultEntitySchemaFactory schemaFactory = new DefaultEntitySchemaFactory();
+    schemaFactory.addFactory(GtfsEntitySchemaFactory.createEntitySchemaFactory());
+
+    _transformStrategies.forEach(s -> s.updateReadSchema(schemaFactory));
+
+    _reader.setEntitySchemaFactory(schemaFactory);
+
     _reader.setEntityStore(dao);
 
     if (_agencyId != null)
@@ -248,8 +249,7 @@ private void writeGtfs() throws IOException {
     DefaultEntitySchemaFactory schemaFactory = new DefaultEntitySchemaFactory();
     schemaFactory.addFactory(GtfsEntitySchemaFactory.createEntitySchemaFactory());
 
-    for (SchemaUpdateStrategy strategy : _outputSchemaUpdates)
-      strategy.updateSchema(schemaFactory);
+    _transformStrategies.forEach(s -> s.updateWriteSchema(schemaFactory));
 
     _writer.setEntitySchemaFactory(schemaFactory);
 

diff --git a/...sformer/src/main/java/org/onebusaway/gtfs_transformer/services/GtfsTransformStrategy.java b/...sformer/src/main/java/org/onebusaway/gtfs_transformer/services/GtfsTransformStrategy.java
@@ -15,10 +15,13 @@
  */
 package org.onebusaway.gtfs_transformer.services;
 
+import org.onebusaway.csv_entities.schema.DefaultEntitySchemaFactory;
 import org.onebusaway.gtfs.services.GtfsMutableRelationalDao;
 
 public interface GtfsTransformStrategy {
 
   public String getName();
   public void run(TransformContext context, GtfsMutableRelationalDao dao);
+  public default void updateReadSchema(DefaultEntitySchemaFactory factory) {};
+  public default void updateWriteSchema(DefaultEntitySchemaFactory factory) {};
 }
diff --git a/...nsformer/src/main/java/org/onebusaway/gtfs_transformer/services/SchemaUpdateStrategy.java b/...nsformer/src/main/java/org/onebusaway/gtfs_transformer/services/SchemaUpdateStrategy.java