Repository: xataio/pgstream
Pull request #190 (2025-03-18 11:34:50)

Update snapshot tutorial
This PR updates the snapshot tutorial to add the `pgstream.snapshot_requests`. The snapshot store is updated to create the `pgstream` schema if it doesn't exist, since the snapshot mode can be run without the need to initialise pgstream (`pgstream init`).

github-actions[bot]: ### Merging this branch will **increase** overall coverage

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/pkg/snapshot/store/postgres | 85.71% (**+1.10%**) | :thumbsup: |

---

<details>

<summary>Coverage by file</summary>

### Changed files (no unit tests)

| Changed File | Coverage Δ | Total | Covered | Missed | :robot: |
|--------------|------------|-------|---------|--------|---------|
| github.com/xataio/pgstream/pkg/snapshot/store/postgres/pg_snapshot_store.go | 85.71% (**+1.10%**) | 56 (+4) | 48 (+4) | 8 | :thumbsup: |

_Please note that the "Total", "Covered", and "Missed" counts above refer to ***code statements*** instead of lines of code. The value in brackets refers to the test coverage of that file in the old version of the code._

### Changed unit test files

- github.com/xataio/pgstream/pkg/snapshot/store/postgres/pg_snapshot_store_test.go

</details>

github-actions[bot]: ### Merging this branch will **increase** overall coverage

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/pkg/snapshot/store/postgres | 85.71% (**+1.10%**) | :thumbsup: |

---

<details>

<summary>Coverage by file</summary>

### Changed files (no unit tests)

| Changed File | Coverage Δ | Total | Covered | Missed | :robot: |
|--------------|------------|-------|---------|--------|---------|
| github.com/xataio/pgstream/pkg/snapshot/store/postgres/pg_snapshot_store.go | 85.71% (**+1.10%**) | 56 (+4) | 48 (+4) | 8 | :thumbsup: |

_Please note that the "Total", "Covered", and "Missed" counts above refer to ***code statements*** instead of lines of code. The value in brackets refers to the test coverage of that file in the old version of the code._

### Changed unit test files

- github.com/xataio/pgstream/pkg/snapshot/store/postgres/pg_snapshot_store_test.go

</details>


Patch:
diff --git a/docs/tutorials/postgres_snapshot.md b/docs/tutorials/postgres_snapshot.md
index a43076d..a419099 100644
--- a/docs/tutorials/postgres_snapshot.md
+++ b/docs/tutorials/postgres_snapshot.md
@@ -77,6 +77,12 @@ PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
 ```
 
+The snapshot listener can also be configured to record and update the status of snapshot requests in a dedicated table `snapshot_requests` under the `pgstream` schema. This allows to only perform a given snapshot once by keeping track of what's already been completed. All it's needed is the URL of the database where the table should be created. For this tutorial, we'll use the source database.
+
+```sh
+PGSTREAM_POSTGRES_SNAPSHOT_STORE_URL="postgres://postgres:postgres@localhost:5432?sslmode=disable"
+```
+
 ### Processor
 
 With the listener side ready, the next step is to configure the processor. Since we want the snapshot reach a PostgreSQL database, we will set the PostgreSQL writer configuration variables. The only required value is the URL of the target database, where the snapshotted schema/data from the source database will be streamed. We use the URL of the docker PostgreSQL database we started earlier (note the port is the only difference between the source and the target PostgreSQL databases).
@@ -107,6 +113,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
+PGSTREAM_POSTGRES_SNAPSHOT_STORE_URL="postgres://postgres:postgres@localhost:5432?sslmode=disable"
 
 # Processor config
 PGSTREAM_POSTGRES_WRITER_TARGET_URL="postgres://postgres:postgres@localhost:7654?sslmode=disable"
@@ -167,6 +174,23 @@ postgres@localhost:postgres> SELECT * FROM test;
 +----+---------+
 ```
 
+We can also check the status of the snapshot by querying the `pgstream.snapshot_requests` table on the source database. It should contain the details of the snapshot that has been requested.
+
+```sh
+➜ psql postgresql://postgres:postgres@localhost:5432/postgres
+```
+
+```sql
+postgres@localhost:postgres> SELECT * FROM pgstream.snapshot_requests;
++--------+-------------+-------------+-------------------------------+-------------------------------+-----------+--------+
+| req_id | schema_name | table_names | created_at                    | updated_at                    | status    | errors |
+|--------+-------------+-------------+-------------------------------+-------------------------------+-----------+--------|
+| 1      | public      | ['*']       | 2025-03-18 11:27:21.683361+00 | 2025-03-18 11:27:21.843207+00 | completed | <null> |
++--------+-------------+-------------+-------------------------------+-------------------------------+-----------+--------+
+```
+
+If the snapshot has completed with errors, it will be retried on the next run of pgstream.
+
 ## Troubleshooting
 
 Here are some common issues you might encounter while following this tutorial and how to resolve them:
diff --git a/docs/tutorials/snapshot2pg_tutorial.env b/docs/tutorials/snapshot2pg_tutorial.env
index abb7dac..b33b638 100644
--- a/docs/tutorials/snapshot2pg_tutorial.env
+++ b/docs/tutorials/snapshot2pg_tutorial.env
@@ -5,6 +5,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
+PGSTREAM_POSTGRES_SNAPSHOT_STORE_URL="postgres://postgres:postgres@localhost:5432?sslmode=disable"
 
 # Processor config
 PGSTREAM_POSTGRES_WRITER_TARGET_URL="postgres://postgres:postgres@localhost:7654?sslmode=disable"
diff --git a/pkg/snapshot/store/postgres/pg_snapshot_store.go b/pkg/snapshot/store/postgres/pg_snapshot_store.go
index 229e872..63a0a63 100644
--- a/pkg/snapshot/store/postgres/pg_snapshot_store.go
+++ b/pkg/snapshot/store/postgres/pg_snapshot_store.go
@@ -105,7 +105,13 @@ func (s *Store) GetSnapshotRequestsBySchema(ctx context.Context, schema string)
 }
 
 func (s *Store) createTable(ctx context.Context) error {
-	createQuery := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s(
+	createSchemaQuery := fmt.Sprintf(`CREATE SCHEMA IF NOT EXISTS %s`, store.SchemaName)
+	_, err := s.conn.Exec(ctx, createSchemaQuery)
+	if err != nil {
+		return fmt.Errorf("error creating pgstream schema for snapshot store: %w", err)
+	}
+
+	createTableQuery := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s(
 	req_id SERIAL PRIMARY KEY,
 	schema_name TEXT,
 	table_names TEXT[],
@@ -113,7 +119,7 @@ func (s *Store) createTable(ctx context.Context) error {
 	updated_at TIMESTAMP WITH TIME ZONE,
 	status TEXT CHECK (status IN ('requested', 'in progress', 'completed')),
 	errors JSONB )`, snapshotsTable())
-	_, err := s.conn.Exec(ctx, createQuery)
+	_, err = s.conn.Exec(ctx, createTableQuery)
 	if err != nil {
 		return fmt.Errorf("error creating snapshots postgres table: %w", err)
 	}


Test patch:
diff --git a/pkg/snapshot/store/postgres/pg_snapshot_store_test.go b/pkg/snapshot/store/postgres/pg_snapshot_store_test.go
index 808d66b..188f7bb 100644
--- a/pkg/snapshot/store/postgres/pg_snapshot_store_test.go
+++ b/pkg/snapshot/store/postgres/pg_snapshot_store_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/xataio/pgstream/internal/postgres"
 	postgresmocks "github.com/xataio/pgstream/internal/postgres/mocks"
 	"github.com/xataio/pgstream/pkg/snapshot"
+	"github.com/xataio/pgstream/pkg/snapshot/store"
 )
 
 func TestStore_CreateSnapshotRequest(t *testing.T) {
@@ -428,6 +429,10 @@ func TestStore_createTable(t *testing.T) {
 				ExecFn: func(ctx context.Context, i uint, s string, a ...any) (postgres.CommandTag, error) {
 					switch i {
 					case 1:
+						wantQuery := fmt.Sprintf(`CREATE SCHEMA IF NOT EXISTS %s`, store.SchemaName)
+						require.Equal(t, wantQuery, s)
+						require.Empty(t, a)
+					case 2:
 						wantQuery := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s(
 	req_id SERIAL PRIMARY KEY,
 	schema_name TEXT,
@@ -438,7 +443,7 @@ func TestStore_createTable(t *testing.T) {
 	errors JSONB )`, snapshotsTable())
 						require.Equal(t, wantQuery, s)
 						require.Empty(t, a)
-					case 2:
+					case 3:
 						wantQuery := fmt.Sprintf(`CREATE UNIQUE INDEX IF NOT EXISTS schema_table_status_unique_index
 	ON %s(schema_name,table_names) WHERE status != 'completed'`, snapshotsTable())
 						require.Equal(t, wantQuery, s)
@@ -452,12 +457,28 @@ func TestStore_createTable(t *testing.T) {
 			},
 			wantErr: nil,
 		},
+		{
+			name: "error - creating schema",
+			querier: &postgresmocks.Querier{
+				ExecFn: func(ctx context.Context, i uint, s string, a ...any) (postgres.CommandTag, error) {
+					switch i {
+					case 1:
+						return postgres.CommandTag{}, errTest
+					default:
+						return postgres.CommandTag{}, fmt.Errorf("unexpected Exec call: %d", i)
+					}
+				},
+			},
+			wantErr: errTest,
+		},
 		{
 			name: "error - creating table",
 			querier: &postgresmocks.Querier{
 				ExecFn: func(ctx context.Context, i uint, s string, a ...any) (postgres.CommandTag, error) {
 					switch i {
 					case 1:
+						return postgres.CommandTag{}, nil
+					case 2:
 						return postgres.CommandTag{}, errTest
 					default:
 						return postgres.CommandTag{}, fmt.Errorf("unexpected Exec call: %d", i)
@@ -474,6 +495,8 @@ func TestStore_createTable(t *testing.T) {
 					case 1:
 						return postgres.CommandTag{}, nil
 					case 2:
+						return postgres.CommandTag{}, nil
+					case 3:
 						return postgres.CommandTag{}, errTest
 					default:
 						return postgres.CommandTag{}, fmt.Errorf("unexpected Exec call: %d", i)