test_postgresql.py 27 KB

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