From 34ff4433fdc5041c19e43fbbc25b51c4f11592e0 Mon Sep 17 00:00:00 2001 From: jross Date: Thu, 19 Mar 2026 10:26:41 -0600 Subject: [PATCH 1/2] Handle pg_dump restrict commands in local restore --- cli/db_restore.py | 15 +++++++++++++-- tests/test_cli_commands.py | 8 ++++++-- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/cli/db_restore.py b/cli/db_restore.py index c746a18e..5290452e 100644 --- a/cli/db_restore.py +++ b/cli/db_restore.py @@ -21,6 +21,12 @@ re.compile(r"^\s*REVOKE\b", re.IGNORECASE), re.compile(r"^\s*ALTER\s+DEFAULT\s+PRIVILEGES\b", re.IGNORECASE), ) +PSQL_META_COMMAND_PATTERNS = ( + # Newer pg_dump versions emit these psql-only commands for safer restores. + # Older local psql clients reject them, so drop them from staged restores. + re.compile(r"^\s*\\restrict\b", re.IGNORECASE), + re.compile(r"^\s*\\unrestrict\b", re.IGNORECASE), +) class LocalDbRestoreError(RuntimeError): @@ -81,9 +87,13 @@ def _sanitize_sql_dump(source_path: Path, target_path: Path) -> None: with open(source_path, "r", encoding="utf-8") as infile: with open(target_path, "w", encoding="utf-8") as outfile: for line in infile: - if any( + matches_role_sql = any( pattern.search(line) for pattern in ROLE_DEPENDENT_SQL_PATTERNS - ): + ) + matches_psql_meta = any( + pattern.search(line) for pattern in PSQL_META_COMMAND_PATTERNS + ) + if matches_role_sql or matches_psql_meta: continue outfile.write(line) except UnicodeError as exc: @@ -235,6 +245,7 @@ def restore_local_db_from_sql( ) from exc return LocalDbRestoreResult( + sql_file=staged_sql_file, source=source_description, host=host, port=port, diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py index 47f451ec..6714098e 100644 --- a/tests/test_cli_commands.py +++ b/tests/test_cli_commands.py @@ -164,10 +164,12 @@ def fake_associate(source_directory): def test_restore_local_db_invokes_psql(monkeypatch, tmp_path): sql_file = tmp_path / "restore.sql" sql_file.write_text( + "\\restrict abc123\n" "SET ROLE ocotillo;\n" "ALTER TABLE public.sample OWNER TO ocotillo;\n" "GRANT ALL ON TABLE public.sample TO ocotillo;\n" "select 1;\n" + "\\unrestrict abc123\n" ) captured: dict[str, object] = {} call_order: list[str] = [] @@ -532,10 +534,12 @@ def test_water_levels_cli_persists_observations(tmp_path, water_well_thing): """ def _write_csv(path: Path, *, well_name: str, notes: str): - csv_text = textwrap.dedent(f"""\ + csv_text = textwrap.dedent( + f"""\ field_staff,well_name_point_id,field_event_date_time,measurement_date_time,sampler,sample_method,mp_height,level_status,depth_to_water_ft,data_quality,water_level_notes CLI Tester,{well_name},2025-02-15T08:00:00-07:00,2025-02-15T10:30:00-07:00,Groundwater Team,electric tape,1.5,stable,42.5,approved,{notes} - """) + """ + ) path.write_text(csv_text) unique_notes = f"pytest-{uuid.uuid4()}" From 20aafdd21dfc127a01485ccc76c9c9a3c5146bb6 Mon Sep 17 00:00:00 2001 From: jirhiker <2035568+jirhiker@users.noreply.github.com> Date: Thu, 19 Mar 2026 16:27:11 +0000 Subject: [PATCH 2/2] Formatting changes --- tests/test_cli_commands.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py index 6714098e..53d136d5 100644 --- a/tests/test_cli_commands.py +++ b/tests/test_cli_commands.py @@ -534,12 +534,10 @@ def test_water_levels_cli_persists_observations(tmp_path, water_well_thing): """ def _write_csv(path: Path, *, well_name: str, notes: str): - csv_text = textwrap.dedent( - f"""\ + csv_text = textwrap.dedent(f"""\ field_staff,well_name_point_id,field_event_date_time,measurement_date_time,sampler,sample_method,mp_height,level_status,depth_to_water_ft,data_quality,water_level_notes CLI Tester,{well_name},2025-02-15T08:00:00-07:00,2025-02-15T10:30:00-07:00,Groundwater Team,electric tape,1.5,stable,42.5,approved,{notes} - """ - ) + """) path.write_text(csv_text) unique_notes = f"pytest-{uuid.uuid4()}"