Ogr2OgrOutputFormat.java

/* Copyright (c) 2001 - 2007 TOPP - www.openplans.org. All rights reserved.
 * This code is licensed under the GPL 2.0 license, available at the root
 * application directory.
 */
package org.geoserver.wfs.response;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.logging.Level;
import java.util.zip.ZipOutputStream;

import javax.xml.namespace.QName;

import net.opengis.wfs.FeatureCollectionType;
import net.opengis.wfs.GetFeatureType;
import net.opengis.wfs.QueryType;

import org.geoserver.data.util.IOUtils;
import org.geoserver.ows.util.OwsUtils;
import org.geoserver.platform.Operation;
import org.geoserver.platform.ServiceException;
import org.geoserver.wfs.WFSException;
import org.geoserver.wfs.WFSGetFeatureOutputFormat;
import org.geotools.data.DataStore;
import org.geotools.data.FeatureStore;
import org.geotools.data.shapefile.ShapefileDataStore;
import org.geotools.feature.FeatureCollection;
import org.geotools.gml.producer.FeatureTransformer;
import org.geotools.gml2.bindings.GML2EncodingUtils;
import org.opengis.feature.GeometryAttribute;
import org.opengis.feature.simple.SimpleFeature;
import org.opengis.feature.simple.SimpleFeatureType;
import org.opengis.feature.type.AttributeDescriptor;
import org.opengis.feature.type.GeometryDescriptor;
import org.opengis.feature.type.GeometryType;
import org.opengis.referencing.crs.CoordinateReferenceSystem;

import com.vividsolutions.jts.geom.LineString;
import com.vividsolutions.jts.geom.LinearRing;
import com.vividsolutions.jts.geom.MultiLineString;
import com.vividsolutions.jts.geom.MultiPoint;
import com.vividsolutions.jts.geom.MultiPolygon;
import com.vividsolutions.jts.geom.Point;
import com.vividsolutions.jts.geom.Polygon;

public class Ogr2OgrOutputFormat extends WFSGetFeatureOutputFormat {
    
    /**
     * The types of geometries a shapefile can handle
     */
    private static final Set<Class> SHAPEFILE_GEOM_TYPES = new HashSet<Class>() {
        {
            add(Point.class);
            add(LineString.class);
            add(LinearRing.class);
            add(Polygon.class);
            add(MultiPoint.class);
            add(MultiLineString.class);
            add(MultiPolygon.class);
        }
    };
    
    /**
     * The fs path to ogr2ogr. If null, we'll assume ogr2ogr is in the PATH and
     * that we can execute it just by running ogr2ogr
     */
    String ogrPath = null;

    /**
     * The full path to ogr2ogr
     */
    String ogrExecutable = "ogr2ogr";
    
    /**
     * The GDAL_DATA folder
     */
    String gdalData = null;

    /**
     * The output formats we can generate using ogr2ogr. Using a concurrent
     * one so that it can be reconfigured while the output format is working
     */
    static Map<String, OgrFormat> formats = new ConcurrentHashMap<String, OgrFormat>();

    public Ogr2OgrOutputFormat() {
        // initialize with the key set of formats, so that it will change as
        // we register new formats
        super(formats.keySet());
    }

    /**
     * Returns the ogr2ogr executable full path
     * 
     * @return
     */
    public String getOgrExecutable() {
        return ogrExecutable;
    }

    /**
     * Sets the ogr2ogr executable full path. The default value is simply
     * "ogr2ogr", which will work if ogr2ogr is in the path
     * 
     * @param ogrExecutable
     */
    public void setOgrExecutable(String ogrExecutable) {
        this.ogrExecutable = ogrExecutable;
    }
    
    /**
     * Returns the location of the gdal data folder (required to set the output srs)
     * @return
     */
    public String getGdalData() {
        return gdalData;
    }

    /**
     * Sets the location of the gdal data folder (requierd to set the output srs)
     * @param gdalData
     */
    public void setGdalData(String gdalData) {
        this.gdalData = gdalData;
    }

    /**
     * @see WFSGetFeatureOutputFormat#getMimeType(Object, Operation)
     */
    public String getMimeType(Object value, Operation operation) throws ServiceException {
        return "application/zip";
    }
    
    @Override
    public String[][] getHeaders(Object value, Operation operation) throws ServiceException {
        GetFeatureType request = (GetFeatureType) OwsUtils.parameter(operation.getParameters(),
                GetFeatureType.class);
        String outputFileName = ((QName) ((QueryType) request.getQuery().get(0)).getTypeName().get(0))
            .getLocalPart();
        return (String[][]) new String[][] {
                { "Content-Disposition", "attachment; filename=" + outputFileName + ".zip" }
            };
    }

    /**
     * Adds a ogr format among the supported ones
     * 
     * @param parameters
     */
    public void addFormat(OgrFormat parameters) {
        formats.put(parameters.formatName, parameters);
    }

    /**
     * Programmatically removes all formats
     * 
     * @param parameters
     */
    public void clearFormats() {
        formats.clear();
    }

    /**
     * Writes out the data to an OGR known format (GML/shapefile) to disk and
     * then ogr2ogr each generated file into the destination format. Finally,
     * zips up all the resulting files.
     */
    @Override
    protected void write(FeatureCollectionType featureCollection, OutputStream output,
            Operation getFeature) throws IOException, ServiceException {

        // figure out which output format we're going to generate
        GetFeatureType gft = (GetFeatureType) getFeature.getParameters()[0];
        OgrFormat format = formats.get(gft.getOutputFormat());
        if (format == null)
            throw new WFSException("Unknown output format " + gft.getOutputFormat());

        // create the first temp directory, used for dumping gs generated
        // content
        File tempGS = org.geoserver.data.util.IOUtils.createTempDirectory("ogrtmpin");
        File tempOGR = org.geoserver.data.util.IOUtils.createTempDirectory("ogrtmpout");

        // build the ogr wrapper used to run the ogr2ogr commands
        OGRWrapper wrapper = new OGRWrapper(ogrExecutable, gdalData);

        // actually export each feature collection
        try {
            Iterator outputFeatureCollections = featureCollection.getFeature().iterator();
            FeatureCollection<SimpleFeatureType, SimpleFeature> curCollection;

            while (outputFeatureCollections.hasNext()) {
                curCollection = (FeatureCollection<SimpleFeatureType, SimpleFeature>) outputFeatureCollections
                        .next();

                // write out the gml
                File intermediate = writeToDisk(tempGS, curCollection);

                // convert with ogr2ogr
                String epsgCode = null;
                final SimpleFeatureType schema = curCollection.getSchema();
                final CoordinateReferenceSystem crs = schema.getCoordinateReferenceSystem();
                wrapper.convert(intermediate, tempOGR, schema.getTypeName(), format, crs);

                // wipe out the input dir contents
                IOUtils.emptyDirectory(tempGS);
            }

            // scan the output directory and zip it all
            ZipOutputStream zipOut = new ZipOutputStream(output);
            IOUtils.zipDirectory(tempOGR, zipOut, null);

            // delete the input and output directories
            IOUtils.delete(tempGS);
            IOUtils.delete(tempOGR);
        } catch (Exception e) {
            throw new ServiceException("Exception occurred during output generation", e);
        }
    }
    
    /**
     * Writes to disk using shapefile if the feature type allows for it, GML otherwise
     * @param tempDir
     * @param curCollection
     * @return
     */
    private File writeToDisk(File tempDir,
            FeatureCollection<SimpleFeatureType, SimpleFeature> curCollection) throws Exception {
        if(isShapefileCompatible(curCollection.getSchema())) 
            return writeShapefile(tempDir, curCollection);
        else
            return writeGML(tempDir, curCollection);
    }

    private File writeShapefile(File tempDir,
            FeatureCollection<SimpleFeatureType, SimpleFeature> collection) {
        SimpleFeatureType schema = collection.getSchema();

        FeatureStore<SimpleFeatureType, SimpleFeature> fstore = null;
        DataStore dstore = null;
        File file = null;
        try {
            file = new File(tempDir, schema.getTypeName() + ".shp");
            dstore = new ShapefileDataStore(file.toURL());
            dstore.createSchema(schema);
            
            fstore = (FeatureStore<SimpleFeatureType, SimpleFeature>) dstore.getFeatureSource(schema.getTypeName());
            fstore.addFeatures(collection);
        } catch (IOException ioe) {
            LOGGER.log(Level.WARNING,
                "Error while writing featuretype '" + schema.getTypeName() + "' to shapefile.", ioe);
            throw new ServiceException(ioe);
        } finally {
            if(dstore != null) {
                dstore.dispose();
            }
        }
        
        return file; 
    }

    /**
     * Returns true if the schema has just one geometry and the geom type is known
     * @param schema
     * @return
     */
    private boolean isShapefileCompatible(SimpleFeatureType schema) {
        GeometryType gt = null;
        for (AttributeDescriptor at : schema.getAttributeDescriptors()) {
            if(at instanceof GeometryDescriptor) {
                if(gt == null)
                    gt = ((GeometryDescriptor) at).getType();
                else
                    // more than one geometry 
                    return false;
            }
        } 
        
        return gt != null && SHAPEFILE_GEOM_TYPES.contains(gt.getBinding()); 
    }

    private File writeGML(File tempDir,
            FeatureCollection<SimpleFeatureType, SimpleFeature> curCollection) throws Exception {
        // create the temp file for this output
        File outFile = new File(tempDir, curCollection.getSchema().getTypeName() + ".gml");

        // write out
        OutputStream os = null;
        try {
            os = new FileOutputStream(outFile);

            // let's invoke the transformer
            FeatureTransformer ft = new FeatureTransformer();
            ft.setNumDecimals(16);
            ft.setNamespaceDeclarationEnabled(false);
            ft.getFeatureNamespaces().declarePrefix("topp",
                    curCollection.getSchema().getName().getNamespaceURI());
            ft.transform(curCollection, os);
        } finally {
            os.close();
        }

        return outFile;
    }

}