/
restore_pg.py
150 lines (127 loc) · 7.23 KB
/
restore_pg.py
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
139
140
141
142
143
144
145
146
147
148
149
150
#!/usr/bin/python
# ****************************************************************#
# ScriptName: restore_from_downloads
# Author: @alibaba-inc.com
# Function:
# ***************************************************************#
import os
import sys
COLUMNS_STRUCTURE_FILE_NAME = "column_structure.sql"
CONSTRAINT_STRUCTURE_FILE_NAME = "constraint_structure.sql"
INDEX_STRUCTURE_FILE_NAME = "index_structure.sql"
INHERIT_STRUCTURE_FILE_NAME = "inherit_structure.sql"
STRUCTURE_FILE_NAME = "structure.sql"
DATA_PATH_PREFIX = "data"
def create_database(db_host, db_port, db_user, db_pass, db_name, create_stmt_file):
if db_name.lower() == "postgres": # 将 db_name 转换为小写,以便检查是否为 "postgres"
return
cmd = 'PGPASSWORD="' + db_pass + '" psql -h ' + db_host + ' -p ' + db_port + ' -U ' + db_user + ' -d postgres -c "CREATE DATABASE \\"' + db_name + '\\";"'
if os.system(cmd) != 0:
print("[ERROR]: execute SQL failed. command: " + cmd)
exit(1)
#
# def create_database(db_host, db_port, db_user, db_pass, db_name, create_stmt_file):
# if db_name == "postgres":
# return
# cmd = "PGPASSWORD=" + db_pass + " psql -h" + db_host + " -p" + db_port + " -U" + db_user + " -d" + "postgres" + " <" + create_stmt_file
# if os.system(cmd) != 0:
# print("[ERROR]: execute SQL failed. command: " + cmd)
# exit(1)
def create_schema(db_host, db_port, db_user, db_pass, db_name, create_stmt_file):
cmd = "PGPASSWORD=" + db_pass + " psql -h" + db_host + " -p" + db_port + " -U" + db_user + " -d" + db_name + " <" + create_stmt_file
if os.system(cmd) != 0:
print("[ERROR]: execute SQL failed. command: " + cmd)
exit(1)
def import_file_csv(db_host, db_port, db_user, db_pass, csv_file, db_name, db_schema, table):
load_cmd = " -c \"\copy " + "\"" + db_schema + "\".\"" + table + "\" " + " from '" + csv_file + "' delimiter ',' csv\""
cmd = "PGPASSWORD=" + db_pass + " psql -h" + db_host + " -p" + db_port + " -U" + db_user + " -d" + db_name + load_cmd
print("[INFO]: trying to exec: " + cmd)
if os.system(cmd) != 0:
print("[ERROR]: execute SQL failed. command: " + cmd)
exit(1)
def import_file_sql(db_host, db_port, db_user, db_pass, db_name, sql_file):
cmd = "PGPASSWORD=" + db_pass + " psql -h" + db_host + " -p" + db_port + " -U" + db_user + " -d" + db_name + " <" + sql_file
print("[INFO]: trying to exec: " + cmd)
if os.system(cmd) != 0:
print("[ERROR]: execute SQL failed. command: " + cmd)
exit(1)
def print_usage():
print(
"Usage: python ./restore_mysql.py [backupset_directory] [database_host] [database_port] [database_username] [database_password]")
if __name__ == '__main__':
if len(sys.argv) != 6:
print_usage()
exit()
root_dir = os.path.abspath(sys.argv[1])
db_host = sys.argv[2]
db_port = sys.argv[3]
db_user = sys.argv[4]
db_pass = sys.argv[5]
print("[INFO]: restore data from " + root_dir + " to " + db_host + ":" + db_port)
db_dirs = os.listdir(root_dir)
for db_dir in db_dirs:
db_dir_path = os.path.join(root_dir, db_dir)
if not os.path.isdir(db_dir_path):
continue
db_structure_file = os.path.join(db_dir_path, STRUCTURE_FILE_NAME)
print("[INFO]: restore create database: " + db_dir + " begin")
create_database(db_host, db_port, db_user, db_pass, db_dir, db_structure_file)
print("[INFO]: restore structure database: " + db_dir + " ends")
constraint_ddl_file_list = []
index_ddl_file_list = []
inherit_ddl_file_list = []
schema_dirs = os.listdir(db_dir_path)
for schema_dir in schema_dirs:
schema_dir_path = os.path.join(db_dir_path, schema_dir)
if not os.path.isdir(schema_dir_path):
continue
db_structure_file = os.path.join(schema_dir_path, STRUCTURE_FILE_NAME)
create_schema(db_host, db_port, db_user, db_pass, db_dir, db_structure_file)
print("[INFO]: restore structure schema: " + schema_dir + " ends")
table_dirs = os.listdir(schema_dir_path)
for table_dir in table_dirs:
table_dir_path = os.path.join(schema_dir_path, table_dir)
if not os.path.isdir(table_dir_path):
continue
column_ddl_file = os.path.join(table_dir_path, COLUMNS_STRUCTURE_FILE_NAME)
constraint_ddl_file = os.path.join(table_dir_path, CONSTRAINT_STRUCTURE_FILE_NAME)
index_ddl_file = os.path.join(table_dir_path, INDEX_STRUCTURE_FILE_NAME)
inherit_ddl_file = os.path.join(table_dir_path, INHERIT_STRUCTURE_FILE_NAME)
if os.path.exists(constraint_ddl_file): constraint_ddl_file_list.append(constraint_ddl_file)
if os.path.exists(index_ddl_file): index_ddl_file_list.append(index_ddl_file)
if os.path.exists(inherit_ddl_file): inherit_ddl_file_list.append(inherit_ddl_file)
import_file_sql(db_host, db_port, db_user, db_pass, db_dir, column_ddl_file)
print("[INFO]: restoring structure table: " + table_dir)
table_data_dir_path = os.path.join(table_dir_path, DATA_PATH_PREFIX)
if not os.path.isdir(table_data_dir_path):
continue
files_format = os.listdir(table_data_dir_path)[0].split(".")[-1]
if files_format == "csv":
csv_files = os.listdir(table_data_dir_path)
csv_count = 0
for csv_file in csv_files:
csv_file_path = os.path.join(table_data_dir_path, csv_file)
file_size = os.path.getsize(csv_file_path)
if file_size > 0:
import_file_csv(db_host, db_port, db_user, db_pass, csv_file_path, db_dir, schema_dir,
table_dir)
csv_count = csv_count + 1
print("[INFO]: restore data [" + str(csv_count) + "/" + str(
len(csv_files)) + "] of table " + db_dir + ": " + schema_dir + "." + table_dir)
elif files_format == "sql":
sql_files = os.listdir(table_data_dir_path)
sql_count = 0
for sql_file in sql_files:
sql_file_path = os.path.join(table_data_dir_path, sql_file)
file_size = os.path.getsize(sql_file_path)
if file_size > 0:
import_file_sql(db_host, db_port, db_user, db_pass, db_dir, sql_file_path)
sql_count = sql_count + 1
print("[INFO]: restore data [" + str(sql_count) + "/" + str(
len(sql_files)) + "] of table " + db_dir + ": " + schema_dir + "." + table_dir)
for file in constraint_ddl_file_list:
import_file_sql(db_host, db_port, db_user, db_pass, db_dir, file)
for file in index_ddl_file_list:
import_file_sql(db_host, db_port, db_user, db_pass, db_dir, file)
for file in inherit_ddl_file_list:
import_file_sql(db_host, db_port, db_user, db_pass, db_dir, file)