pgslice.rb in pgslice-0.3.6

- old
+ new

@@ -110,11 +110,11 @@
       abort "Usage: pgslice unprep <table>" if arguments.length != 1
       abort "Table not found: #{intermediate_table}" unless table_exists?(intermediate_table)
 
       queries = [
         "DROP TABLE #{quote_ident(intermediate_table)} CASCADE;",
-        "DROP FUNCTION IF EXISTS #{trigger_name}();"
+        "DROP FUNCTION IF EXISTS #{quote_ident(trigger_name)}();"
       ]
       run_queries(queries)
     end
 
     def add_partitions
@@ -130,11 +130,11 @@
       range = (-1 * past)..future
 
       # ensure table has trigger
       abort "No trigger on table: #{table}\nDid you mean to use --intermediate?" unless has_trigger?(trigger_name, table)
 
-      index_defs = execute("SELECT pg_get_indexdef(indexrelid) FROM pg_index INNER JOIN pg_stat_user_indexes USING (indexrelid) WHERE relname = $1 AND schemaname = $2 AND indisprimary = 'f'", [original_table, schema]).map { |r| r["pg_get_indexdef"] }
+      index_defs = execute("SELECT pg_get_indexdef(indexrelid) FROM pg_index WHERE indrelid = #{regclass(schema, original_table)} AND indisprimary = 'f'").map { |r| r["pg_get_indexdef"] }
       primary_key = self.primary_key(table)
 
       queries = []
 
       period, field, cast, needs_comment = settings_from_trigger(original_table, table)
@@ -161,20 +161,20 @@
         SQL
 
         queries << "ALTER TABLE #{quote_ident(partition_name)} ADD PRIMARY KEY (#{quote_ident(primary_key)});" if primary_key
 
         index_defs.each do |index_def|
-          queries << index_def.sub(" ON #{original_table} USING ", " ON #{partition_name} USING ").sub(/ INDEX .+ ON /, " INDEX ON ") + ";"
+          queries << index_def.sub(/ ON \S+ USING /, " ON #{quote_ident(partition_name)} USING ").sub(/ INDEX .+ ON /, " INDEX ON ") + ";"
         end
       end
 
       # update trigger based on existing partitions
       current_defs = []
       future_defs = []
       past_defs = []
       name_format = self.name_format(period)
-      existing_tables = self.existing_tables(like: "#{original_table}_%").select { |t| /\A#{Regexp.escape("#{original_table}_")}\d{6,8}\z/.match(t) }
+      existing_tables = existing_partitions(original_table)
       existing_tables = (existing_tables + added_partitions).uniq.sort
 
       existing_tables.each do |table|
         day = DateTime.strptime(table.split("_").last, name_format)
         partition_name = "#{original_table}_#{day.strftime(name_format(period))}"
@@ -234,11 +234,11 @@
       period, field, cast, needs_comment = settings_from_trigger(table, dest_table)
 
       if period
         name_format = self.name_format(period)
 
-        existing_tables = self.existing_tables(like: "#{table}_%").select { |t| /\A#{Regexp.escape("#{table}_")}\d{6,8}\z/.match(t) }.sort
+        existing_tables = existing_partitions(table)
         if existing_tables.any?
           starting_time = DateTime.strptime(existing_tables.first.split("_").last, name_format)
           ending_time = advance_date(DateTime.strptime(existing_tables.last.split("_").last, name_format), period, 1)
         end
       end
@@ -313,11 +313,11 @@
         "ALTER TABLE #{quote_ident(table)} RENAME TO #{quote_ident(retired_table)};",
         "ALTER TABLE #{quote_ident(intermediate_table)} RENAME TO #{quote_ident(table)};"
       ]
 
       self.sequences(table).each do |sequence|
-        queries << "ALTER SEQUENCE #{quote_ident(sequence["sequence_name"])} OWNED BY #{table}.#{sequence["related_column"]};"
+        queries << "ALTER SEQUENCE #{quote_ident(sequence["sequence_name"])} OWNED BY #{quote_ident(table)}.#{quote_ident(sequence["related_column"])};"
       end
 
       queries.unshift("SET LOCAL lock_timeout = '#{options[:lock_timeout]}';") if server_version_num >= 90300
 
       run_queries(queries)
@@ -337,11 +337,11 @@
         "ALTER TABLE #{quote_ident(table)} RENAME TO #{quote_ident(intermediate_table)};",
         "ALTER TABLE #{quote_ident(retired_table)} RENAME TO #{quote_ident(table)};"
       ]
 
       self.sequences(table).each do |sequence|
-        queries << "ALTER SEQUENCE #{quote_ident(sequence["sequence_name"])} OWNED BY #{table}.#{sequence["related_column"]};"
+        queries << "ALTER SEQUENCE #{quote_ident(sequence["sequence_name"])} OWNED BY #{quote_ident(table)}.#{quote_ident(sequence["related_column"])};"
       end
 
       run_queries(queries)
     end
 
@@ -349,11 +349,11 @@
       table = arguments.first
       parent_table = options[:swapped] ? table : intermediate_name(table)
 
       abort "Usage: pgslice analyze <table>" if arguments.length != 1
 
-      existing_tables = self.existing_tables(like: "#{table}_%").select { |t| /\A#{Regexp.escape("#{table}_")}\d{6,8}\z/.match(t) }
+      existing_tables = existing_partitions(table)
       analyze_list = existing_tables + [parent_table]
       run_queries_without_transaction analyze_list.map { |t| "ANALYZE VERBOSE #{quote_ident(t)};" }
     end
 
     # arguments
@@ -462,10 +462,14 @@
 
     def server_version_num
       execute("SHOW server_version_num")[0]["server_version_num"].to_i
     end
 
+    def existing_partitions(table)
+      existing_tables(like: "#{table}_%").select { |t| /\A#{Regexp.escape("#{table}_")}\d{6,8}\z/.match(t) }
+    end
+
     def existing_tables(like:)
       query = "SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname = $1 AND tablename LIKE $2"
       execute(query, [schema, like]).map { |r| r["tablename"] }.sort
     end
 
@@ -597,11 +601,15 @@
 
     def quote_ident(value)
       PG::Connection.quote_ident(value)
     end
 
+    def regclass(schema, table)
+      "'#{quote_ident(schema)}.#{quote_ident(table)}'::regclass"
+    end
+
     def fetch_trigger(trigger_name, table)
-      execute("SELECT obj_description(oid, 'pg_trigger') AS comment FROM pg_trigger WHERE tgname = $1 AND EXISTS (SELECT 1 FROM pg_stat_user_tables WHERE relid = tgrelid AND relname = $2 AND schemaname = $3)", [trigger_name, table, schema])[0]
+      execute("SELECT obj_description(oid, 'pg_trigger') AS comment FROM pg_trigger WHERE tgname = $1 AND tgrelid = #{regclass(schema, table)}", [trigger_name])[0]
     end
 
     def settings_from_trigger(original_table, table)
       trigger_name = self.trigger_name(original_table)
 