Merge remote-tracking branch 'origin/release-3.2.14'

CHANGELOG.md

@@ -1,3 +1,7 @@
+### 3.2.14
+
+* Allow to split data into multiple shapefiles
+
 ### 3.2.13
 
 * Add country, pame_country and marine stats (WDPA June 2020 release).

lib/modules/download/generators/shapefile.rb

@@ -13,42 +13,58 @@ class Download::Generators::Shapefile < Download::Generators::Base
     }
   }
 
-  def initialize zip_path, wdpa_ids
+  def initialize zip_path, wdpa_ids, number_of_pieces=3
     @path = File.dirname(zip_path)
     @filename = File.basename(zip_path, File.extname(zip_path))
     @wdpa_ids = wdpa_ids
+    @number_of_pieces = number_of_pieces
   end
 
   def generate
     return false if @wdpa_ids.is_a?(Array) && @wdpa_ids.empty?
 
     shapefile_paths = []
 
-    clean_up_after do
-      QUERY_CONDITIONS.each do |name, props|
-        shapefile_paths |= export_component name, props
-      end
+    @number_of_pieces.times do |i|
+      clean_up_after do
+        QUERY_CONDITIONS.each do |name, props|
+          shapefile_paths |= export_component name, props, i
+        end
 
-      system("zip -j #{zip_path} #{shapefile_paths.join(' ')}") and system("zip -ru #{zip_path} *", chdir: ATTACHMENTS_PATH)
+        return false unless system("zip -j #{zip_path(i)} #{shapefile_paths.join(' ')}")
+      end
     end
+    merge_files
   rescue Ogr::Postgres::ExportError
     return false
   end
 
   private
 
-  def export_component name, props
+  def export_component name, props, piece_index
     component_paths = shapefile_components(name)
     view_name = create_view query(props[:select], props[:where])
 
-    return [] if ActiveRecord::Base.connection.select_value("""
+    total_count = ActiveRecord::Base.connection.select_value("""
       SELECT COUNT(*) FROM #{view_name}
-    """).to_i.zero?
+    """).to_i
+
+    return [] if total_count.zero?
+
+    limit = (total_count / @number_of_pieces.to_f).ceil
+    offset = limit * piece_index
+    order_by = 'ORDER BY \""WDPAID"\" ASC'
+    sql = """
+      SELECT *
+      FROM #{view_name}
+      #{order_by if name.to_s == 'polygons'}
+      LIMIT #{limit} OFFSET #{offset}
+    """.squish
 
     export_success = Ogr::Postgres.export(
       :shapefile,
       component_paths.first,
-      "SELECT * FROM #{view_name}"
+      sql
     )
 
     raise Ogr::Postgres::ExportError unless export_success
@@ -74,13 +90,20 @@ def clean_up
     end
   end
 
-  def zip_path
-    File.join(@path, "#{@filename}.zip")
+  def zip_path(index='')
+    File.join(@path, "#{@filename}#{index}.zip")
   end
 
   def shapefile_components name
     SHAPEFILE_PARTS.collect do |ext|
       File.join(@path, "#{@filename}-#{name}.#{ext}")
     end
   end
+
+  def merge_files
+    range = (0..@number_of_pieces-1)
+    files_paths = range.map { |i| zip_path(i) }.join(' ')
+    system("zip -j #{zip_path} #{files_paths}") and system("zip -ru #{zip_path} *", chdir: ATTACHMENTS_PATH)
+    range.each { |i| FileUtils.rm_rf(zip_path(i)) }
+  end
 end

test/unit/download/generators/shapefile_test.rb

@@ -4,6 +4,9 @@ class DownloadShapefileTest < ActiveSupport::TestCase
   test '#generate, given a zip file path, exports shapefiles for each
    geometry table, and returns them as a single zip' do
     zip_file_path = './all-shp.zip'
+    zip_file_path0 = './all-shp0.zip'
+    zip_file_path1 = './all-shp1.zip'
+    zip_file_path2 = './all-shp2.zip'
     shp_polygon_file_path = './all-shp-polygons.shp'
     shp_polygon_joined_files = './all-shp-polygons.shp ./all-shp-polygons.shx ./all-shp-polygons.dbf ./all-shp-polygons.prj ./all-shp-polygons.cpg'
 
@@ -26,13 +29,25 @@ class DownloadShapefileTest < ActiveSupport::TestCase
     view_name_point = 'temporary_view_456'
     Download::Generators::Shapefile.any_instance.stubs(:create_view).with(shp_point_query).returns(view_name_point)
 
-    ActiveRecord::Base.connection.stubs(:select_value).returns(1).twice
-    Ogr::Postgres.expects(:export).with(:shapefile, shp_polygon_file_path, "SELECT * FROM #{view_name_poly}").returns(true)
-    Ogr::Postgres.expects(:export).with(:shapefile, shp_point_file_path, "SELECT * FROM #{view_name_point}").returns(true)
-
-    create_zip_command = "zip -j #{zip_file_path} #{shp_polygon_joined_files} #{shp_point_joined_files}"
+    ActiveRecord::Base.connection.stubs(:select_value).returns(3)
+    poly_query = "SELECT * FROM #{view_name_poly}" << ' ORDER BY \""WDPAID"\" ASC'
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_polygon_file_path, "#{poly_query} LIMIT 1 OFFSET 0").returns(true)
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_polygon_file_path, "#{poly_query} LIMIT 1 OFFSET 1").returns(true)
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_polygon_file_path, "#{poly_query} LIMIT 1 OFFSET 2").returns(true)
+    point_query = "SELECT * FROM #{view_name_point}"
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_point_file_path, "#{point_query} LIMIT 1 OFFSET 0").returns(true)
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_point_file_path, "#{point_query} LIMIT 1 OFFSET 1").returns(true)
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_point_file_path, "#{point_query} LIMIT 1 OFFSET 2").returns(true)
+
+    create_zip_command = "zip -j #{zip_file_path0} #{shp_polygon_joined_files} #{shp_point_joined_files}"
+    Download::Generators::Shapefile.any_instance.expects(:system).with(create_zip_command).returns(true)
+    create_zip_command = "zip -j #{zip_file_path1} #{shp_polygon_joined_files} #{shp_point_joined_files}"
+    Download::Generators::Shapefile.any_instance.expects(:system).with(create_zip_command).returns(true)
+    create_zip_command = "zip -j #{zip_file_path2} #{shp_polygon_joined_files} #{shp_point_joined_files}"
     Download::Generators::Shapefile.any_instance.expects(:system).with(create_zip_command).returns(true)
 
+    merge_zip_command = "zip -j #{zip_file_path} #{zip_file_path0} #{zip_file_path1} #{zip_file_path2}"
+    Download::Generators::Shapefile.any_instance.expects(:system).with(merge_zip_command).returns(true)
     update_zip_command = "zip -ru #{zip_file_path} *"
     opts = {chdir: Download::Generators::Base::ATTACHMENTS_PATH}
     Download::Generators::Shapefile.any_instance.expects(:system).with(update_zip_command, opts).returns(true)
@@ -52,7 +67,7 @@ class DownloadShapefileTest < ActiveSupport::TestCase
 
   test '#generate returns false if the zip fails' do
     ActiveRecord::Base.connection.stubs(:execute)
-    ActiveRecord::Base.connection.stubs(:select_value).returns(1).twice
+    ActiveRecord::Base.connection.stubs(:select_value).returns(2)
     Ogr::Postgres.expects(:export).twice.returns(true)
     Download::Generators::Shapefile.any_instance.expects(:system).returns(false)
 
@@ -77,24 +92,35 @@ class DownloadShapefileTest < ActiveSupport::TestCase
       './all-points.cpg'
     ]
 
+    zip_file_path  = './all.zip'
+    zip_file_path0 = './all0.zip'
+    zip_file_path1 = './all1.zip'
+    zip_file_path2 = './all2.zip'
+
     ActiveRecord::Base.connection.stubs(:execute)
-    ActiveRecord::Base.connection.stubs(:select_value).returns(1).twice
-    Ogr::Postgres.expects(:export).twice.returns(true)
+    ActiveRecord::Base.connection.stubs(:select_value).returns(3).times(6)
+    Ogr::Postgres.expects(:export).times(6).returns(true)
     Download::Generators::Shapefile.
       any_instance.
       expects(:system).
       returns(true).
-      twice
+      times(5) # number_of_pieces + 2 related to merge_files method
 
-    FileUtils.expects(:rm_rf).with(shp_polygons_paths)
-    FileUtils.expects(:rm_rf).with(shp_points_paths)
+    FileUtils.expects(:rm_rf).with(shp_polygons_paths).times(3)
+    FileUtils.expects(:rm_rf).with(shp_points_paths).times(3)
+    FileUtils.expects(:rm_rf).with(zip_file_path0)
+    FileUtils.expects(:rm_rf).with(zip_file_path1)
+    FileUtils.expects(:rm_rf).with(zip_file_path2)
 
-    Download::Generators::Shapefile.generate('./all.zip')
+    Download::Generators::Shapefile.generate(zip_file_path)
   end
 
   test '#generate, given a zip file path and WDPA IDs, exports
    shapefiles for each geometry table, and returns them as a single zip' do
-    zip_file_path = './all-shp.zip'
+    zip_file_path  = './all-shp.zip'
+    zip_file_path0 = './all-shp0.zip'
+    zip_file_path1 = './all-shp1.zip'
+    zip_file_path2 = './all-shp2.zip'
     shp_polygon_file_path = './all-shp-polygons.shp'
     shp_polygon_joined_files = './all-shp-polygons.shp ./all-shp-polygons.shx ./all-shp-polygons.dbf ./all-shp-polygons.prj ./all-shp-polygons.cpg'
 
@@ -121,13 +147,25 @@ class DownloadShapefileTest < ActiveSupport::TestCase
     view_name_point = 'temporary_view_456'
     Download::Generators::Shapefile.any_instance.stubs(:create_view).with(shp_point_query).returns(view_name_point)
 
-    ActiveRecord::Base.connection.stubs(:select_value).returns(1).twice
-    Ogr::Postgres.expects(:export).with(:shapefile, shp_polygon_file_path, "SELECT * FROM #{view_name_poly}").returns(true)
-    Ogr::Postgres.expects(:export).with(:shapefile, shp_point_file_path, "SELECT * FROM #{view_name_point}").returns(true)
-
-    create_zip_command = "zip -j #{zip_file_path} #{shp_polygon_joined_files} #{shp_point_joined_files}"
+    ActiveRecord::Base.connection.stubs(:select_value).returns(1).times(6)
+    poly_query = "SELECT * FROM #{view_name_poly}" << ' ORDER BY \""WDPAID"\" ASC'
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_polygon_file_path, "#{poly_query} LIMIT 1 OFFSET 0").returns(true)
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_polygon_file_path, "#{poly_query} LIMIT 1 OFFSET 1").returns(true)
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_polygon_file_path, "#{poly_query} LIMIT 1 OFFSET 2").returns(true)
+    point_query = "SELECT * FROM #{view_name_point}"
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_point_file_path, "#{point_query} LIMIT 1 OFFSET 0").returns(true)
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_point_file_path, "#{point_query} LIMIT 1 OFFSET 1").returns(true)
+    Ogr::Postgres.expects(:export).with(:shapefile, shp_point_file_path, "#{point_query} LIMIT 1 OFFSET 2").returns(true)
+
+    create_zip_command = "zip -j #{zip_file_path0} #{shp_polygon_joined_files} #{shp_point_joined_files}"
+    Download::Generators::Shapefile.any_instance.expects(:system).with(create_zip_command).returns(true)
+    create_zip_command = "zip -j #{zip_file_path1} #{shp_polygon_joined_files} #{shp_point_joined_files}"
+    Download::Generators::Shapefile.any_instance.expects(:system).with(create_zip_command).returns(true)
+    create_zip_command = "zip -j #{zip_file_path2} #{shp_polygon_joined_files} #{shp_point_joined_files}"
     Download::Generators::Shapefile.any_instance.expects(:system).with(create_zip_command).returns(true)
 
+    merge_zip_command = "zip -j #{zip_file_path} #{zip_file_path0} #{zip_file_path1} #{zip_file_path2}"
+    Download::Generators::Shapefile.any_instance.expects(:system).with(merge_zip_command).returns(true)
     update_zip_command = "zip -ru #{zip_file_path} *"
     opts = {chdir: Download::Generators::Base::ATTACHMENTS_PATH}
     Download::Generators::Shapefile.any_instance.expects(:system).with(update_zip_command, opts).returns(true)