Fix issues and add extra option to quarantine_dynamodb cli

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    quarantine (2.2.2)
+    quarantine (2.3.0)
       rspec (~> 3.0)
       rspec-retry (~> 0.6)
       sorbet-runtime (= 0.5.6338)

bin/quarantine_dynamodb

@@ -1,6 +1,7 @@
 #!/usr/bin/env ruby
 
 require 'optparse'
+require 'sorbet-runtime'
 
 require_relative '../lib/quarantine/cli'
 

lib/quarantine/cli.rb

@@ -30,6 +30,10 @@ def parse
           @options[:region] = region
         end
 
+        parser.on('-e', '--endpoint [ENDPOINT]', String, 'Specify the aws endpoint for DynamoDB') do |endpoint|
+          @options[:endpoint] = endpoint
+        end
+
         parser.on(
           '-qTABLE',
           '--quarantine_table=TABLE',
@@ -55,7 +59,8 @@ def parse
     # TODO: eventually move to a separate file & create_table by db type when my db adapters
     sig { void }
     def create_tables
-      dynamodb = Quarantine::Databases::DynamoDB.new(region: @options[:region])
+      dynamodb_params = @options.slice(:region, :endpoint).compact
+      dynamodb = Quarantine::Databases::DynamoDB.new(dynamodb_params)
 
       attributes = [
         { attribute_name: 'id', attribute_type: 'S', key_type: 'HASH' }

lib/quarantine/databases/dynamo_db.rb

@@ -12,6 +12,8 @@ module Databases
     class DynamoDB < Base
       extend T::Sig
 
+      BATCH_WRITE_MAXIMUM_ITEMS = T.let(25, Integer)
+
       Attribute = T.type_alias { { attribute_name: String, attribute_type: String, key_type: String } }
 
       sig { returns(Aws::DynamoDB::Client) }
@@ -42,17 +44,19 @@ def fetch_items(table_name)
         ).void
       end
       def write_items(table_name, items)
-        @dynamodb.batch_write_item(
-          request_items: {
-            table_name => items.map do |item|
-              {
-                put_request: {
-                  item: item
+        items.each_slice(BATCH_WRITE_MAXIMUM_ITEMS) do |batch_items|
+          @dynamodb.batch_write_item(
+            request_items: {
+              table_name => batch_items.map do |item|
+                {
+                  put_request: {
+                    item: item
+                  }
                 }
-              }
-            end
-          }
-        )
+              end
+            }
+          )
+        end
       rescue Aws::DynamoDB::Errors::ServiceError
         raise Quarantine::DatabaseError
       end

lib/quarantine/version.rb

@@ -1,3 +1,3 @@
 class Quarantine
-  VERSION = '2.2.2'.freeze
+  VERSION = '2.3.0'.freeze
 end

spec/quarantine/cli_spec.rb

@@ -37,6 +37,18 @@
       expect(cli.options[:test_statuses_table_name]).to eq('foo')
     end
 
+    it 'define the endpoint' do
+      endpoint = 'http://localhost:4569'
+      cli = Quarantine::CLI.new
+
+      ARGV << '-r' << 'us-west-1'
+      ARGV << '-e' << endpoint
+
+      cli.parse
+
+      expect(cli.options[:endpoint]).to eq(endpoint)
+    end
+
     context '#create_tables' do
       let(:dynamodb) { Quarantine::Databases::DynamoDB.new(region: 'us-west-1') }
       let(:cli) { Quarantine::CLI.new }
@@ -63,6 +75,26 @@
         cli.create_tables
       end
 
+      it 'called with endpoint' do
+        endpoint = 'localhost:4569'
+        dynamodb_double = double('dynamodb instance')
+
+        allow(dynamodb_double).to receive(:create_table).with(any_args)
+        allow(Quarantine::Databases::DynamoDB)
+          .to receive(:new).and_return(dynamodb_double)
+
+        ARGV << '-r' << 'us-west-1'
+        ARGV << '-e' << endpoint
+
+        cli.parse
+        cli.create_tables
+
+        expect(dynamodb_double).to have_received(:create_table)
+        expect(Quarantine::Databases::DynamoDB)
+          .to have_received(:new)
+          .with(hash_including(endpoint: endpoint))
+      end
+
       it 'raises exception if Quarantine::DatabaseError exception occurs' do
         allow(Quarantine::Databases::DynamoDB).to receive(:new).and_return(dynamodb)
         allow(dynamodb).to receive(:create_table).and_raise(Quarantine::DatabaseError.new('db error'))