test_postgresql.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972
  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. (
  58. 'psql',
  59. '--list',
  60. '--no-password',
  61. '--no-psqlrc',
  62. '--csv',
  63. '--tuples-only',
  64. '--username',
  65. 'postgres',
  66. ),
  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_lists_databases_with_options():
  74. database = {'name': 'all', 'format': 'custom', 'list_options': '--harder'}
  75. flexmock(module).should_receive('execute_command_and_capture_output').with_args(
  76. ('psql', '--list', '--no-password', '--no-psqlrc', '--csv', '--tuples-only', '--harder'),
  77. extra_environment=object,
  78. ).and_return('foo,test,\nbar,test,"stuff and such"')
  79. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  80. 'foo',
  81. 'bar',
  82. )
  83. def test_database_names_to_dump_with_all_and_format_excludes_particular_databases():
  84. database = {'name': 'all', 'format': 'custom'}
  85. flexmock(module).should_receive('execute_command_and_capture_output').and_return(
  86. 'foo,test,\ntemplate0,test,blah'
  87. )
  88. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  89. 'foo',
  90. )
  91. def test_database_names_to_dump_with_all_and_psql_command_uses_custom_command():
  92. database = {'name': 'all', 'format': 'custom', 'psql_command': 'docker exec mycontainer psql'}
  93. flexmock(module).should_receive('execute_command_and_capture_output').with_args(
  94. (
  95. 'docker',
  96. 'exec',
  97. 'mycontainer',
  98. 'psql',
  99. '--list',
  100. '--no-password',
  101. '--no-psqlrc',
  102. '--csv',
  103. '--tuples-only',
  104. ),
  105. extra_environment=object,
  106. ).and_return('foo,text').once()
  107. assert module.database_names_to_dump(database, flexmock(), flexmock(), dry_run=False) == (
  108. 'foo',
  109. )
  110. def test_dump_databases_runs_pg_dump_for_each_database():
  111. databases = [{'name': 'foo'}, {'name': 'bar'}]
  112. processes = [flexmock(), flexmock()]
  113. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  114. flexmock(module).should_receive('make_dump_path').and_return('')
  115. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',)).and_return(
  116. ('bar',)
  117. )
  118. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  119. 'databases/localhost/foo'
  120. ).and_return('databases/localhost/bar')
  121. flexmock(module.os.path).should_receive('exists').and_return(False)
  122. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  123. for name, process in zip(('foo', 'bar'), processes):
  124. flexmock(module).should_receive('execute_command').with_args(
  125. (
  126. 'pg_dump',
  127. '--no-password',
  128. '--clean',
  129. '--if-exists',
  130. '--format',
  131. 'custom',
  132. name,
  133. '>',
  134. f'databases/localhost/{name}',
  135. ),
  136. shell=True,
  137. extra_environment={'PGSSLMODE': 'disable'},
  138. run_to_completion=False,
  139. ).and_return(process).once()
  140. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == processes
  141. def test_dump_databases_raises_when_no_database_names_to_dump():
  142. databases = [{'name': 'foo'}, {'name': 'bar'}]
  143. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  144. flexmock(module).should_receive('make_dump_path').and_return('')
  145. flexmock(module).should_receive('database_names_to_dump').and_return(())
  146. with pytest.raises(ValueError):
  147. module.dump_databases(databases, 'test.yaml', {}, dry_run=False)
  148. def test_dump_databases_does_not_raise_when_no_database_names_to_dump():
  149. databases = [{'name': 'foo'}, {'name': 'bar'}]
  150. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  151. flexmock(module).should_receive('make_dump_path').and_return('')
  152. flexmock(module).should_receive('database_names_to_dump').and_return(())
  153. module.dump_databases(databases, 'test.yaml', {}, dry_run=True) == []
  154. def test_dump_databases_with_duplicate_dump_skips_pg_dump():
  155. databases = [{'name': 'foo'}, {'name': 'bar'}]
  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',)).and_return(
  159. ('bar',)
  160. )
  161. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  162. 'databases/localhost/foo'
  163. ).and_return('databases/localhost/bar')
  164. flexmock(module.os.path).should_receive('exists').and_return(True)
  165. flexmock(module.dump).should_receive('create_named_pipe_for_dump').never()
  166. flexmock(module).should_receive('execute_command').never()
  167. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == []
  168. def test_dump_databases_with_dry_run_skips_pg_dump():
  169. databases = [{'name': 'foo'}, {'name': 'bar'}]
  170. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  171. flexmock(module).should_receive('make_dump_path').and_return('')
  172. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',)).and_return(
  173. ('bar',)
  174. )
  175. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  176. 'databases/localhost/foo'
  177. ).and_return('databases/localhost/bar')
  178. flexmock(module.os.path).should_receive('exists').and_return(False)
  179. flexmock(module.dump).should_receive('create_named_pipe_for_dump').never()
  180. flexmock(module).should_receive('execute_command').never()
  181. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=True) == []
  182. def test_dump_databases_runs_pg_dump_with_hostname_and_port():
  183. databases = [{'name': 'foo', 'hostname': 'database.example.org', 'port': 5433}]
  184. process = flexmock()
  185. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  186. flexmock(module).should_receive('make_dump_path').and_return('')
  187. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',))
  188. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  189. 'databases/database.example.org/foo'
  190. )
  191. flexmock(module.os.path).should_receive('exists').and_return(False)
  192. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  193. flexmock(module).should_receive('execute_command').with_args(
  194. (
  195. 'pg_dump',
  196. '--no-password',
  197. '--clean',
  198. '--if-exists',
  199. '--host',
  200. 'database.example.org',
  201. '--port',
  202. '5433',
  203. '--format',
  204. 'custom',
  205. 'foo',
  206. '>',
  207. 'databases/database.example.org/foo',
  208. ),
  209. shell=True,
  210. extra_environment={'PGSSLMODE': 'disable'},
  211. run_to_completion=False,
  212. ).and_return(process).once()
  213. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
  214. def test_dump_databases_runs_pg_dump_with_username_and_password():
  215. databases = [{'name': 'foo', 'username': 'postgres', 'password': 'trustsome1'}]
  216. process = flexmock()
  217. flexmock(module).should_receive('make_extra_environment').and_return(
  218. {'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'}
  219. )
  220. flexmock(module).should_receive('make_dump_path').and_return('')
  221. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',))
  222. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  223. 'databases/localhost/foo'
  224. )
  225. flexmock(module.os.path).should_receive('exists').and_return(False)
  226. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  227. flexmock(module).should_receive('execute_command').with_args(
  228. (
  229. 'pg_dump',
  230. '--no-password',
  231. '--clean',
  232. '--if-exists',
  233. '--username',
  234. 'postgres',
  235. '--format',
  236. 'custom',
  237. 'foo',
  238. '>',
  239. 'databases/localhost/foo',
  240. ),
  241. shell=True,
  242. extra_environment={'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'},
  243. run_to_completion=False,
  244. ).and_return(process).once()
  245. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
  246. def test_make_extra_environment_maps_options_to_environment():
  247. database = {
  248. 'name': 'foo',
  249. 'password': 'pass',
  250. 'ssl_mode': 'require',
  251. 'ssl_cert': 'cert.crt',
  252. 'ssl_key': 'key.key',
  253. 'ssl_root_cert': 'root.crt',
  254. 'ssl_crl': 'crl.crl',
  255. }
  256. expected = {
  257. 'PGPASSWORD': 'pass',
  258. 'PGSSLMODE': 'require',
  259. 'PGSSLCERT': 'cert.crt',
  260. 'PGSSLKEY': 'key.key',
  261. 'PGSSLROOTCERT': 'root.crt',
  262. 'PGSSLCRL': 'crl.crl',
  263. }
  264. extra_env = module.make_extra_environment(database)
  265. assert extra_env == expected
  266. def test_dump_databases_runs_pg_dump_with_directory_format():
  267. databases = [{'name': 'foo', 'format': 'directory'}]
  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_parent_directory_for_dump')
  276. flexmock(module.dump).should_receive('create_named_pipe_for_dump').never()
  277. flexmock(module).should_receive('execute_command').with_args(
  278. (
  279. 'pg_dump',
  280. '--no-password',
  281. '--clean',
  282. '--if-exists',
  283. '--format',
  284. 'directory',
  285. '--file',
  286. 'databases/localhost/foo',
  287. 'foo',
  288. ),
  289. shell=True,
  290. extra_environment={'PGSSLMODE': 'disable'},
  291. ).and_return(flexmock()).once()
  292. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == []
  293. def test_dump_databases_runs_pg_dump_with_options():
  294. databases = [{'name': 'foo', 'options': '--stuff=such'}]
  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(('foo',))
  299. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  300. 'databases/localhost/foo'
  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. (
  306. 'pg_dump',
  307. '--no-password',
  308. '--clean',
  309. '--if-exists',
  310. '--format',
  311. 'custom',
  312. '--stuff=such',
  313. 'foo',
  314. '>',
  315. 'databases/localhost/foo',
  316. ),
  317. shell=True,
  318. extra_environment={'PGSSLMODE': 'disable'},
  319. run_to_completion=False,
  320. ).and_return(process).once()
  321. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
  322. def test_dump_databases_runs_pg_dumpall_for_all_databases():
  323. databases = [{'name': 'all'}]
  324. process = flexmock()
  325. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  326. flexmock(module).should_receive('make_dump_path').and_return('')
  327. flexmock(module).should_receive('database_names_to_dump').and_return(('all',))
  328. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  329. 'databases/localhost/all'
  330. )
  331. flexmock(module.os.path).should_receive('exists').and_return(False)
  332. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  333. flexmock(module).should_receive('execute_command').with_args(
  334. ('pg_dumpall', '--no-password', '--clean', '--if-exists', '>', 'databases/localhost/all'),
  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_dump_databases_runs_non_default_pg_dump():
  341. databases = [{'name': 'foo', 'pg_dump_command': 'special_pg_dump'}]
  342. process = flexmock()
  343. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  344. flexmock(module).should_receive('make_dump_path').and_return('')
  345. flexmock(module).should_receive('database_names_to_dump').and_return(('foo',))
  346. flexmock(module.dump).should_receive('make_database_dump_filename').and_return(
  347. 'databases/localhost/foo'
  348. )
  349. flexmock(module.os.path).should_receive('exists').and_return(False)
  350. flexmock(module.dump).should_receive('create_named_pipe_for_dump')
  351. flexmock(module).should_receive('execute_command').with_args(
  352. (
  353. 'special_pg_dump',
  354. '--no-password',
  355. '--clean',
  356. '--if-exists',
  357. '--format',
  358. 'custom',
  359. 'foo',
  360. '>',
  361. 'databases/localhost/foo',
  362. ),
  363. shell=True,
  364. extra_environment={'PGSSLMODE': 'disable'},
  365. run_to_completion=False,
  366. ).and_return(process).once()
  367. assert module.dump_databases(databases, 'test.yaml', {}, dry_run=False) == [process]
  368. def test_restore_database_dump_runs_pg_restore():
  369. database_config = [{'name': 'foo', 'schemas': None}]
  370. extract_process = flexmock(stdout=flexmock())
  371. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  372. flexmock(module).should_receive('make_dump_path')
  373. flexmock(module.dump).should_receive('make_database_dump_filename')
  374. flexmock(module).should_receive('execute_command_with_processes').with_args(
  375. (
  376. 'pg_restore',
  377. '--no-password',
  378. '--if-exists',
  379. '--exit-on-error',
  380. '--clean',
  381. '--dbname',
  382. 'foo',
  383. ),
  384. processes=[extract_process],
  385. output_log_level=logging.DEBUG,
  386. input_file=extract_process.stdout,
  387. extra_environment={'PGSSLMODE': 'disable'},
  388. ).once()
  389. flexmock(module).should_receive('execute_command').with_args(
  390. (
  391. 'psql',
  392. '--no-password',
  393. '--no-psqlrc',
  394. '--quiet',
  395. '--dbname',
  396. 'foo',
  397. '--command',
  398. 'ANALYZE',
  399. ),
  400. extra_environment={'PGSSLMODE': 'disable'},
  401. ).once()
  402. module.restore_database_dump(
  403. database_config,
  404. 'test.yaml',
  405. {},
  406. dry_run=False,
  407. extract_process=extract_process,
  408. connection_params={
  409. 'hostname': None,
  410. 'port': None,
  411. 'username': None,
  412. 'password': None,
  413. },
  414. )
  415. def test_restore_database_dump_errors_on_multiple_database_config():
  416. database_config = [{'name': 'foo'}, {'name': 'bar'}]
  417. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  418. flexmock(module).should_receive('make_dump_path')
  419. flexmock(module.dump).should_receive('make_database_dump_filename')
  420. flexmock(module).should_receive('execute_command_with_processes').never()
  421. flexmock(module).should_receive('execute_command').never()
  422. with pytest.raises(ValueError):
  423. module.restore_database_dump(
  424. database_config,
  425. 'test.yaml',
  426. {},
  427. dry_run=False,
  428. extract_process=flexmock(),
  429. connection_params={
  430. 'restore_hostname': None,
  431. 'restore_port': None,
  432. 'restore_username': None,
  433. 'restore_password': None,
  434. },
  435. )
  436. def test_restore_database_dump_runs_pg_restore_with_hostname_and_port():
  437. database_config = [
  438. {'name': 'foo', 'hostname': 'database.example.org', 'port': 5433, 'schemas': None}
  439. ]
  440. extract_process = flexmock(stdout=flexmock())
  441. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  442. flexmock(module).should_receive('make_dump_path')
  443. flexmock(module.dump).should_receive('make_database_dump_filename')
  444. flexmock(module).should_receive('execute_command_with_processes').with_args(
  445. (
  446. 'pg_restore',
  447. '--no-password',
  448. '--if-exists',
  449. '--exit-on-error',
  450. '--clean',
  451. '--dbname',
  452. 'foo',
  453. '--host',
  454. 'database.example.org',
  455. '--port',
  456. '5433',
  457. ),
  458. processes=[extract_process],
  459. output_log_level=logging.DEBUG,
  460. input_file=extract_process.stdout,
  461. extra_environment={'PGSSLMODE': 'disable'},
  462. ).once()
  463. flexmock(module).should_receive('execute_command').with_args(
  464. (
  465. 'psql',
  466. '--no-password',
  467. '--no-psqlrc',
  468. '--quiet',
  469. '--host',
  470. 'database.example.org',
  471. '--port',
  472. '5433',
  473. '--dbname',
  474. 'foo',
  475. '--command',
  476. 'ANALYZE',
  477. ),
  478. extra_environment={'PGSSLMODE': 'disable'},
  479. ).once()
  480. module.restore_database_dump(
  481. database_config,
  482. 'test.yaml',
  483. {},
  484. dry_run=False,
  485. extract_process=extract_process,
  486. connection_params={
  487. 'hostname': None,
  488. 'port': None,
  489. 'username': None,
  490. 'password': None,
  491. },
  492. )
  493. def test_restore_database_dump_runs_pg_restore_with_username_and_password():
  494. database_config = [
  495. {'name': 'foo', 'username': 'postgres', 'password': 'trustsome1', 'schemas': None}
  496. ]
  497. extract_process = flexmock(stdout=flexmock())
  498. flexmock(module).should_receive('make_extra_environment').and_return(
  499. {'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'}
  500. )
  501. flexmock(module).should_receive('make_dump_path')
  502. flexmock(module.dump).should_receive('make_database_dump_filename')
  503. flexmock(module).should_receive('execute_command_with_processes').with_args(
  504. (
  505. 'pg_restore',
  506. '--no-password',
  507. '--if-exists',
  508. '--exit-on-error',
  509. '--clean',
  510. '--dbname',
  511. 'foo',
  512. '--username',
  513. 'postgres',
  514. ),
  515. processes=[extract_process],
  516. output_log_level=logging.DEBUG,
  517. input_file=extract_process.stdout,
  518. extra_environment={'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'},
  519. ).once()
  520. flexmock(module).should_receive('execute_command').with_args(
  521. (
  522. 'psql',
  523. '--no-password',
  524. '--no-psqlrc',
  525. '--quiet',
  526. '--username',
  527. 'postgres',
  528. '--dbname',
  529. 'foo',
  530. '--command',
  531. 'ANALYZE',
  532. ),
  533. extra_environment={'PGPASSWORD': 'trustsome1', 'PGSSLMODE': 'disable'},
  534. ).once()
  535. module.restore_database_dump(
  536. database_config,
  537. 'test.yaml',
  538. {},
  539. dry_run=False,
  540. extract_process=extract_process,
  541. connection_params={
  542. 'hostname': None,
  543. 'port': None,
  544. 'username': None,
  545. 'password': None,
  546. },
  547. )
  548. def test_restore_database_dump_runs_pg_restore_with_options():
  549. database_config = [
  550. {
  551. 'name': 'foo',
  552. 'restore_options': '--harder',
  553. 'analyze_options': '--smarter',
  554. 'schemas': None,
  555. }
  556. ]
  557. extract_process = flexmock(stdout=flexmock())
  558. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  559. flexmock(module).should_receive('make_dump_path')
  560. flexmock(module.dump).should_receive('make_database_dump_filename')
  561. flexmock(module).should_receive('execute_command_with_processes').with_args(
  562. (
  563. 'pg_restore',
  564. '--no-password',
  565. '--if-exists',
  566. '--exit-on-error',
  567. '--clean',
  568. '--dbname',
  569. 'foo',
  570. '--harder',
  571. ),
  572. processes=[extract_process],
  573. output_log_level=logging.DEBUG,
  574. input_file=extract_process.stdout,
  575. extra_environment={'PGSSLMODE': 'disable'},
  576. ).once()
  577. flexmock(module).should_receive('execute_command').with_args(
  578. (
  579. 'psql',
  580. '--no-password',
  581. '--no-psqlrc',
  582. '--quiet',
  583. '--dbname',
  584. 'foo',
  585. '--smarter',
  586. '--command',
  587. 'ANALYZE',
  588. ),
  589. extra_environment={'PGSSLMODE': 'disable'},
  590. ).once()
  591. module.restore_database_dump(
  592. database_config,
  593. 'test.yaml',
  594. {},
  595. dry_run=False,
  596. extract_process=extract_process,
  597. connection_params={
  598. 'hostname': None,
  599. 'port': None,
  600. 'username': None,
  601. 'password': None,
  602. },
  603. )
  604. def test_restore_database_dump_runs_psql_for_all_database_dump():
  605. database_config = [{'name': 'all', 'schemas': None}]
  606. extract_process = flexmock(stdout=flexmock())
  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')
  610. flexmock(module).should_receive('execute_command_with_processes').with_args(
  611. (
  612. 'psql',
  613. '--no-password',
  614. '--no-psqlrc',
  615. ),
  616. processes=[extract_process],
  617. output_log_level=logging.DEBUG,
  618. input_file=extract_process.stdout,
  619. extra_environment={'PGSSLMODE': 'disable'},
  620. ).once()
  621. flexmock(module).should_receive('execute_command').with_args(
  622. ('psql', '--no-password', '--no-psqlrc', '--quiet', '--command', 'ANALYZE'),
  623. extra_environment={'PGSSLMODE': 'disable'},
  624. ).once()
  625. module.restore_database_dump(
  626. database_config,
  627. 'test.yaml',
  628. {},
  629. dry_run=False,
  630. extract_process=extract_process,
  631. connection_params={
  632. 'hostname': None,
  633. 'port': None,
  634. 'username': None,
  635. 'password': None,
  636. },
  637. )
  638. def test_restore_database_dump_runs_psql_for_plain_database_dump():
  639. database_config = [{'name': 'foo', 'format': 'plain', 'schemas': None}]
  640. extract_process = flexmock(stdout=flexmock())
  641. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  642. flexmock(module).should_receive('make_dump_path')
  643. flexmock(module.dump).should_receive('make_database_dump_filename')
  644. flexmock(module).should_receive('execute_command_with_processes').with_args(
  645. ('psql', '--no-password', '--no-psqlrc', '--dbname', 'foo'),
  646. processes=[extract_process],
  647. output_log_level=logging.DEBUG,
  648. input_file=extract_process.stdout,
  649. extra_environment={'PGSSLMODE': 'disable'},
  650. ).once()
  651. flexmock(module).should_receive('execute_command').with_args(
  652. (
  653. 'psql',
  654. '--no-password',
  655. '--no-psqlrc',
  656. '--quiet',
  657. '--dbname',
  658. 'foo',
  659. '--command',
  660. 'ANALYZE',
  661. ),
  662. extra_environment={'PGSSLMODE': 'disable'},
  663. ).once()
  664. module.restore_database_dump(
  665. database_config,
  666. 'test.yaml',
  667. {},
  668. dry_run=False,
  669. extract_process=extract_process,
  670. connection_params={
  671. 'hostname': None,
  672. 'port': None,
  673. 'username': None,
  674. 'password': None,
  675. },
  676. )
  677. def test_restore_database_dump_runs_non_default_pg_restore_and_psql():
  678. database_config = [
  679. {
  680. 'name': 'foo',
  681. 'pg_restore_command': 'docker exec mycontainer pg_restore',
  682. 'psql_command': 'docker exec mycontainer psql',
  683. 'schemas': None,
  684. }
  685. ]
  686. extract_process = flexmock(stdout=flexmock())
  687. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  688. flexmock(module).should_receive('make_dump_path')
  689. flexmock(module.dump).should_receive('make_database_dump_filename')
  690. flexmock(module).should_receive('execute_command_with_processes').with_args(
  691. (
  692. 'docker',
  693. 'exec',
  694. 'mycontainer',
  695. 'pg_restore',
  696. '--no-password',
  697. '--if-exists',
  698. '--exit-on-error',
  699. '--clean',
  700. '--dbname',
  701. 'foo',
  702. ),
  703. processes=[extract_process],
  704. output_log_level=logging.DEBUG,
  705. input_file=extract_process.stdout,
  706. extra_environment={'PGSSLMODE': 'disable'},
  707. ).once()
  708. flexmock(module).should_receive('execute_command').with_args(
  709. (
  710. 'docker',
  711. 'exec',
  712. 'mycontainer',
  713. 'psql',
  714. '--no-password',
  715. '--no-psqlrc',
  716. '--quiet',
  717. '--dbname',
  718. 'foo',
  719. '--command',
  720. 'ANALYZE',
  721. ),
  722. extra_environment={'PGSSLMODE': 'disable'},
  723. ).once()
  724. module.restore_database_dump(
  725. database_config,
  726. 'test.yaml',
  727. {},
  728. dry_run=False,
  729. extract_process=extract_process,
  730. connection_params={
  731. 'hostname': None,
  732. 'port': None,
  733. 'username': None,
  734. 'password': None,
  735. },
  736. )
  737. def test_restore_database_dump_with_dry_run_skips_restore():
  738. database_config = [{'name': 'foo', 'schemas': None}]
  739. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  740. flexmock(module).should_receive('make_dump_path')
  741. flexmock(module.dump).should_receive('make_database_dump_filename')
  742. flexmock(module).should_receive('execute_command_with_processes').never()
  743. module.restore_database_dump(
  744. database_config,
  745. 'test.yaml',
  746. {},
  747. dry_run=True,
  748. extract_process=flexmock(),
  749. connection_params={
  750. 'hostname': None,
  751. 'port': None,
  752. 'username': None,
  753. 'password': None,
  754. },
  755. )
  756. def test_restore_database_dump_without_extract_process_restores_from_disk():
  757. database_config = [{'name': 'foo', 'schemas': None}]
  758. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  759. flexmock(module).should_receive('make_dump_path')
  760. flexmock(module.dump).should_receive('make_database_dump_filename').and_return('/dump/path')
  761. flexmock(module).should_receive('execute_command_with_processes').with_args(
  762. (
  763. 'pg_restore',
  764. '--no-password',
  765. '--if-exists',
  766. '--exit-on-error',
  767. '--clean',
  768. '--dbname',
  769. 'foo',
  770. '/dump/path',
  771. ),
  772. processes=[],
  773. output_log_level=logging.DEBUG,
  774. input_file=None,
  775. extra_environment={'PGSSLMODE': 'disable'},
  776. ).once()
  777. flexmock(module).should_receive('execute_command').with_args(
  778. (
  779. 'psql',
  780. '--no-password',
  781. '--no-psqlrc',
  782. '--quiet',
  783. '--dbname',
  784. 'foo',
  785. '--command',
  786. 'ANALYZE',
  787. ),
  788. extra_environment={'PGSSLMODE': 'disable'},
  789. ).once()
  790. module.restore_database_dump(
  791. database_config,
  792. 'test.yaml',
  793. {},
  794. dry_run=False,
  795. extract_process=None,
  796. connection_params={
  797. 'hostname': None,
  798. 'port': None,
  799. 'username': None,
  800. 'password': None,
  801. },
  802. )
  803. def test_restore_database_dump_with_schemas_restores_schemas():
  804. database_config = [{'name': 'foo', 'schemas': ['bar', 'baz']}]
  805. flexmock(module).should_receive('make_extra_environment').and_return({'PGSSLMODE': 'disable'})
  806. flexmock(module).should_receive('make_dump_path')
  807. flexmock(module.dump).should_receive('make_database_dump_filename').and_return('/dump/path')
  808. flexmock(module).should_receive('execute_command_with_processes').with_args(
  809. (
  810. 'pg_restore',
  811. '--no-password',
  812. '--if-exists',
  813. '--exit-on-error',
  814. '--clean',
  815. '--dbname',
  816. 'foo',
  817. '/dump/path',
  818. '--schema',
  819. 'bar',
  820. '--schema',
  821. 'baz',
  822. ),
  823. processes=[],
  824. output_log_level=logging.DEBUG,
  825. input_file=None,
  826. extra_environment={'PGSSLMODE': 'disable'},
  827. ).once()
  828. flexmock(module).should_receive('execute_command').with_args(
  829. (
  830. 'psql',
  831. '--no-password',
  832. '--no-psqlrc',
  833. '--quiet',
  834. '--dbname',
  835. 'foo',
  836. '--command',
  837. 'ANALYZE',
  838. ),
  839. extra_environment={'PGSSLMODE': 'disable'},
  840. ).once()
  841. module.restore_database_dump(
  842. database_config,
  843. 'test.yaml',
  844. {},
  845. dry_run=False,
  846. extract_process=None,
  847. connection_params={
  848. 'hostname': None,
  849. 'port': None,
  850. 'username': None,
  851. 'password': None,
  852. },
  853. )