storcon: squash all migrations into one

Problem

Neon and Hadron deployments have the same database schema, but different
migration histories. Some transactions have different identifiers too.
If we don't do anything about it, then the storage controller would fail
to apply the merged set of transactions.

Summary of Changes

We squash all migrations into a single one. If the schema already
matches, then the new transaction gets applie without doing anything.
For new regions, this migration will bootstrap the database schema.

This should be merged in both neon and hadron codebases.

Note that after deploying this change, the `__diesel_schema_migrations`
table will still contain entries for the old pre-squash transactions.
This is fine because diesel only considers the tranasactions embedded
in the repo for application. Once we are certain that we are not going
to roll back, we can clean up the `__diesel_schema_migrations` tables in
prod, but this is manual and error prone, so I'd skip it.

Rolling back

Rolling back to a previous deployment which embeds the non-squashed
transactions is safe. The old transactions are still present in
`__diesel_schema_migrations` (i.e. considered applied), so no migrations
will be run, so no migrations will be run.

Note that this assumes that all transactions are applied and squashed into
the new migration before deployment.

compute_tools/src/sql/anon_ext_fn_reassign.sql

@@ -1,13 +0,0 @@
-DO $$
-DECLARE
-    query varchar;
-BEGIN
-    FOR query IN
-    SELECT pg_catalog.format('ALTER FUNCTION %I(%s) OWNER TO {db_owner};', p.oid::regproc, pg_catalog.pg_get_function_identity_arguments(p.oid))
-    FROM pg_catalog.pg_proc p
-        WHERE p.pronamespace OPERATOR(pg_catalog.=) 'anon'::regnamespace::oid
-    LOOP
-        EXECUTE query;
-    END LOOP;
-END
-$$;

storage_controller/migrations/2024-01-07-211257_create_tenant_shards/down.sql

@@ -1 +0,0 @@
-DROP TABLE tenant_shards;

storage_controller/migrations/2024-01-07-211257_create_tenant_shards/up.sql

@@ -1,13 +0,0 @@
-CREATE TABLE tenant_shards (
-  tenant_id VARCHAR NOT NULL,
-  shard_number INTEGER NOT NULL,
-  shard_count INTEGER NOT NULL,
-  PRIMARY KEY(tenant_id, shard_number, shard_count),
-  shard_stripe_size INTEGER NOT NULL,
-  generation INTEGER NOT NULL,
-  generation_pageserver BIGINT NOT NULL,
-  placement_policy VARCHAR NOT NULL,
-  splitting SMALLINT NOT NULL,
-  -- config is JSON encoded, opaque to the database.
-  config TEXT NOT NULL
-);

storage_controller/migrations/2024-01-07-212945_create_nodes/down.sql

@@ -1 +0,0 @@
-DROP TABLE nodes;

storage_controller/migrations/2024-01-07-212945_create_nodes/up.sql

@@ -1,10 +0,0 @@
-CREATE TABLE nodes (
-  node_id BIGINT PRIMARY KEY NOT NULL,
-
-  scheduling_policy VARCHAR NOT NULL,
-
-  listen_http_addr VARCHAR NOT NULL,
-  listen_http_port INTEGER NOT NULL,
-  listen_pg_addr VARCHAR NOT NULL,
-  listen_pg_port INTEGER NOT NULL
-);

storage_controller/migrations/2024-02-29-094122_generations_null/down.sql

@@ -1,2 +0,0 @@
-ALTER TABLE tenant_shards ALTER generation SET NOT NULL;
-ALTER TABLE tenant_shards ALTER generation_pageserver SET NOT NULL;

storage_controller/migrations/2024-02-29-094122_generations_null/up.sql

@@ -1,4 +0,0 @@
-
-
-ALTER TABLE tenant_shards ALTER generation DROP NOT NULL;
-ALTER TABLE tenant_shards ALTER generation_pageserver DROP NOT NULL;

storage_controller/migrations/2024-03-18-184429_rename_policy/down.sql

@@ -1,3 +0,0 @@
-
-UPDATE tenant_shards set placement_policy='{"Double": 1}' where placement_policy='{"Attached": 1}';
-UPDATE tenant_shards set placement_policy='"Single"' where placement_policy='{"Attached": 0}';

storage_controller/migrations/2024-03-18-184429_rename_policy/up.sql

@@ -1,3 +0,0 @@
-
-UPDATE tenant_shards set placement_policy='{"Attached": 1}' where placement_policy='{"Double": 1}';
-UPDATE tenant_shards set placement_policy='{"Attached": 0}' where placement_policy='"Single"';

storage_controller/migrations/2024-03-27-133204_tenant_policies/down.sql

@@ -1,3 +0,0 @@
--- This file should undo anything in `up.sql`
-
-ALTER TABLE tenant_shards drop scheduling_policy;

storage_controller/migrations/2024-03-27-133204_tenant_policies/up.sql

@@ -1,2 +0,0 @@
-
-ALTER TABLE tenant_shards add scheduling_policy VARCHAR NOT NULL DEFAULT '"Active"';

storage_controller/migrations/2024-07-23-191537_create_metadata_health/down.sql

@@ -1 +0,0 @@
-DROP TABLE metadata_health;

storage_controller/migrations/2024-07-23-191537_create_metadata_health/up.sql

@@ -1,14 +0,0 @@
-CREATE TABLE metadata_health (
-  tenant_id VARCHAR NOT NULL,
-  shard_number INTEGER NOT NULL,
-  shard_count INTEGER NOT NULL,
-  PRIMARY KEY(tenant_id, shard_number, shard_count),
-  -- Rely on cascade behavior for delete
-  FOREIGN KEY(tenant_id, shard_number, shard_count) REFERENCES tenant_shards ON DELETE CASCADE,
-  healthy BOOLEAN NOT NULL DEFAULT TRUE,
-  last_scrubbed_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
-);
-
-
-INSERT INTO metadata_health(tenant_id, shard_number, shard_count)
-SELECT tenant_id, shard_number, shard_count FROM tenant_shards;

storage_controller/migrations/2024-07-26-140924_create_leader/down.sql

@@ -1 +0,0 @@
-DROP TABLE controllers;

storage_controller/migrations/2024-07-26-140924_create_leader/up.sql

@@ -1,5 +0,0 @@
-CREATE TABLE controllers (
-  address VARCHAR NOT NULL,
-  started_at TIMESTAMPTZ NOT NULL,
-  PRIMARY KEY(address, started_at)
-);

storage_controller/migrations/2024-08-23-102952_safekeepers/down.sql

@@ -1,2 +0,0 @@
--- This file should undo anything in `up.sql`
-DROP TABLE safekeepers;

storage_controller/migrations/2024-08-23-102952_safekeepers/up.sql

@@ -1,15 +0,0 @@
--- started out as a copy of cplane schema, removed the unnecessary columns.
-CREATE TABLE safekeepers (
-	-- the surrogate identifier defined by control plane database sequence
-	id BIGINT PRIMARY KEY,
-	region_id TEXT NOT NULL,
-	version BIGINT NOT NULL,
-	-- the natural id on whatever cloud platform, not needed in storage controller
-	-- instance_id TEXT UNIQUE NOT NULL,
-	host TEXT NOT NULL,
-	port INTEGER NOT NULL,
-	active BOOLEAN NOT NULL DEFAULT false,
-	-- projects_count INTEGER NOT NULL DEFAULT 0,
-	http_port INTEGER NOT NULL,
-	availability_zone_id TEXT NOT NULL
-);

storage_controller/migrations/2024-08-23-170149_tenant_id_index/down.sql

@@ -1,2 +0,0 @@
--- This file should undo anything in `up.sql`
-DROP INDEX tenant_shards_tenant_id;

storage_controller/migrations/2024-08-23-170149_tenant_id_index/up.sql

@@ -1,2 +0,0 @@
--- Your SQL goes here
-CREATE INDEX tenant_shards_tenant_id ON tenant_shards (tenant_id);

storage_controller/migrations/2024-08-27-184400_pageserver_az/down.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes DROP availability_zone_id;

storage_controller/migrations/2024-08-27-184400_pageserver_az/up.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes ADD availability_zone_id VARCHAR;

storage_controller/migrations/2024-08-28-150530_pageserver_az_not_null/down.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes ALTER availability_zone_id DROP NOT NULL;

storage_controller/migrations/2024-08-28-150530_pageserver_az_not_null/up.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes ALTER availability_zone_id SET NOT NULL;

storage_controller/migrations/2024-09-05-104500_tenant_shard_preferred_az/down.sql

@@ -1 +0,0 @@
-ALTER TABLE tenant_shards DROP preferred_az_id;

storage_controller/migrations/2024-09-05-104500_tenant_shard_preferred_az/up.sql

@@ -1 +0,0 @@
-ALTER TABLE tenant_shards ADD preferred_az_id VARCHAR;

storage_controller/migrations/2024-12-12-212515_safekeepers_scheduling_policy/down.sql

@@ -1 +0,0 @@
-ALTER TABLE safekeepers DROP scheduling_policy;

storage_controller/migrations/2024-12-12-212515_safekeepers_scheduling_policy/up.sql

@@ -1 +0,0 @@
-ALTER TABLE safekeepers ADD scheduling_policy VARCHAR NOT NULL DEFAULT 'disabled';

storage_controller/migrations/2025-01-09-160454_safekeepers_remove_active/down.sql

@@ -1,4 +0,0 @@
--- this sadly isn't a "true" revert of the migration, as the column is now at the end of the table.
--- But preserving order is not a trivial operation.
--- https://wiki.postgresql.org/wiki/Alter_column_position
-ALTER TABLE safekeepers ADD active BOOLEAN NOT NULL DEFAULT false;

storage_controller/migrations/2025-01-09-160454_safekeepers_remove_active/up.sql

@@ -1 +0,0 @@
-ALTER TABLE safekeepers DROP active;

storage_controller/migrations/2025-01-15-181207_safekeepers_disabled_to_pause/down.sql

@@ -1,2 +0,0 @@
-ALTER TABLE safekeepers ALTER COLUMN scheduling_policy SET DEFAULT 'disabled';
-UPDATE safekeepers SET scheduling_policy = 'disabled' WHERE scheduling_policy = 'pause';

storage_controller/migrations/2025-01-15-181207_safekeepers_disabled_to_pause/up.sql

@@ -1,2 +0,0 @@
-ALTER TABLE safekeepers ALTER COLUMN scheduling_policy SET DEFAULT 'pause';
-UPDATE safekeepers SET scheduling_policy = 'pause' WHERE scheduling_policy = 'disabled';

storage_controller/migrations/2025-02-11-144848_pageserver_use_https/down.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes DROP listen_https_port;

storage_controller/migrations/2025-02-11-144848_pageserver_use_https/up.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes ADD listen_https_port INTEGER;

storage_controller/migrations/2025-02-14-160526_safekeeper_timelines/down.sql

@@ -1,2 +0,0 @@
-DROP TABLE timelines;
-DROP TABLE safekeeper_timeline_pending_ops;

storage_controller/migrations/2025-02-14-160526_safekeeper_timelines/up.sql

@@ -1,19 +0,0 @@
-CREATE TABLE timelines (
-  tenant_id VARCHAR NOT NULL,
-  timeline_id VARCHAR NOT NULL,
-  start_lsn pg_lsn NOT NULL,
-  generation INTEGER NOT NULL,
-  sk_set BIGINT[] NOT NULL,
-  new_sk_set BIGINT[],
-  cplane_notified_generation INTEGER NOT NULL,
-  deleted_at timestamptz,
-  PRIMARY KEY(tenant_id, timeline_id)
-);
-CREATE TABLE safekeeper_timeline_pending_ops (
-  sk_id BIGINT NOT NULL,
-  tenant_id VARCHAR NOT NULL,
-  timeline_id VARCHAR NOT NULL,
-  generation INTEGER NOT NULL,
-  op_kind VARCHAR NOT NULL,
-  PRIMARY KEY(tenant_id, timeline_id, sk_id)
-);

storage_controller/migrations/2025-02-28-141741_safekeeper_use_https/down.sql

@@ -1 +0,0 @@
-ALTER TABLE safekeepers DROP https_port;

storage_controller/migrations/2025-02-28-141741_safekeeper_use_https/up.sql

@@ -1 +0,0 @@
-ALTER TABLE safekeepers ADD https_port INTEGER;

storage_controller/migrations/2025-03-18-103700_timeline_imports/down.sql

@@ -1 +0,0 @@
-DROP TABLE timeline_imports;

storage_controller/migrations/2025-03-18-103700_timeline_imports/up.sql

@@ -1,6 +0,0 @@
-CREATE TABLE timeline_imports (
-  tenant_id VARCHAR NOT NULL,
-  timeline_id VARCHAR NOT NULL,
-  shard_statuses JSONB NOT NULL,
-  PRIMARY KEY(tenant_id, timeline_id)
-);

storage_controller/migrations/2025-06-01-201442_add_lifecycle_to_nodes/down.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes DROP COLUMN lifecycle;

storage_controller/migrations/2025-06-01-201442_add_lifecycle_to_nodes/up.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes ADD COLUMN lifecycle VARCHAR NOT NULL DEFAULT 'active';

storage_controller/migrations/2025-06-17-082247_pageserver_grpc_addr/down.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes DROP listen_grpc_addr, listen_grpc_port;

storage_controller/migrations/2025-06-17-082247_pageserver_grpc_addr/up.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes ADD listen_grpc_addr VARCHAR NULL, ADD listen_grpc_port INTEGER NULL;

storage_controller/migrations/2025-07-02-170751_safekeeper_default_no_pause/down.sql

@@ -1 +0,0 @@
-ALTER TABLE safekeepers ALTER COLUMN scheduling_policy SET DEFAULT 'pause';

storage_controller/migrations/2025-07-02-170751_safekeeper_default_no_pause/up.sql

@@ -1 +0,0 @@
-ALTER TABLE safekeepers ALTER COLUMN scheduling_policy SET DEFAULT 'activating';

storage_controller/migrations/2025-07-08-114340_sk_set_notified_generation/down.sql

@@ -1 +0,0 @@
-ALTER TABLE timelines DROP sk_set_notified_generation;

storage_controller/migrations/2025-07-08-114340_sk_set_notified_generation/up.sql

@@ -1 +0,0 @@
-ALTER TABLE timelines ADD sk_set_notified_generation INTEGER NOT NULL DEFAULT 1;

storage_controller/migrations/2025-07-17-000001_hadron_safekeepers/down.sql

@@ -1,2 +0,0 @@
-DROP TABLE hadron_safekeepers;
-DROP TABLE hadron_timeline_safekeepers;

storage_controller/migrations/2025-07-17-000001_hadron_safekeepers/up.sql

@@ -1,17 +0,0 @@
--- hadron_safekeepers keep track of all Safe Keeper nodes that exist in the system.
--- Upon startup, each Safe Keeper reaches out to the hadron cluster coordinator to register its node ID and listen addresses.
-
-CREATE TABLE hadron_safekeepers (
-  sk_node_id BIGINT PRIMARY KEY NOT NULL,
-  listen_http_addr VARCHAR NOT NULL,
-  listen_http_port INTEGER NOT NULL,
-  listen_pg_addr VARCHAR NOT NULL,
-  listen_pg_port INTEGER NOT NULL
-);
-
-CREATE TABLE hadron_timeline_safekeepers (
-  timeline_id VARCHAR NOT NULL,
-  sk_node_id BIGINT NOT NULL,
-  legacy_endpoint_id UUID DEFAULT NULL,
-  PRIMARY KEY(timeline_id, sk_node_id)
-);

storage_controller/migrations/2025-07-31-115100_squash_migrations/down.sql

@@ -0,0 +1,10 @@
+DROP TABLE IF EXISTS timelines CASCADE;
+DROP TABLE IF EXISTS timeline_imports CASCADE;
+DROP TABLE IF EXISTS tenant_shards CASCADE;
+DROP TABLE IF EXISTS safekeepers CASCADE;
+DROP TABLE IF EXISTS safekeeper_timeline_pending_ops CASCADE;
+DROP TABLE IF EXISTS nodes CASCADE;
+DROP TABLE IF EXISTS metadata_health CASCADE;
+DROP TABLE IF EXISTS hadron_timeline_safekeepers CASCADE;
+DROP TABLE IF EXISTS hadron_safekeepers CASCADE;
+DROP TABLE IF EXISTS controllers CASCADE;

storage_controller/migrations/2025-07-31-115100_squash_migrations/up.sql

@@ -0,0 +1,112 @@
+/*
+ * This migration squashes all previous migrations into a single one.
+ * It will be applied in two cases:
+ * 1. Spinning up a new region. There's no previous migrations apart from the diesel setup in this case
+ * 2. For existing regions.
+ *
+ * Hence, this migration does nothing if the schema is already up to date.
+ */
+
+CREATE TABLE IF NOT EXISTS controllers (
+    address character varying NOT NULL,
+    started_at timestamp with time zone NOT NULL,
+    PRIMARY KEY(address, started_at)
+);
+
+CREATE TABLE IF NOT EXISTS hadron_safekeepers (
+    sk_node_id bigint PRIMARY KEY NOT NULL,
+    listen_http_addr character varying NOT NULL,
+    listen_http_port integer NOT NULL,
+    listen_pg_addr character varying NOT NULL,
+    listen_pg_port integer NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS hadron_timeline_safekeepers (
+    timeline_id character varying NOT NULL,
+    sk_node_id bigint NOT NULL,
+    legacy_endpoint_id uuid,
+    PRIMARY KEY(timeline_id, sk_node_id)
+);
+
+CREATE TABLE IF NOT EXISTS tenant_shards (
+    tenant_id character varying NOT NULL,
+    shard_number integer NOT NULL,
+    shard_count integer NOT NULL,
+    PRIMARY KEY(tenant_id, shard_number, shard_count),
+    shard_stripe_size integer NOT NULL,
+    generation integer,
+    generation_pageserver bigint,
+    placement_policy character varying NOT NULL,
+    splitting smallint NOT NULL,
+    config text NOT NULL,
+    scheduling_policy character varying DEFAULT '"Active"'::character varying NOT NULL,
+    preferred_az_id character varying
+);
+
+CREATE INDEX IF NOT EXISTS tenant_shards_tenant_id ON tenant_shards USING btree (tenant_id);
+
+CREATE TABLE IF NOT EXISTS metadata_health (
+    tenant_id character varying NOT NULL,
+    shard_number integer NOT NULL,
+    shard_count integer NOT NULL,
+    PRIMARY KEY(tenant_id, shard_number, shard_count),
+    -- Rely on cascade behavior for delete
+    FOREIGN KEY(tenant_id, shard_number, shard_count) REFERENCES tenant_shards ON DELETE CASCADE,
+    healthy boolean DEFAULT true NOT NULL,
+    last_scrubbed_at timestamp with time zone DEFAULT now() NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS nodes (
+    node_id bigint PRIMARY KEY NOT NULL,
+    scheduling_policy character varying NOT NULL,
+    listen_http_addr character varying NOT NULL,
+    listen_http_port integer NOT NULL,
+    listen_pg_addr character varying NOT NULL,
+    listen_pg_port integer NOT NULL,
+    availability_zone_id character varying NOT NULL,
+    listen_https_port integer,
+    lifecycle character varying DEFAULT 'active'::character varying NOT NULL,
+    listen_grpc_addr character varying,
+    listen_grpc_port integer
+);
+
+CREATE TABLE IF NOT EXISTS safekeeper_timeline_pending_ops (
+    sk_id bigint NOT NULL,
+    tenant_id character varying NOT NULL,
+    timeline_id character varying NOT NULL,
+    generation integer NOT NULL,
+    op_kind character varying NOT NULL,
+    PRIMARY KEY(tenant_id, timeline_id, sk_id)
+);
+
+CREATE TABLE IF NOT EXISTS safekeepers (
+    id bigint PRIMARY KEY NOT NULL,
+    region_id text NOT NULL,
+    version bigint NOT NULL,
+    host text NOT NULL,
+    port integer NOT NULL,
+    http_port integer NOT NULL,
+    availability_zone_id text NOT NULL,
+    scheduling_policy character varying DEFAULT 'activating'::character varying NOT NULL,
+    https_port integer
+);
+
+CREATE TABLE IF NOT EXISTS timeline_imports (
+    tenant_id character varying NOT NULL,
+    timeline_id character varying NOT NULL,
+    shard_statuses jsonb NOT NULL,
+    PRIMARY KEY(tenant_id, timeline_id)
+);
+
+CREATE TABLE IF NOT EXISTS timelines (
+    tenant_id character varying NOT NULL,
+    timeline_id character varying NOT NULL,
+    start_lsn pg_lsn NOT NULL,
+    generation integer NOT NULL,
+    sk_set bigint[] NOT NULL,
+    new_sk_set bigint[],
+    cplane_notified_generation integer NOT NULL,
+    deleted_at timestamp with time zone,
+    sk_set_notified_generation integer DEFAULT 1 NOT NULL,
+    PRIMARY KEY(tenant_id, timeline_id)
+);

test_runner/fixtures/neon_api.py

@@ -79,7 +79,6 @@ class NeonAPI:
                 elif resp.status_code == 423 and resp.json()["message"] in {
                     "endpoint is in some transitive state, could not suspend",
                     "project already has running conflicting operations, scheduling of new ones is prohibited",
-                    "snapshot is in transition",
                 }:
                     retry = True
                     self.retries4xx += 1
@@ -106,7 +105,6 @@ class NeonAPI:
         branch_name: str | None = None,
         branch_role_name: str | None = None,
         branch_database_name: str | None = None,
-        project_settings: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         data: dict[str, Any] = {
             "project": {
@@ -123,8 +121,6 @@ class NeonAPI:
             data["project"]["branch"]["role_name"] = branch_role_name
         if branch_database_name:
             data["project"]["branch"]["database_name"] = branch_database_name
-        if project_settings:
-            data["project"]["settings"] = project_settings
 
         resp = self.__request(
             "POST",
@@ -359,63 +355,6 @@ class NeonAPI:
 
         return cast("dict[str, Any]", resp.json())
 
-    def create_snapshot(
-        self,
-        project_id: str,
-        branch_id: str,
-        lsn: str | None = None,
-        timestamp: str | None = None,
-        name: str | None = None,
-        expires_at: str | None = None,
-    ) -> dict[str, Any]:
-        params: dict[str, Any] = {
-            "lsn": lsn,
-            "timestamp": timestamp,
-            "name": name,
-            "expires_at": expires_at,
-        }
-        params = {key: value for key, value in params.items() if value is not None}
-        resp = self.__request(
-            "POST",
-            f"/projects/{project_id}/branches/{branch_id}/snapshot",
-            params=params,
-            json={},
-            headers={
-                "Accept": "application/json",
-            },
-        )
-        return cast("dict[str, Any]", resp.json())
-
-    def delete_snapshot(self, project_id: str, snapshot_id: str) -> dict[str, Any]:
-        resp = self.__request("DELETE", f"/projects/{project_id}/snapshots/{snapshot_id}")
-        return cast("dict[str, Any]", resp.json())
-
-    def restore_snapshot(
-        self,
-        project_id: str,
-        snapshot_id: str,
-        target_branch_id: str,
-        name: str | None = None,
-        finalize_restore: bool = True,
-    ) -> dict[str, Any]:
-        data: dict[str, Any] = {
-            "target_branch_id": target_branch_id,
-            "finalize_restore": finalize_restore,
-        }
-        if name is not None:
-            data["name"] = name
-        log.info("Restore snapshot data: %s", data)
-        resp = self.__request(
-            "POST",
-            f"/projects/{project_id}/snapshots/{snapshot_id}/restore",
-            json=data,
-            headers={
-                "Accept": "application/json",
-                "Content-Type": "application/json",
-            },
-        )
-        return cast("dict[str, Any]", resp.json())
-
     def delete_endpoint(self, project_id: str, endpoint_id: str) -> dict[str, Any]:
         resp = self.__request("DELETE", f"/projects/{project_id}/endpoints/{endpoint_id}")
         return cast("dict[str,Any]", resp.json())
@@ -457,14 +396,6 @@ class NeonAPI:
 
         return cast("dict[str, Any]", resp.json())
 
-    def get_branch_endpoints(self, project_id: str, branch_id: str) -> dict[str, Any]:
-        resp = self.__request(
-            "GET",
-            f"/projects/{project_id}/branches/{branch_id}/endpoints",
-            headers={"Accept": "application/json", "Content-Type": "application/json"},
-        )
-        return cast("dict[str, Any]", resp.json())
-
     def get_endpoints(self, project_id: str) -> dict[str, Any]:
         resp = self.__request(
             "GET",

test_runner/random_ops/test_random_ops.py

@@ -11,7 +11,6 @@ import time
 from datetime import UTC, datetime, timedelta
 from typing import TYPE_CHECKING, Any
 
-import psycopg2
 import pytest
 from fixtures.log_helper import log
 
@@ -23,29 +22,6 @@ if TYPE_CHECKING:
     from fixtures.pg_version import PgVersion
 
 
-class NeonSnapshot:
-    """
-    A snapshot of the Neon Branch
-    Gets the output of the API call af a snapshot creation
-    """
-
-    def __init__(self, project: NeonProject, snapshot: dict[str, Any]):
-        self.project: NeonProject = project
-        snapshot = snapshot["snapshot"]
-        self.id: str = snapshot["id"]
-        self.name: str = snapshot["name"]
-        self.created_at: datetime = datetime.fromisoformat(snapshot["created_at"])
-        self.source_branch: NeonBranch = project.branches[snapshot["source_branch_id"]]
-        project.snapshots[self.id] = self
-        self.restored: bool = False
-
-    def __str__(self) -> str:
-        return f"id: {self.id}, name: {self.name}, created_at: {self.created_at}"
-
-    def delete(self) -> None:
-        self.project.delete_snapshot(self.id)
-
-
 class NeonEndpoint:
     """
     Neon Endpoint
@@ -91,21 +67,9 @@ class NeonBranch:
     is_reset defines if the branch is a reset one i.e. created as a result of the reset API Call
     """
 
-    def __init__(
-        self,
-        project,
-        branch: dict[str, Any],
-        is_reset=False,
-        primary_branch: NeonBranch | None = None,
-    ):
+    def __init__(self, project, branch: dict[str, Any], is_reset=False):
         self.id: str = branch["branch"]["id"]
         self.desc = branch
-        self.name: str | None = None
-        if "name" in branch["branch"]:
-            self.name = branch["branch"]["name"]
-        self.restored_from: str | None = None
-        if "restored_from" in branch["branch"]:
-            self.restored_from = branch["branch"]["restored_from"]
         self.project: NeonProject = project
         self.neon_api: NeonAPI = project.neon_api
         self.project_id: str = branch["branch"]["project_id"]
@@ -146,36 +110,13 @@ class NeonBranch:
                 "PGPASSWORD": self.connection_parameters["password"],
                 "PGSSLMODE": "require",
             }
-        self.replicas: dict[str, NeonBranch] = {}
-        self.primary_branch: NeonBranch | None = primary_branch
-        if primary_branch:
-            if not self.connection_parameters:
-                raise ValueError(
-                    "connection_parameters is required when primary_branch is specified"
-                )
-            self.project.replicas[self.id] = self
-            primary_branch.replicas[self.id] = self
-            with psycopg2.connect(primary_branch.connstr()) as conn:
-                with conn.cursor() as cur:
-                    cur.execute(f"CREATE PUBLICATION {self.id} FOR ALL TABLES")
-                conn.commit()
-            with psycopg2.connect(self.connstr()) as conn:
-                with conn.cursor() as cur:
-                    cur.execute(
-                        f"CREATE SUBSCRIPTION {self.id} CONNECTION '{primary_branch.connstr()}' PUBLICATION {self.id}"
-                    )
-                conn.commit()
 
     def __str__(self):
         """
-        Prints the branch's information with all the predecessors
+        Prints the branch's name with all the predecessors
+        (r) means the branch is a reset one
         """
-        name = f"({self.name})" if self.name and self.name != self.id else ""
-        restored_from = f"(restored_from: {self.restored_from})" if self.restored_from else ""
-        ancestor = (
-            f" <- {self.primary_branch}" if self.primary_branch else f", parent: {self.parent}"
-        )
-        return f"{self.id}{name}{restored_from}{ancestor}"
+        return f"{self.id}{'(r)' if self.id in self.project.reset_branches else ''}, parent: {self.parent}"
 
     def random_time(self) -> datetime:
         min_time = max(
@@ -187,10 +128,8 @@ class NeonBranch:
         log.info("min_time: %s, max_time: %s", min_time, max_time)
         return (min_time + (max_time - min_time) * random.random()).replace(microsecond=0)
 
-    def create_child_branch(
-        self, parent_timestamp: datetime | None = None, primary_branch: NeonBranch | None = None
-    ) -> NeonBranch | None:
-        return self.project.create_branch(self.id, parent_timestamp, primary_branch=primary_branch)
+    def create_child_branch(self, parent_timestamp: datetime | None = None) -> NeonBranch | None:
+        return self.project.create_branch(self.id, parent_timestamp)
 
     def create_ro_endpoint(self) -> NeonEndpoint | None:
         if not self.project.check_limit_endpoints():
@@ -213,9 +152,6 @@ class NeonBranch:
         self.project.terminate_benchmark(self.id)
 
     def reset_to_parent(self) -> None:
-        """
-        Resets the branch to the parent branch
-        """
         for ep in self.project.endpoints.values():
             if ep.type == "read_only":
                 ep.terminate_benchmark()
@@ -281,19 +217,6 @@ class NeonBranch:
             ep.start_benchmark()
         return res
 
-    def create_logical_replica(self) -> NeonBranch | None:
-        if self.primary_branch is not None:
-            raise RuntimeError("The primary branch cannot be a logical replica")
-        if self.id in self.project.reset_branches:
-            raise RuntimeError("Reset branch cannot be a primary branch")
-        replica = self.create_child_branch(primary_branch=self)
-        return replica
-
-    def connstr(self):
-        if self.connection_parameters is None:
-            raise RuntimeError("Connection parameters are not defined")
-        return " ".join([f"{key}={value}" for key, value in self.connection_parameters.items()])
-
 
 class NeonProject:
     """
@@ -305,9 +228,7 @@ class NeonProject:
         self.neon_api = neon_api
         self.pg_bin = pg_bin
         proj = self.neon_api.create_project(
-            pg_version,
-            f"Automatic random API test GITHUB_RUN_ID={os.getenv('GITHUB_RUN_ID')}",
-            project_settings={"enable_logical_replication": True},
+            pg_version, f"Automatic random API test GITHUB_RUN_ID={os.getenv('GITHUB_RUN_ID')}"
         )
         self.id: str = proj["project"]["id"]
         self.name: str = proj["project"]["name"]
@@ -319,7 +240,6 @@ class NeonProject:
         # Leaf branches are the branches, which do not have children
         self.leaf_branches: dict[str, NeonBranch] = {}
         self.branches: dict[str, NeonBranch] = {}
-        self.branch_num: int = 0
         self.reset_branches: set[str] = set()
         self.main_branch: NeonBranch = NeonBranch(self, proj)
         self.main_branch.connection_parameters = self.connection_parameters
@@ -333,9 +253,6 @@ class NeonProject:
         self.limits: dict[str, Any] = self.get_limits()["limits"]
         self.read_only_endpoints_total: int = 0
         self.min_time: datetime = datetime.now(UTC)
-        self.snapshots: dict[str, NeonSnapshot] = {}
-        self.snapshot_num: int = 0
-        self.replicas: dict[str, NeonBranch] = {}
 
     def get_limits(self) -> dict[str, Any]:
         return self.neon_api.get_project_limits(self.id)
@@ -363,11 +280,7 @@ class NeonProject:
         return False
 
     def create_branch(
-        self,
-        parent_id: str | None = None,
-        parent_timestamp: datetime | None = None,
-        is_reset: bool = False,
-        primary_branch: NeonBranch | None = None,
+        self, parent_id: str | None = None, parent_timestamp: datetime | None = None
     ) -> NeonBranch | None:
         self.wait()
         if not self.check_limit_branches():
@@ -380,14 +293,14 @@ class NeonProject:
         branch_def = self.neon_api.create_branch(
             self.id, parent_id=parent_id, parent_timestamp=parent_timestamp_str
         )
-        new_branch = NeonBranch(self, branch_def, is_reset, primary_branch)
+        new_branch = NeonBranch(self, branch_def)
         self.wait()
         return new_branch
 
     def delete_branch(self, branch_id: str) -> None:
         parent = self.branches[branch_id].parent
         if not parent or branch_id == self.main_branch.id:
-            raise RuntimeError("Cannot delete the main branch or a branch restored from a snapshot")
+            raise RuntimeError("Cannot delete the main branch")
         if branch_id not in self.leaf_branches and branch_id not in self.reset_branches:
             raise RuntimeError(f"The branch {branch_id}, probably, has ancestors")
         if branch_id not in self.branches:
@@ -400,18 +313,7 @@ class NeonProject:
         if branch_id not in self.reset_branches:
             self.terminate_benchmark(branch_id)
         self.neon_api.delete_branch(self.id, branch_id)
-        primary_branch = self.branches[branch_id].primary_branch
-        if primary_branch is not None:
-            with psycopg2.connect(primary_branch.connstr()) as conn:
-                with conn.cursor() as cur:
-                    cur.execute(f"DROP PUBLICATION {branch_id}")
-                conn.commit()
-            parent.replicas.pop(branch_id)
-            self.replicas.pop(branch_id)
-        else:
-            for replica in self.branches[branch_id].replicas.values():
-                replica.delete()
-        if len(parent.children) == 1 and parent.parent is not None:
+        if len(parent.children) == 1 and parent.id != self.main_branch.id:
             self.leaf_branches[parent.id] = parent
         parent.children.pop(branch_id)
         if branch_id in self.leaf_branches:
@@ -431,26 +333,6 @@ class NeonProject:
             log.info("No leaf branches found")
         return target
 
-    def get_random_parent_branch(self) -> NeonBranch:
-        return self.branches[
-            random.choice(
-                list(set(self.branches.keys()) - self.reset_branches - set(self.replicas.keys()))
-            )
-        ]
-
-    def gen_branch_name(self) -> str:
-        self.branch_num += 1
-        return f"branch{self.branch_num}"
-
-    def get_random_snapshot(self) -> NeonSnapshot | None:
-        snapshot: NeonSnapshot | None = None
-        avail_snapshots = [sn for sn in self.snapshots.values() if not sn.restored]
-        if avail_snapshots:
-            snapshot = random.choice(avail_snapshots)
-        else:
-            log.info("No snapshots found")
-        return snapshot
-
     def delete_endpoint(self, endpoint_id: str) -> None:
         self.terminate_benchmark(endpoint_id)
         self.neon_api.delete_endpoint(self.id, endpoint_id)
@@ -527,116 +409,6 @@ class NeonProject:
         self.restore_num += 1
         return f"restore{self.restore_num}"
 
-    def gen_snapshot_name(self) -> str:
-        self.snapshot_num += 1
-        return f"snapshot{self.snapshot_num}"
-
-    def create_snapshot(
-        self,
-        lsn: str | None = None,
-        timestamp: datetime | None = None,
-    ) -> NeonSnapshot:
-        """
-        Create a new Neon snapshot for the current project
-        Two optional arguments: lsn and timestamp are mutually exclusive
-        they instruct to create a snapshot with the specific lns or timestamp
-        """
-        snapshot_name = self.gen_snapshot_name()
-        with psycopg2.connect(self.connection_uri) as conn:
-            with conn.cursor() as cur:
-                # We will check the value we set now after the snapshot restored to verify consistency
-                cur.execute(
-                    f"INSERT INTO sanity_check (name, value) VALUES "
-                    f"('snapsot_name', '{snapshot_name}') ON CONFLICT (name) DO UPDATE SET value = EXCLUDED.value"
-                )
-                conn.commit()
-                snapshot = NeonSnapshot(
-                    self,
-                    self.neon_api.create_snapshot(
-                        self.id,
-                        self.main_branch.id,
-                        lsn,
-                        timestamp.isoformat().replace("+00:00", "Z") if timestamp else None,
-                        snapshot_name,
-                    ),
-                )
-                self.wait()
-                # Now we taint the value after the snapshot was taken
-                cur.execute("UPDATE sanity_check SET value = 'tainted' || value")
-                conn.commit()
-        return snapshot
-
-    def delete_snapshot(self, snapshot_id: str) -> None:
-        """
-        Deletes the snapshot with the given id
-        """
-        self.wait()
-        self.neon_api.delete_snapshot(self.id, snapshot_id)
-        self.snapshots.pop(snapshot_id)
-        self.wait()
-
-    def restore_snapshot(self, snapshot_id: str) -> NeonBranch | None:
-        """
-        Creates a new Neon branch for the current project, then restores the snapshot
-        with the given id
-        """
-        target_branch = self.get_random_parent_branch().create_child_branch()
-        if not target_branch:
-            return None
-        self.snapshots[snapshot_id].restored = True
-        new_branch_def: dict[str, Any] = self.neon_api.restore_snapshot(
-            self.id,
-            snapshot_id,
-            target_branch.id,
-            self.gen_branch_name(),
-        )
-        self.wait()
-        new_branch_def = self.neon_api.get_branch_details(self.id, new_branch_def["branch"]["id"])
-        # The restored branch will lose the parent afterward, but it has it during the restoration.
-        # So, we delete parent_id
-        new_branch_def["branch"].pop("parent_id")
-        new_branch = NeonBranch(self, new_branch_def)
-        log.info("Restored snapshot to the branch: %s", new_branch)
-        target_branch_def = self.neon_api.get_branch_details(self.id, target_branch.id)
-        if "name" in target_branch_def["branch"]:
-            target_branch.name = target_branch_def["branch"]["name"]
-        if new_branch.connection_parameters is None:
-            if not new_branch.endpoints:
-                for ep in self.neon_api.get_branch_endpoints(self.id, new_branch.id)["endpoints"]:
-                    if ep["id"] not in self.endpoints:
-                        NeonEndpoint(self, ep)
-            new_branch.connection_parameters = self.connection_parameters.copy()
-            for ep in new_branch.endpoints.values():
-                if ep.type == "read_write":
-                    new_branch.connection_parameters["host"] = ep.host
-                    break
-            new_branch.connect_env = {
-                "PGHOST": new_branch.connection_parameters["host"],
-                "PGUSER": new_branch.connection_parameters["role"],
-                "PGDATABASE": new_branch.connection_parameters["database"],
-                "PGPASSWORD": new_branch.connection_parameters["password"],
-                "PGSSLMODE": "require",
-            }
-        with psycopg2.connect(
-            host=new_branch.connection_parameters["host"],
-            port=5432,
-            user=new_branch.connection_parameters["role"],
-            password=new_branch.connection_parameters["password"],
-            database=new_branch.connection_parameters["database"],
-        ) as conn:
-            with conn.cursor() as cur:
-                cur.execute("SELECT value FROM sanity_check WHERE name = 'snapsot_name'")
-                snapshot_name = None
-                if row := cur.fetchone():
-                    snapshot_name = row[0]
-                # We verify here that the value we select from the table matches with the snapshot name
-                # To ensure consistency
-                assert snapshot_name == self.snapshots[snapshot_id].name
-        self.wait()
-        target_branch.start_benchmark()
-        new_branch.start_benchmark()
-        return new_branch
-
 
 @pytest.fixture()
 def setup_class(
@@ -666,7 +438,9 @@ def do_action(project: NeonProject, action: str) -> bool:
     if action == "new_branch" or action == "new_branch_random_time":
         use_random_time: bool = action == "new_branch_random_time"
         log.info("Trying to create a new branch %s", "random time" if use_random_time else "")
-        parent = project.get_random_parent_branch()
+        parent = project.branches[
+            random.choice(list(set(project.branches.keys()) - project.reset_branches))
+        ]
         child = parent.create_child_branch(parent.random_time() if use_random_time else None)
         if child is None:
             return False
@@ -705,31 +479,6 @@ def do_action(project: NeonProject, action: str) -> bool:
             return False
         log.info("Reset to parent %s", target)
         target.reset_to_parent()
-    elif action == "create_snapshot":
-        snapshot = project.create_snapshot()
-        if snapshot is None:
-            return False
-        log.info("Created snapshot %s", snapshot)
-    elif action == "restore_snapshot":
-        if (snapshot_to_restore := project.get_random_snapshot()) is None:
-            return False
-        log.info("Restoring snapshot %s", snapshot_to_restore)
-        if project.restore_snapshot(snapshot_to_restore.id) is None:
-            return False
-    elif action == "delete_snapshot":
-        snapshot_to_delete = project.get_random_snapshot()
-        if snapshot_to_delete is None:
-            return False
-        snapshot_to_delete.delete()
-        log.info("Deleted snapshot %s", snapshot_to_delete)
-    elif action == "create_logical_replica":
-        primary: NeonBranch | None = project.get_random_parent_branch()
-        if primary is None:
-            return False
-        replica: NeonBranch | None = primary.create_logical_replica()
-        if replica is None:
-            return False
-        log.info("Created logical replica %s", replica)
     else:
         raise ValueError(f"The action {action} is unknown")
     return True
@@ -763,28 +512,12 @@ def test_api_random(
         ("delete_branch", 1.2),
         ("restore_random_time", 0.9),
         ("reset_to_parent", 0.3),
-        ("create_snapshot", 0.2),
-        ("restore_snapshot", 0.1),
-        ("delete_snapshot", 0.1),
     )
     if num_ops_env := os.getenv("NUM_OPERATIONS"):
         num_operations = int(num_ops_env)
     else:
         num_operations = 250
     pg_bin.run(["pgbench", "-i", "-I", "dtGvp", "-s100"], env=project.main_branch.connect_env)
-    # Create a table for sanity check
-    # We are going to leve some control values there to check, e.g., after restoring a snapshot
-    pg_bin.run(
-        [
-            "psql",
-            "-c",
-            "CREATE TABLE IF NOT EXISTS sanity_check (name VARCHAR NOT NULL PRIMARY KEY, value VARCHAR)",
-        ],
-        env=project.main_branch.connect_env,
-    )
-    # To not go to the past where pgbench tables do not exist
-    time.sleep(1)
-    project.min_time = datetime.now(UTC)
     # To not go to the past where pgbench tables do not exist
     time.sleep(1)
     project.min_time = datetime.now(UTC)