/
postgresql_database_tasks.rb
138 lines (115 loc) · 4.18 KB
/
postgresql_database_tasks.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
# frozen_string_literal: true
require "tempfile"
module ActiveRecord
module Tasks # :nodoc:
class PostgreSQLDatabaseTasks # :nodoc:
DEFAULT_ENCODING = ENV["CHARSET"] || "utf8"
ON_ERROR_STOP_1 = "ON_ERROR_STOP=1"
SQL_COMMENT_BEGIN = "--"
delegate :connection, :establish_connection, :clear_active_connections!,
to: ActiveRecord::Base
def self.using_database_configurations?
true
end
def initialize(db_config)
@db_config = db_config
@configuration_hash = db_config.configuration_hash
end
def create(master_established = false)
establish_master_connection unless master_established
connection.create_database(db_config.database, configuration_hash.merge(encoding: encoding))
establish_connection(db_config)
end
def drop
establish_master_connection
connection.drop_database(db_config.database)
end
def charset
connection.encoding
end
def collation
connection.collation
end
def purge
clear_active_connections!
drop
create true
end
def structure_dump(filename, extra_flags)
set_psql_env
search_path = \
case ActiveRecord::Base.dump_schemas
when :schema_search_path
configuration_hash[:schema_search_path]
when :all
nil
when String
ActiveRecord::Base.dump_schemas
end
args = ["--schema-only", "--no-privileges", "--no-owner", "--file", filename]
args.concat(Array(extra_flags)) if extra_flags
unless search_path.blank?
args += search_path.split(",").map do |part|
"--schema=#{part.strip}"
end
end
ignore_tables = ActiveRecord::SchemaDumper.ignore_tables
if ignore_tables.any?
args += ignore_tables.flat_map { |table| ["-T", table] }
end
args << db_config.database
run_cmd("pg_dump", args, "dumping")
remove_sql_header_comments(filename)
File.open(filename, "a") { |f| f << "SET search_path TO #{connection.schema_search_path};\n\n" }
end
def structure_load(filename, extra_flags)
set_psql_env
args = ["--set", ON_ERROR_STOP_1, "--quiet", "--no-psqlrc", "--file", filename]
args.concat(Array(extra_flags)) if extra_flags
args << db_config.database
run_cmd("psql", args, "loading")
end
private
attr_reader :db_config, :configuration_hash
def encoding
configuration_hash[:encoding] || DEFAULT_ENCODING
end
def establish_master_connection
establish_connection configuration_hash.merge(
database: "postgres",
schema_search_path: "public"
)
end
def set_psql_env
ENV["PGHOST"] = db_config.host if db_config.host
ENV["PGPORT"] = configuration_hash[:port].to_s if configuration_hash[:port]
ENV["PGPASSWORD"] = configuration_hash[:password].to_s if configuration_hash[:password]
ENV["PGUSER"] = configuration_hash[:username].to_s if configuration_hash[:username]
end
def run_cmd(cmd, args, action)
fail run_cmd_error(cmd, args, action) unless Kernel.system(cmd, *args)
end
def run_cmd_error(cmd, args, action)
msg = +"failed to execute:\n"
msg << "#{cmd} #{args.join(' ')}\n\n"
msg << "Please check the output above for any errors and make sure that `#{cmd}` is installed in your PATH and has proper permissions.\n\n"
msg
end
def remove_sql_header_comments(filename)
removing_comments = true
tempfile = Tempfile.open("uncommented_structure.sql")
begin
File.foreach(filename) do |line|
unless removing_comments && (line.start_with?(SQL_COMMENT_BEGIN) || line.blank?)
tempfile << line
removing_comments = false
end
end
ensure
tempfile.close
end
FileUtils.cp(tempfile.path, filename)
end
end
end
end