lib/pgslice.rb in pgslice-0.2.1 vs lib/pgslice.rb in pgslice-0.2.2

@@ -247,11 +247,11 @@
 
       if max_dest_id == 0 && !options[:swapped]
         if options[:start]
           max_dest_id = options[:start]
         else
-          min_source_id = min_id(source_table, primary_key, field, starting_time)
+          min_source_id = min_id(source_table, primary_key, field, starting_time, options[:where])
           max_dest_id = min_source_id - 1 if min_source_id
         end
       end
 
       starting_id = max_dest_id
@@ -306,10 +306,12 @@
 
       self.sequences(table).each do |sequence|
         queries << "ALTER SEQUENCE #{sequence["sequence_name"]} OWNED BY #{table}.#{sequence["related_column"]};"
       end
 
+      queries.unshift("SET LOCAL lock_timeout = '#{options[:lock_timeout]}';") if server_version_num >= 90300
+
       run_queries(queries)
     end
 
     def unswap
       table = arguments.first
@@ -347,10 +349,11 @@
         o.boolean "--no-partition", default: false
         o.integer "--start"
         o.string "--url"
         o.string "--source-table"
         o.string "--where"
+        o.string "--lock-timeout", default: "5s"
         o.on "-v", "--version", "print the version" do
           log PgSlice::VERSION
           @exit = true
         end
       end
@@ -399,22 +402,32 @@
       connection.exec_params(query, params).to_a
     end
 
     def run_queries(queries)
       connection.transaction do
-        execute("SET client_min_messages TO warning") unless options[:dry_run]
+        execute("SET LOCAL client_min_messages TO warning") unless options[:dry_run]
         log_sql "BEGIN;"
         log_sql
         queries.each do |query|
           log_sql query
           log_sql
-          execute(query) unless options[:dry_run]
+          unless options[:dry_run]
+            begin
+              execute(query)
+            rescue PG::ServerError => e
+              abort("#{e.class.name}: #{e.message}")
+            end
+          end
         end
         log_sql "COMMIT;"
       end
     end
 
+    def server_version_num
+      execute("SHOW server_version_num")[0]["server_version_num"].to_i
+    end
+
     def existing_tables(like:)
       query = "SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname = $1 AND tablename LIKE $2"
       execute(query, ["public", like]).map { |r| r["tablename"] }.sort
     end
 
@@ -451,12 +464,15 @@
       query = "SELECT MAX(#{primary_key}) FROM #{table}"
       query << " WHERE #{primary_key} <= #{below}" if below
       execute(query)[0]["max"].to_i
     end
 
-    def min_id(table, primary_key, column, starting_time)
+    def min_id(table, primary_key, column, starting_time, where)
       query = "SELECT MIN(#{primary_key}) FROM #{table}"
-      query << " WHERE #{column} >= #{sql_date(starting_time)}" if starting_time
+      conditions = []
+      conditions << "#{column} >= #{sql_date(starting_time)}" if starting_time
+      conditions << where if where
+      query << " WHERE #{conditions.join(" AND ")}" if conditions.any?
       (execute(query)[0]["min"] || 1).to_i
     end
 
     def has_trigger?(trigger_name, table)
       execute("SELECT 1 FROM pg_trigger WHERE tgname = $1 AND tgrelid = $2::regclass", [trigger_name, table]).any?