@@ -355,9 +355,10 @@ def delete_user(self, user: str) -> None:
355
355
# Existing objects need to be reassigned in each database
356
356
# before the user can be deleted.
357
357
for database in databases :
358
- with self ._connect_to_database (
359
- database
360
- ) as connection , connection .cursor () as cursor :
358
+ with (
359
+ self ._connect_to_database (database ) as connection ,
360
+ connection .cursor () as cursor ,
361
+ ):
361
362
cursor .execute (
362
363
SQL ("REASSIGN OWNED BY {} TO {};" ).format (
363
364
Identifier (user ), Identifier (self .user )
@@ -448,9 +449,10 @@ def enable_disable_extensions(
448
449
449
450
# Enable/disabled the extension in each database.
450
451
for database in databases :
451
- with self ._connect_to_database (
452
- database = database
453
- ) as connection , connection .cursor () as cursor :
452
+ with (
453
+ self ._connect_to_database (database = database ) as connection ,
454
+ connection .cursor () as cursor ,
455
+ ):
454
456
for extension , enable in ordered_extensions .items ():
455
457
cursor .execute (
456
458
f"CREATE EXTENSION IF NOT EXISTS { extension } ;"
@@ -560,9 +562,10 @@ def get_postgresql_text_search_configs(self) -> Set[str]:
560
562
Returns:
561
563
Set of PostgreSQL text search configs.
562
564
"""
563
- with self ._connect_to_database (
564
- database_host = self .current_host
565
- ) as connection , connection .cursor () as cursor :
565
+ with (
566
+ self ._connect_to_database (database_host = self .current_host ) as connection ,
567
+ connection .cursor () as cursor ,
568
+ ):
566
569
cursor .execute ("SELECT CONCAT('pg_catalog.', cfgname) FROM pg_ts_config;" )
567
570
text_search_configs = cursor .fetchall ()
568
571
return {text_search_config [0 ] for text_search_config in text_search_configs }
@@ -573,9 +576,10 @@ def get_postgresql_timezones(self) -> Set[str]:
573
576
Returns:
574
577
Set of PostgreSQL timezones.
575
578
"""
576
- with self ._connect_to_database (
577
- database_host = self .current_host
578
- ) as connection , connection .cursor () as cursor :
579
+ with (
580
+ self ._connect_to_database (database_host = self .current_host ) as connection ,
581
+ connection .cursor () as cursor ,
582
+ ):
579
583
cursor .execute ("SELECT name FROM pg_timezone_names;" )
580
584
timezones = cursor .fetchall ()
581
585
return {timezone [0 ] for timezone in timezones }
@@ -586,9 +590,10 @@ def get_postgresql_default_table_access_methods(self) -> Set[str]:
586
590
Returns:
587
591
Set of PostgreSQL table access methods.
588
592
"""
589
- with self ._connect_to_database (
590
- database_host = self .current_host
591
- ) as connection , connection .cursor () as cursor :
593
+ with (
594
+ self ._connect_to_database (database_host = self .current_host ) as connection ,
595
+ connection .cursor () as cursor ,
596
+ ):
592
597
cursor .execute ("SELECT amname FROM pg_am WHERE amtype = 't';" )
593
598
access_methods = cursor .fetchall ()
594
599
return {access_method [0 ] for access_method in access_methods }
@@ -601,9 +606,10 @@ def get_postgresql_version(self, current_host=True) -> str:
601
606
"""
602
607
host = self .current_host if current_host else None
603
608
try :
604
- with self ._connect_to_database (
605
- database_host = host
606
- ) as connection , connection .cursor () as cursor :
609
+ with (
610
+ self ._connect_to_database (database_host = host ) as connection ,
611
+ connection .cursor () as cursor ,
612
+ ):
607
613
cursor .execute ("SELECT version();" )
608
614
# Split to get only the version number.
609
615
return cursor .fetchone ()[0 ].split (" " )[1 ]
@@ -622,9 +628,12 @@ def is_tls_enabled(self, check_current_host: bool = False) -> bool:
622
628
whether TLS is enabled.
623
629
"""
624
630
try :
625
- with self ._connect_to_database (
626
- database_host = self .current_host if check_current_host else None
627
- ) as connection , connection .cursor () as cursor :
631
+ with (
632
+ self ._connect_to_database (
633
+ database_host = self .current_host if check_current_host else None
634
+ ) as connection ,
635
+ connection .cursor () as cursor ,
636
+ ):
628
637
cursor .execute ("SHOW ssl;" )
629
638
return "on" in cursor .fetchone ()[0 ]
630
639
except psycopg2 .Error :
@@ -644,9 +653,10 @@ def list_access_groups(self, current_host=False) -> Set[str]:
644
653
connection = None
645
654
host = self .current_host if current_host else None
646
655
try :
647
- with self ._connect_to_database (
648
- database_host = host
649
- ) as connection , connection .cursor () as cursor :
656
+ with (
657
+ self ._connect_to_database (database_host = host ) as connection ,
658
+ connection .cursor () as cursor ,
659
+ ):
650
660
cursor .execute (
651
661
"SELECT groname FROM pg_catalog.pg_group WHERE groname LIKE '%_access';"
652
662
)
@@ -674,9 +684,10 @@ def list_accessible_databases_for_user(self, user: str, current_host=False) -> S
674
684
connection = None
675
685
host = self .current_host if current_host else None
676
686
try :
677
- with self ._connect_to_database (
678
- database_host = host
679
- ) as connection , connection .cursor () as cursor :
687
+ with (
688
+ self ._connect_to_database (database_host = host ) as connection ,
689
+ connection .cursor () as cursor ,
690
+ ):
680
691
cursor .execute (
681
692
SQL (
682
693
"SELECT TRUE FROM pg_catalog.pg_user WHERE usename = {} AND usesuper;"
@@ -712,9 +723,10 @@ def list_users(self, group: Optional[str] = None, current_host=False) -> Set[str
712
723
connection = None
713
724
host = self .current_host if current_host else None
714
725
try :
715
- with self ._connect_to_database (
716
- database_host = host
717
- ) as connection , connection .cursor () as cursor :
726
+ with (
727
+ self ._connect_to_database (database_host = host ) as connection ,
728
+ connection .cursor () as cursor ,
729
+ ):
718
730
if group :
719
731
query = SQL (
720
732
"SELECT usename FROM (SELECT UNNEST(grolist) AS user_id FROM pg_catalog.pg_group WHERE groname = {}) AS g JOIN pg_catalog.pg_user AS u ON g.user_id = u.usesysid;"
@@ -744,9 +756,10 @@ def list_users_from_relation(self, current_host=False) -> Set[str]:
744
756
connection = None
745
757
host = self .current_host if current_host else None
746
758
try :
747
- with self ._connect_to_database (
748
- database_host = host
749
- ) as connection , connection .cursor () as cursor :
759
+ with (
760
+ self ._connect_to_database (database_host = host ) as connection ,
761
+ connection .cursor () as cursor ,
762
+ ):
750
763
cursor .execute (
751
764
"SELECT usename "
752
765
"FROM pg_catalog.pg_user "
@@ -782,9 +795,10 @@ def set_up_database(self, temp_location: Optional[str] = None) -> None:
782
795
connection = None
783
796
cursor = None
784
797
try :
785
- with self ._connect_to_database (
786
- database = "template1"
787
- ) as connection , connection .cursor () as cursor :
798
+ with (
799
+ self ._connect_to_database (database = "template1" ) as connection ,
800
+ connection .cursor () as cursor ,
801
+ ):
788
802
# Create database function and event trigger to identify users created by PgBouncer.
789
803
cursor .execute (
790
804
"SELECT TRUE FROM pg_event_trigger WHERE evtname = 'update_pg_hba_on_create_schema';"
@@ -900,9 +914,10 @@ def update_user_password(
900
914
"""
901
915
connection = None
902
916
try :
903
- with self ._connect_to_database (
904
- database_host = database_host
905
- ) as connection , connection .cursor () as cursor :
917
+ with (
918
+ self ._connect_to_database (database_host = database_host ) as connection ,
919
+ connection .cursor () as cursor ,
920
+ ):
906
921
cursor .execute (SQL ("BEGIN;" ))
907
922
cursor .execute (SQL ("SET LOCAL log_statement = 'none';" ))
908
923
cursor .execute (
@@ -1039,9 +1054,10 @@ def validate_date_style(self, date_style: str) -> bool:
1039
1054
Whether the date style is valid.
1040
1055
"""
1041
1056
try :
1042
- with self ._connect_to_database (
1043
- database_host = self .current_host
1044
- ) as connection , connection .cursor () as cursor :
1057
+ with (
1058
+ self ._connect_to_database (database_host = self .current_host ) as connection ,
1059
+ connection .cursor () as cursor ,
1060
+ ):
1045
1061
cursor .execute (
1046
1062
SQL (
1047
1063
"SET DateStyle to {};" ,
0 commit comments