gooderp18绿色标准版
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

441 lines
14KB

  1. #########################################################################
  2. #
  3. # pgAdmin 4 - PostgreSQL Tools
  4. #
  5. # Copyright (C) 2013 - 2020, The pgAdmin Development Team
  6. # This software is released under the PostgreSQL Licence
  7. #
  8. ##########################################################################
  9. """Perform the initial setup of the application, by creating the auth
  10. and settings database."""
  11. import argparse
  12. import json
  13. import os
  14. import sys
  15. import builtins
  16. from pgadmin.model import db, User, Version, ServerGroup, Server, \
  17. SCHEMA_VERSION as CURRENT_SCHEMA_VERSION
  18. # Grab the SERVER_MODE if it's been set by the runtime
  19. if 'SERVER_MODE' in globals():
  20. builtins.SERVER_MODE = globals()['SERVER_MODE']
  21. else:
  22. builtins.SERVER_MODE = None
  23. # We need to include the root directory in sys.path to ensure that we can
  24. # find everything we need when running in the standalone runtime.
  25. root = os.path.dirname(os.path.realpath(__file__))
  26. if sys.path[0] != root:
  27. sys.path.insert(0, root)
  28. from pgadmin import create_app
  29. def add_value(attr_dict, key, value):
  30. """Add a value to the attribute dict if non-empty.
  31. Args:
  32. attr_dict (dict): The dictionary to add the values to
  33. key (str): The key for the new value
  34. value (str): The value to add
  35. Returns:
  36. The updated attribute dictionary
  37. """
  38. if value != "" and value is not None:
  39. attr_dict[key] = value
  40. return attr_dict
  41. def dump_servers(args):
  42. """Dump the server groups and servers.
  43. Args:
  44. args (ArgParser): The parsed command line options
  45. """
  46. # What user?
  47. if args.user is not None:
  48. dump_user = args.user
  49. else:
  50. dump_user = config.DESKTOP_USER
  51. # And the sqlite path
  52. if args.sqlite_path is not None:
  53. config.SQLITE_PATH = args.sqlite_path
  54. print('----------')
  55. print('Dumping servers with:')
  56. print('User:', dump_user)
  57. print('SQLite pgAdmin config:', config.SQLITE_PATH)
  58. print('----------')
  59. app = create_app(config.APP_NAME + '-cli')
  60. with app.app_context():
  61. user = User.query.filter_by(email=dump_user).first()
  62. if user is None:
  63. print("The specified user ID (%s) could not be found." %
  64. dump_user)
  65. sys.exit(1)
  66. user_id = user.id
  67. # Dict to collect the output
  68. object_dict = {}
  69. # Counters
  70. servers_dumped = 0
  71. # Dump servers
  72. servers = Server.query.filter_by(user_id=user_id).all()
  73. server_dict = {}
  74. for server in servers:
  75. if args.servers is None or str(server.id) in args.servers:
  76. # Get the group name
  77. group_name = ServerGroup.query.filter_by(
  78. user_id=user_id, id=server.servergroup_id).first().name
  79. attr_dict = {}
  80. add_value(attr_dict, "Name", server.name)
  81. add_value(attr_dict, "Group", group_name)
  82. add_value(attr_dict, "Host", server.host)
  83. add_value(attr_dict, "HostAddr", server.hostaddr)
  84. add_value(attr_dict, "Port", server.port)
  85. add_value(attr_dict, "MaintenanceDB", server.maintenance_db)
  86. add_value(attr_dict, "Username", server.username)
  87. add_value(attr_dict, "Role", server.role)
  88. add_value(attr_dict, "SSLMode", server.ssl_mode)
  89. add_value(attr_dict, "Comment", server.comment)
  90. add_value(attr_dict, "DBRestriction", server.db_res)
  91. add_value(attr_dict, "PassFile", server.passfile)
  92. add_value(attr_dict, "SSLCert", server.sslcert)
  93. add_value(attr_dict, "SSLKey", server.sslkey)
  94. add_value(attr_dict, "SSLRootCert", server.sslrootcert)
  95. add_value(attr_dict, "SSLCrl", server.sslcrl)
  96. add_value(attr_dict, "SSLCompression", server.sslcompression)
  97. add_value(attr_dict, "BGColor", server.bgcolor)
  98. add_value(attr_dict, "FGColor", server.fgcolor)
  99. add_value(attr_dict, "Service", server.service)
  100. add_value(attr_dict, "Timeout", server.connect_timeout)
  101. add_value(attr_dict, "UseSSHTunnel", server.use_ssh_tunnel)
  102. add_value(attr_dict, "TunnelHost", server.tunnel_host)
  103. add_value(attr_dict, "TunnelPort", server.tunnel_port)
  104. add_value(attr_dict, "TunnelUsername", server.tunnel_username)
  105. add_value(attr_dict, "TunnelAuthentication",
  106. server.tunnel_authentication)
  107. servers_dumped = servers_dumped + 1
  108. server_dict[servers_dumped] = attr_dict
  109. object_dict["Servers"] = server_dict
  110. try:
  111. f = open(args.dump_servers, "w")
  112. except Exception as e:
  113. print("Error opening output file %s: [%d] %s" %
  114. (args.dump_servers, e.errno, e.strerror))
  115. sys.exit(1)
  116. try:
  117. f.write(json.dumps(object_dict, indent=4))
  118. except Exception as e:
  119. print("Error writing output file %s: [%d] %s" %
  120. (args.dump_servers, e.errno, e.strerror))
  121. sys.exit(1)
  122. f.close()
  123. print("Configuration for %s servers dumped to %s." %
  124. (servers_dumped, args.dump_servers))
  125. def load_servers(args):
  126. """Load server groups and servers.
  127. Args:
  128. args (ArgParser): The parsed command line options
  129. """
  130. # What user?
  131. if args.user is not None:
  132. load_user = args.user
  133. else:
  134. load_user = config.DESKTOP_USER
  135. # And the sqlite path
  136. if args.sqlite_path is not None:
  137. config.SQLITE_PATH = args.sqlite_path
  138. print('----------')
  139. print('Loading servers with:')
  140. print('User:', load_user)
  141. print('SQLite pgAdmin config:', config.SQLITE_PATH)
  142. print('----------')
  143. try:
  144. with open(args.load_servers) as f:
  145. data = json.load(f)
  146. except json.decoder.JSONDecodeError as e:
  147. print("Error parsing input file %s: %s" %
  148. (args.load_servers, e))
  149. sys.exit(1)
  150. except Exception as e:
  151. print("Error reading input file %s: [%d] %s" %
  152. (args.load_servers, e.errno, e.strerror))
  153. sys.exit(1)
  154. f.close()
  155. app = create_app(config.APP_NAME + '-cli')
  156. with app.app_context():
  157. user = User.query.filter_by(email=load_user).first()
  158. if user is None:
  159. print("The specified user ID (%s) could not be found." %
  160. load_user)
  161. sys.exit(1)
  162. user_id = user.id
  163. # Counters
  164. groups_added = 0
  165. servers_added = 0
  166. # Get the server groups
  167. groups = ServerGroup.query.filter_by(user_id=user_id)
  168. def print_summary():
  169. print("Added %d Server Group(s) and %d Server(s)." %
  170. (groups_added, servers_added))
  171. # Loop through the servers...
  172. if "Servers" not in data:
  173. print("'Servers' attribute not found in file '%s'" %
  174. args.load_servers)
  175. print_summary()
  176. sys.exit(1)
  177. for server in data["Servers"]:
  178. obj = data["Servers"][server]
  179. def check_attrib(attrib):
  180. if attrib not in obj:
  181. print("'%s' attribute not found for server '%s'" %
  182. (attrib, server))
  183. print_summary()
  184. sys.exit(1)
  185. check_attrib("Name")
  186. check_attrib("Group")
  187. check_attrib("Port")
  188. check_attrib("Username")
  189. check_attrib("SSLMode")
  190. check_attrib("MaintenanceDB")
  191. if "Host" not in obj and \
  192. "HostAddr" not in obj and \
  193. "Service" not in obj:
  194. print("'Host', 'HostAddr' or 'Service' attribute not found "
  195. "for server '%s'" % server)
  196. print_summary()
  197. sys.exit(1)
  198. # Get the group. Create if necessary
  199. group_id = -1
  200. for g in groups:
  201. if g.name == obj["Group"]:
  202. group_id = g.id
  203. break
  204. if group_id == -1:
  205. new_group = ServerGroup()
  206. new_group.name = obj["Group"]
  207. new_group.user_id = user_id
  208. db.session.add(new_group)
  209. try:
  210. db.session.commit()
  211. except Exception as e:
  212. print("Error creating server group '%s': %s" %
  213. (new_group.name, e))
  214. print_summary()
  215. sys.exit(1)
  216. group_id = new_group.id
  217. groups_added = groups_added + 1
  218. groups = ServerGroup.query.filter_by(user_id=user_id)
  219. # Create the server
  220. new_server = Server()
  221. new_server.name = obj["Name"]
  222. new_server.servergroup_id = group_id
  223. new_server.user_id = user_id
  224. new_server.host = obj["Host"]
  225. if "HostAddr" in obj:
  226. new_server.hostaddr = obj["HostAddr"]
  227. new_server.port = obj["Port"]
  228. new_server.maintenance_db = obj["MaintenanceDB"]
  229. new_server.username = obj["Username"]
  230. if "Role" in obj:
  231. new_server.role = obj["Role"]
  232. new_server.ssl_mode = obj["SSLMode"]
  233. if "Comment" in obj:
  234. new_server.comment = obj["Comment"]
  235. if "DBRestriction" in obj:
  236. new_server.db_res = obj["DBRestriction"]
  237. if "PassFile" in obj:
  238. new_server.passfile = obj["PassFile"]
  239. if "SSLCert" in obj:
  240. new_server.sslcert = obj["SSLCert"]
  241. if "SSLKey" in obj:
  242. new_server.sslkey = obj["SSLKey"]
  243. if "SSLRootCert" in obj:
  244. new_server.sslrootcert = obj["SSLRootCert"]
  245. if "SSLCrl" in obj:
  246. new_server.sslcrl = obj["SSLCrl"]
  247. if "SSLCompression" in obj:
  248. new_server.sslcompression = obj["SSLCompression"]
  249. if "BGColor" in obj:
  250. new_server.bgcolor = obj["BGColor"]
  251. if "FGColor" in obj:
  252. new_server.fgcolor = obj["FGColor"]
  253. if "Service" in obj:
  254. new_server.service = obj["Service"]
  255. if "Timeout" in obj:
  256. new_server.connect_timeout = obj["Timeout"]
  257. if "UseSSHTunnel" in obj:
  258. new_server.use_ssh_tunnel = obj["UseSSHTunnel"]
  259. if "TunnelHost" in obj:
  260. new_server.tunnel_host = obj["TunnelHost"]
  261. if "TunnelPort" in obj:
  262. new_server.tunnel_port = obj["TunnelPort"]
  263. if "TunnelUsername" in obj:
  264. new_server.tunnel_username = obj["TunnelUsername"]
  265. if "TunnelAuthentication" in obj:
  266. new_server.tunnel_authentication = obj["TunnelAuthentication"]
  267. db.session.add(new_server)
  268. try:
  269. db.session.commit()
  270. except Exception as e:
  271. print("Error creating server '%s': %s" %
  272. (new_server.name, e))
  273. print_summary()
  274. sys.exit(1)
  275. servers_added = servers_added + 1
  276. print_summary()
  277. def setup_db():
  278. """Setup the configuration database."""
  279. create_app_data_directory(config)
  280. app = create_app()
  281. print(u"pgAdmin 4 - Application Initialisation")
  282. print(u"======================================\n")
  283. with app.app_context():
  284. # Run migration for the first time i.e. create database
  285. from config import SQLITE_PATH
  286. if not os.path.exists(SQLITE_PATH):
  287. db_upgrade(app)
  288. else:
  289. version = Version.query.filter_by(name='ConfigDB').first()
  290. schema_version = version.value
  291. # Run migration if current schema version is greater than the
  292. # schema version stored in version table
  293. if CURRENT_SCHEMA_VERSION >= schema_version:
  294. db_upgrade(app)
  295. # Update schema version to the latest
  296. if CURRENT_SCHEMA_VERSION > schema_version:
  297. version = Version.query.filter_by(name='ConfigDB').first()
  298. version.value = CURRENT_SCHEMA_VERSION
  299. db.session.commit()
  300. if os.name != 'nt':
  301. os.chmod(config.SQLITE_PATH, 0o600)
  302. if __name__ == '__main__':
  303. # Configuration settings
  304. import config
  305. from pgadmin.model import SCHEMA_VERSION
  306. from pgadmin.setup import db_upgrade, create_app_data_directory
  307. parser = argparse.ArgumentParser(description='Setup the pgAdmin config DB')
  308. exp_group = parser.add_argument_group('Dump server config')
  309. exp_group.add_argument('--dump-servers', metavar="OUTPUT_FILE",
  310. help='Dump the servers in the DB', required=False)
  311. exp_group.add_argument('--servers', metavar="SERVERS", nargs='*',
  312. help='One or more servers to dump', required=False)
  313. imp_group = parser.add_argument_group('Load server config')
  314. imp_group.add_argument('--load-servers', metavar="INPUT_FILE",
  315. help='Load servers into the DB', required=False)
  316. # Common args
  317. parser.add_argument('--sqlite-path', metavar="PATH",
  318. help='Dump/load with the specified pgAdmin config DB'
  319. ' file. This is particularly helpful when there'
  320. ' are multiple pgAdmin configurations. It is also'
  321. ' recommended to use this option when running'
  322. ' pgAdmin in desktop mode.', required=False)
  323. parser.add_argument('--user', metavar="USER_NAME",
  324. help='Dump/load servers for the specified username',
  325. required=False)
  326. args, extra = parser.parse_known_args()
  327. config.SETTINGS_SCHEMA_VERSION = SCHEMA_VERSION
  328. if "PGADMIN_TESTING_MODE" in os.environ and \
  329. os.environ["PGADMIN_TESTING_MODE"] == "1":
  330. config.SQLITE_PATH = config.TEST_SQLITE_PATH
  331. # What to do?
  332. if args.dump_servers is not None:
  333. try:
  334. dump_servers(args)
  335. except Exception as e:
  336. print(str(e))
  337. elif args.load_servers is not None:
  338. try:
  339. load_servers(args)
  340. except Exception as e:
  341. print(str(e))
  342. else:
  343. setup_db()
上海开阖软件有限公司 沪ICP备12045867号-1