test_postgresql.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748
  1. import logging
  2. import pytest
  3. from flexmock import flexmock
  4. from borgmatic.hooks import postgresql as module
  5. def test_database_names_to_dump_passes_through_individual_database_name():
  6. database = {'name': 'foo'}
  7. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  8. 'foo',
  9. )
  10. def test_database_names_to_dump_passes_through_individual_database_name_with_format():
  11. database = {'name': 'foo', 'format': 'custom'}
  12. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  13. 'foo',
  14. )
  15. def test_database_names_to_dump_passes_through_all_without_format():
  16. database = {'name': 'all'}
  17. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  18. 'all',
  19. )
  20. def test_database_names_to_dump_with_all_and_format_and_dry_run_bails():
  21. database = {'name': 'all', 'format': 'custom'}
  22. flexmock(module).should_receive('execute_command_and_capture_output').never()
  23. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=True) == ()
  24. def test_database_names_to_dump_with_all_and_format_lists_databases():
  25. database = {'name': 'all', 'format': 'custom'}
  26. flexmock(module).should_receive('execute_command_and_capture_output').and_return(
  27. 'foo,test,\nbar,test,"stuff and such"'
  28. )
  29. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  30. 'foo',
  31. 'bar',
  32. )
  33. def test_database_names_to_dump_with_all_and_format_lists_databases_with_hostname_and_port():
  34. database = {'name': 'all', 'format': 'custom', 'hostname': 'localhost', 'port': 1234}
  35. flexmock(module).should_receive('execute_command_and_capture_output').with_args(
  36. (
  37. 'psql',
  38. '--list',
  39. '--no-password',
  40. '--csv',
  41. '--tuples-only',
  42. '--host',
  43. 'localhost',
  44. '--port',
  45. '1234',
  46. ),
  47. extra_environment=object,
  48. ).and_return('foo,test,\nbar,test,"stuff and such"')
  49. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  50. 'foo',
  51. 'bar',
  52. )
  53. def test_database_names_to_dump_with_all_and_format_lists_databases_with_username():
  54. database = {'name': 'all', 'format': 'custom', 'username': 'postgres'}
  55. flexmock(module).should_receive('execute_command_and_capture_output').with_args(
  56. ('psql', '--list', '--no-password', '--csv', '--tuples-only', '--username', 'postgres'),
  57. extra_environment=object,
  58. ).and_return('foo,test,\nbar,test,"stuff and such"')
  59. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  60. 'foo',
  61. 'bar',
  62. )
  63. def test_database_names_to_dump_with_all_and_format_lists_databases_with_options():
  64. database = {'name': 'all', 'format': 'custom', 'list_options': '--harder'}
  65. flexmock(module).should_receive('execute_command_and_capture_output').with_args(
  66. ('psql', '--list', '--no-password', '--csv', '--tuples-only', '--harder'),
  67. extra_environment=object,
  68. ).and_return('foo,test,\nbar,test,"stuff and such"')
  69. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  70. 'foo',
  71. 'bar',
  72. )
  73. def test_database_names_to_dump_with_all_and_format_excludes_particular_databases():
  74. database = {'name': 'all', 'format': 'custom'}
  75. flexmock(module).should_receive('execute_command_and_capture_output').and_return(
  76. 'foo,test,\ntemplate0,test,blah'
  77. )
  78. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  79. 'foo',
  80. )
  81. def test_dump_databases_runs_pg_dump_for_each_database():
  82. databases = [{'name': 'foo'}, {'name': 'bar'}]
  83. processes = [flexmock(), flexmock()]
  84. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  85. flexmock(module).should_receive('make_dump_path').and_return('')
  86. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',)).and_return(
  87. ('bar',)
  88. )
  89. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  90. 'databases/localhost/foo'
  91. ).and_return('databases/localhost/bar')
  92. flexmock(module.os.path).should_receive('exists').and_return(False)
  93. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  94. for name, process in zip(('foo', 'bar'), processes):
  95. flexmock(module).should_receive('execute_command').with_args(
  96. (
  97. 'pg_dump',
  98. '--no-password',
  99. '--clean',
  100. '--if-exists',
  101. '--format',
  102. 'custom',
  103. name,
  104. '>',
  105. f'databases/localhost/{name}',
  106. ),
  107. shell=True,
  108. extra_environment={'PGSSLMODE': 'disable'},
  109. run_to_completion=False,
  110. ).and_return(process).once()
  111. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == processes
  112. def test_dump_databases_raises_when_no_database_names_to_dump():
  113. databases = [{'name': 'foo'}, {'name': 'bar'}]
  114. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  115. flexmock(module).should_receive('make_dump_path').and_return('')
  116. flexmock(module).should_receive('database_names_to_dump').and_return(())
  117. with pytest.raises(ValueError):
  118. module.dump_databases(databases, 'test.yaml', {}, dry_run=False)
  119. def test_dump_databases_does_not_raise_when_no_database_names_to_dump():
  120. databases = [{'name': 'foo'}, {'name': 'bar'}]
  121. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  122. flexmock(module).should_receive('make_dump_path').and_return('')
  123. flexmock(module).should_receive('database_names_to_dump').and_return(())
  124. module.dump_databases(databases, 'test.yaml', {}, dry_run=True) == []
  125. def test_dump_databases_with_duplicate_dump_skips_pg_dump():
  126. databases = [{'name': 'foo'}, {'name': 'bar'}]
  127. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  128. flexmock(module).should_receive('make_dump_path').and_return('')
  129. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',)).and_return(
  130. ('bar',)
  131. )
  132. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  133. 'databases/localhost/foo'
  134. ).and_return('databases/localhost/bar')
  135. flexmock(module.os.path).should_receive('exists').and_return(True)
  136. flexmock(module.dump).should_receive('create_named_pipe_for_dump').never()
  137. flexmock(module).should_receive('execute_command').never()
  138. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == []
  139. def test_dump_databases_with_dry_run_skips_pg_dump():
  140. databases = [{'name': 'foo'}, {'name': 'bar'}]
  141. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  142. flexmock(module).should_receive('make_dump_path').and_return('')
  143. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',)).and_return(
  144. ('bar',)
  145. )
  146. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  147. 'databases/localhost/foo'
  148. ).and_return('databases/localhost/bar')
  149. flexmock(module.os.path).should_receive('exists').and_return(False)
  150. flexmock(module.dump).should_receive('create_named_pipe_for_dump').never()
  151. flexmock(module).should_receive('execute_command').never()
  152. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=True) == []
  153. def test_dump_databases_runs_pg_dump_with_hostname_and_port():
  154. databases = [{'name': 'foo', 'hostname': 'database.example.org', 'port': 5433}]
  155. process = flexmock()
  156. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  157. flexmock(module).should_receive('make_dump_path').and_return('')
  158. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',))
  159. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  160. 'databases/database.example.org/foo'
  161. )
  162. flexmock(module.os.path).should_receive('exists').and_return(False)
  163. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  164. flexmock(module).should_receive('execute_command').with_args(
  165. (
  166. 'pg_dump',
  167. '--no-password',
  168. '--clean',
  169. '--if-exists',
  170. '--host',
  171. 'database.example.org',
  172. '--port',
  173. '5433',
  174. '--format',
  175. 'custom',
  176. 'foo',
  177. '>',
  178. 'databases/database.example.org/foo',
  179. ),
  180. shell=True,
  181. extra_environment={'PGSSLMODE': 'disable'},
  182. run_to_completion=False,
  183. ).and_return(process).once()
  184. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
  185. def test_dump_databases_runs_pg_dump_with_username_and_password():
  186. databases = [{'name': 'foo', 'username': 'postgres', 'password': 'trustsome1'}]
  187. process = flexmock()
  188. flexmock(module).should_receive('make_extra_environment').and_return(
  189. {'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'}
  190. )
  191. flexmock(module).should_receive('make_dump_path').and_return('')
  192. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',))
  193. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  194. 'databases/localhost/foo'
  195. )
  196. flexmock(module.os.path).should_receive('exists').and_return(False)
  197. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  198. flexmock(module).should_receive('execute_command').with_args(
  199. (
  200. 'pg_dump',
  201. '--no-password',
  202. '--clean',
  203. '--if-exists',
  204. '--username',
  205. 'postgres',
  206. '--format',
  207. 'custom',
  208. 'foo',
  209. '>',
  210. 'databases/localhost/foo',
  211. ),
  212. shell=True,
  213. extra_environment={'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'},
  214. run_to_completion=False,
  215. ).and_return(process).once()
  216. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
  217. def test_make_extra_environment_maps_options_to_environment():
  218. database = {
  219. 'name': 'foo',
  220. 'password': 'pass',
  221. 'ssl_mode': 'require',
  222. 'ssl_cert': 'cert.crt',
  223. 'ssl_key': 'key.key',
  224. 'ssl_root_cert': 'root.crt',
  225. 'ssl_crl': 'crl.crl',
  226. }
  227. expected = {
  228. 'PGPASSWORD': 'pass',
  229. 'PGSSLMODE': 'require',
  230. 'PGSSLCERT': 'cert.crt',
  231. 'PGSSLKEY': 'key.key',
  232. 'PGSSLROOTCERT': 'root.crt',
  233. 'PGSSLCRL': 'crl.crl',
  234. }
  235. extra_env = module.make_extra_environment(database)
  236. assert extra_env == expected
  237. def test_dump_databases_runs_pg_dump_with_directory_format():
  238. databases = [{'name': 'foo', 'format': 'directory'}]
  239. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  240. flexmock(module).should_receive('make_dump_path').and_return('')
  241. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',))
  242. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  243. 'databases/localhost/foo'
  244. )
  245. flexmock(module.os.path).should_receive('exists').and_return(False)
  246. flexmock(module.dump).should_receive('create_parent_directory_for_dump')
  247. flexmock(module.dump).should_receive('create_named_pipe_for_dump').never()
  248. flexmock(module).should_receive('execute_command').with_args(
  249. (
  250. 'pg_dump',
  251. '--no-password',
  252. '--clean',
  253. '--if-exists',
  254. '--format',
  255. 'directory',
  256. '--file',
  257. 'databases/localhost/foo',
  258. 'foo',
  259. ),
  260. shell=True,
  261. extra_environment={'PGSSLMODE': 'disable'},
  262. ).and_return(flexmock()).once()
  263. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == []
  264. def test_dump_databases_runs_pg_dump_with_options():
  265. databases = [{'name': 'foo', 'options': '--stuff=such'}]
  266. process = flexmock()
  267. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  268. flexmock(module).should_receive('make_dump_path').and_return('')
  269. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',))
  270. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  271. 'databases/localhost/foo'
  272. )
  273. flexmock(module.os.path).should_receive('exists').and_return(False)
  274. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  275. flexmock(module).should_receive('execute_command').with_args(
  276. (
  277. 'pg_dump',
  278. '--no-password',
  279. '--clean',
  280. '--if-exists',
  281. '--format',
  282. 'custom',
  283. '--stuff=such',
  284. 'foo',
  285. '>',
  286. 'databases/localhost/foo',
  287. ),
  288. shell=True,
  289. extra_environment={'PGSSLMODE': 'disable'},
  290. run_to_completion=False,
  291. ).and_return(process).once()
  292. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
  293. def test_dump_databases_runs_pg_dumpall_for_all_databases():
  294. databases = [{'name': 'all'}]
  295. process = flexmock()
  296. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  297. flexmock(module).should_receive('make_dump_path').and_return('')
  298. flexmock(module).should_receive('database_names_to_dump').and_return(('all',))
  299. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  300. 'databases/localhost/all'
  301. )
  302. flexmock(module.os.path).should_receive('exists').and_return(False)
  303. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  304. flexmock(module).should_receive('execute_command').with_args(
  305. ('pg_dumpall', '--no-password', '--clean', '--if-exists', '>', 'databases/localhost/all'),
  306. shell=True,
  307. extra_environment={'PGSSLMODE': 'disable'},
  308. run_to_completion=False,
  309. ).and_return(process).once()
  310. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
  311. def test_dump_databases_runs_non_default_pg_dump():
  312. databases = [{'name': 'foo', 'pg_dump_command': 'special_pg_dump'}]
  313. process = flexmock()
  314. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  315. flexmock(module).should_receive('make_dump_path').and_return('')
  316. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',))
  317. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  318. 'databases/localhost/foo'
  319. )
  320. flexmock(module.os.path).should_receive('exists').and_return(False)
  321. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  322. flexmock(module).should_receive('execute_command').with_args(
  323. (
  324. 'special_pg_dump',
  325. '--no-password',
  326. '--clean',
  327. '--if-exists',
  328. '--format',
  329. 'custom',
  330. 'foo',
  331. '>',
  332. 'databases/localhost/foo',
  333. ),
  334. shell=True,
  335. extra_environment={'PGSSLMODE': 'disable'},
  336. run_to_completion=False,
  337. ).and_return(process).once()
  338. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
  339. def test_restore_database_dump_runs_pg_restore():
  340. database_config = [{'name': 'foo', 'schemas': None}]
  341. extract_process = flexmock(stdout=flexmock())
  342. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  343. flexmock(module).should_receive('make_dump_path')
  344. flexmock(module.dump).should_receive('make_database_dump_filename')
  345. flexmock(module).should_receive('execute_command_with_processes').with_args(
  346. (
  347. 'pg_restore',
  348. '--no-password',
  349. '--if-exists',
  350. '--exit-on-error',
  351. '--clean',
  352. '--dbname',
  353. 'foo',
  354. ),
  355. processes=[extract_process],
  356. output_log_level=logging.DEBUG,
  357. input_file=extract_process.stdout,
  358. extra_environment={'PGSSLMODE': 'disable'},
  359. ).once()
  360. flexmock(module).should_receive('execute_command').with_args(
  361. ('psql', '--no-password', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
  362. extra_environment={'PGSSLMODE': 'disable'},
  363. ).once()
  364. module.restore_database_dump(
  365. database_config, 'test.yaml', {}, dry_run=False, extract_process=extract_process
  366. )
  367. def test_restore_database_dump_errors_on_multiple_database_config():
  368. database_config = [{'name': 'foo'}, {'name': 'bar'}]
  369. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  370. flexmock(module).should_receive('make_dump_path')
  371. flexmock(module.dump).should_receive('make_database_dump_filename')
  372. flexmock(module).should_receive('execute_command_with_processes').never()
  373. flexmock(module).should_receive('execute_command').never()
  374. with pytest.raises(ValueError):
  375. module.restore_database_dump(
  376. database_config, 'test.yaml', {}, dry_run=False, extract_process=flexmock()
  377. )
  378. def test_restore_database_dump_runs_pg_restore_with_hostname_and_port():
  379. database_config = [
  380. {'name': 'foo', 'hostname': 'database.example.org', 'port': 5433, 'schemas': None}
  381. ]
  382. extract_process = flexmock(stdout=flexmock())
  383. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  384. flexmock(module).should_receive('make_dump_path')
  385. flexmock(module.dump).should_receive('make_database_dump_filename')
  386. flexmock(module).should_receive('execute_command_with_processes').with_args(
  387. (
  388. 'pg_restore',
  389. '--no-password',
  390. '--if-exists',
  391. '--exit-on-error',
  392. '--clean',
  393. '--dbname',
  394. 'foo',
  395. '--host',
  396. 'database.example.org',
  397. '--port',
  398. '5433',
  399. ),
  400. processes=[extract_process],
  401. output_log_level=logging.DEBUG,
  402. input_file=extract_process.stdout,
  403. extra_environment={'PGSSLMODE': 'disable'},
  404. ).once()
  405. flexmock(module).should_receive('execute_command').with_args(
  406. (
  407. 'psql',
  408. '--no-password',
  409. '--quiet',
  410. '--host',
  411. 'database.example.org',
  412. '--port',
  413. '5433',
  414. '--dbname',
  415. 'foo',
  416. '--command',
  417. 'ANALYZE',
  418. ),
  419. extra_environment={'PGSSLMODE': 'disable'},
  420. ).once()
  421. module.restore_database_dump(
  422. database_config, 'test.yaml', {}, dry_run=False, extract_process=extract_process
  423. )
  424. def test_restore_database_dump_runs_pg_restore_with_username_and_password():
  425. database_config = [
  426. {'name': 'foo', 'username': 'postgres', 'password': 'trustsome1', 'schemas': None}
  427. ]
  428. extract_process = flexmock(stdout=flexmock())
  429. flexmock(module).should_receive('make_extra_environment').and_return(
  430. {'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'}
  431. )
  432. flexmock(module).should_receive('make_dump_path')
  433. flexmock(module.dump).should_receive('make_database_dump_filename')
  434. flexmock(module).should_receive('execute_command_with_processes').with_args(
  435. (
  436. 'pg_restore',
  437. '--no-password',
  438. '--if-exists',
  439. '--exit-on-error',
  440. '--clean',
  441. '--dbname',
  442. 'foo',
  443. '--username',
  444. 'postgres',
  445. ),
  446. processes=[extract_process],
  447. output_log_level=logging.DEBUG,
  448. input_file=extract_process.stdout,
  449. extra_environment={'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'},
  450. ).once()
  451. flexmock(module).should_receive('execute_command').with_args(
  452. (
  453. 'psql',
  454. '--no-password',
  455. '--quiet',
  456. '--username',
  457. 'postgres',
  458. '--dbname',
  459. 'foo',
  460. '--command',
  461. 'ANALYZE',
  462. ),
  463. extra_environment={'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'},
  464. ).once()
  465. module.restore_database_dump(
  466. database_config, 'test.yaml', {}, dry_run=False, extract_process=extract_process
  467. )
  468. def test_restore_database_dump_runs_pg_restore_with_options():
  469. database_config = [
  470. {
  471. 'name': 'foo',
  472. 'restore_options': '--harder',
  473. 'analyze_options': '--smarter',
  474. 'schemas': None,
  475. }
  476. ]
  477. extract_process = flexmock(stdout=flexmock())
  478. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  479. flexmock(module).should_receive('make_dump_path')
  480. flexmock(module.dump).should_receive('make_database_dump_filename')
  481. flexmock(module).should_receive('execute_command_with_processes').with_args(
  482. (
  483. 'pg_restore',
  484. '--no-password',
  485. '--if-exists',
  486. '--exit-on-error',
  487. '--clean',
  488. '--dbname',
  489. 'foo',
  490. '--harder',
  491. ),
  492. processes=[extract_process],
  493. output_log_level=logging.DEBUG,
  494. input_file=extract_process.stdout,
  495. extra_environment={'PGSSLMODE': 'disable'},
  496. ).once()
  497. flexmock(module).should_receive('execute_command').with_args(
  498. (
  499. 'psql',
  500. '--no-password',
  501. '--quiet',
  502. '--dbname',
  503. 'foo',
  504. '--smarter',
  505. '--command',
  506. 'ANALYZE',
  507. ),
  508. extra_environment={'PGSSLMODE': 'disable'},
  509. ).once()
  510. module.restore_database_dump(
  511. database_config, 'test.yaml', {}, dry_run=False, extract_process=extract_process
  512. )
  513. def test_restore_database_dump_runs_psql_for_all_database_dump():
  514. database_config = [{'name': 'all', 'schemas': None}]
  515. extract_process = flexmock(stdout=flexmock())
  516. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  517. flexmock(module).should_receive('make_dump_path')
  518. flexmock(module.dump).should_receive('make_database_dump_filename')
  519. flexmock(module).should_receive('execute_command_with_processes').with_args(
  520. ('psql', '--no-password'),
  521. processes=[extract_process],
  522. output_log_level=logging.DEBUG,
  523. input_file=extract_process.stdout,
  524. extra_environment={'PGSSLMODE': 'disable'},
  525. ).once()
  526. flexmock(module).should_receive('execute_command').with_args(
  527. ('psql', '--no-password', '--quiet', '--command', 'ANALYZE'),
  528. extra_environment={'PGSSLMODE': 'disable'},
  529. ).once()
  530. module.restore_database_dump(
  531. database_config, 'test.yaml', {}, dry_run=False, extract_process=extract_process
  532. )
  533. def test_restore_database_dump_runs_non_default_pg_restore_and_psql():
  534. database_config = [
  535. {
  536. 'name': 'foo',
  537. 'pg_restore_command': 'special_pg_restore',
  538. 'psql_command': 'special_psql',
  539. 'schemas': None,
  540. }
  541. ]
  542. extract_process = flexmock(stdout=flexmock())
  543. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  544. flexmock(module).should_receive('make_dump_path')
  545. flexmock(module.dump).should_receive('make_database_dump_filename')
  546. flexmock(module).should_receive('execute_command_with_processes').with_args(
  547. (
  548. 'special_pg_restore',
  549. '--no-password',
  550. '--if-exists',
  551. '--exit-on-error',
  552. '--clean',
  553. '--dbname',
  554. 'foo',
  555. ),
  556. processes=[extract_process],
  557. output_log_level=logging.DEBUG,
  558. input_file=extract_process.stdout,
  559. extra_environment={'PGSSLMODE': 'disable'},
  560. ).once()
  561. flexmock(module).should_receive('execute_command').with_args(
  562. ('special_psql', '--no-password', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
  563. extra_environment={'PGSSLMODE': 'disable'},
  564. ).once()
  565. module.restore_database_dump(
  566. database_config, 'test.yaml', {}, dry_run=False, extract_process=extract_process
  567. )
  568. def test_restore_database_dump_with_dry_run_skips_restore():
  569. database_config = [{'name': 'foo', 'schemas': None}]
  570. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  571. flexmock(module).should_receive('make_dump_path')
  572. flexmock(module.dump).should_receive('make_database_dump_filename')
  573. flexmock(module).should_receive('execute_command_with_processes').never()
  574. module.restore_database_dump(
  575. database_config, 'test.yaml', {}, dry_run=True, extract_process=flexmock()
  576. )
  577. def test_restore_database_dump_without_extract_process_restores_from_disk():
  578. database_config = [{'name': 'foo', 'schemas': None}]
  579. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  580. flexmock(module).should_receive('make_dump_path')
  581. flexmock(module.dump).should_receive('make_database_dump_filename').and_return('/dump/path')
  582. flexmock(module).should_receive('execute_command_with_processes').with_args(
  583. (
  584. 'pg_restore',
  585. '--no-password',
  586. '--if-exists',
  587. '--exit-on-error',
  588. '--clean',
  589. '--dbname',
  590. 'foo',
  591. '/dump/path',
  592. ),
  593. processes=[],
  594. output_log_level=logging.DEBUG,
  595. input_file=None,
  596. extra_environment={'PGSSLMODE': 'disable'},
  597. ).once()
  598. flexmock(module).should_receive('execute_command').with_args(
  599. ('psql', '--no-password', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
  600. extra_environment={'PGSSLMODE': 'disable'},
  601. ).once()
  602. module.restore_database_dump(
  603. database_config, 'test.yaml', {}, dry_run=False, extract_process=None
  604. )
  605. def test_restore_database_dump_with_schemas_restores_schemas():
  606. database_config = [{'name': 'foo', 'schemas': ['bar', 'baz']}]
  607. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  608. flexmock(module).should_receive('make_dump_path')
  609. flexmock(module.dump).should_receive('make_database_dump_filename').and_return('/dump/path')
  610. flexmock(module).should_receive('execute_command_with_processes').with_args(
  611. (
  612. 'pg_restore',
  613. '--no-password',
  614. '--if-exists',
  615. '--exit-on-error',
  616. '--clean',
  617. '--dbname',
  618. 'foo',
  619. '/dump/path',
  620. '--schema',
  621. 'bar',
  622. '--schema',
  623. 'baz',
  624. ),
  625. processes=[],
  626. output_log_level=logging.DEBUG,
  627. input_file=None,
  628. extra_environment={'PGSSLMODE': 'disable'},
  629. ).once()
  630. flexmock(module).should_receive('execute_command').with_args(
  631. ('psql', '--no-password', '--quiet', '--dbname', 'foo', '--command', 'ANALYZE'),
  632. extra_environment={'PGSSLMODE': 'disable'},
  633. ).once()
  634. module.restore_database_dump(
  635. database_config, 'test.yaml', {}, dry_run=False, extract_process=None
  636. )