Repository: xataio/pgstream
Pull request #95 (2024-12-03 16:21:24)

Update snapshot store
This PR updates the existing snapshot store, since it was initially designed to work with the backfilling approach which required identity columns. Since we will no longer rely on a primary key/unique column to produce the snapshot, it can be removed. Additionally, the snapshot structs have been refined, separating the snapshot request, and including a more thorough recording of snapshot errors at snapshot/table level, which will be helpful to retry failed snapshots for only the affected relations. 
A new method to retrieve snapshot requests by schema is added.


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

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/pkg/snapshot | 0.00% (ø) |  |
| github.com/xataio/pgstream/pkg/snapshot/store | 0.00% (ø) |  |
| github.com/xataio/pgstream/pkg/snapshot/store/mocks | 0.00% (ø) |  |
| github.com/xataio/pgstream/pkg/snapshot/store/postgres | 84.62% (**+3.22%**) | :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/errors.go | 0.00% (ø) | 66 (+66) | 0 | 66 (+66) |  |
| github.com/xataio/pgstream/pkg/snapshot/snapshot.go | 0.00% (ø) | 22 (+14) | 0 | 22 (+14) |  |
| github.com/xataio/pgstream/pkg/snapshot/store/mocks/mock_snapshot_store.go | 0.00% (ø) | 5 (+2) | 0 | 5 (+2) |  |
| github.com/xataio/pgstream/pkg/snapshot/store/postgres/pg_snapshot_store.go | 84.62% (**+3.22%**) | 52 (+9) | 44 (+9) | 8 | :thumbsup: |
| github.com/xataio/pgstream/pkg/snapshot/store/snapshot_store.go | 0.00% (ø) | 0 | 0 | 0 |  |

_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 | 0.00% (ø) |  |
| github.com/xataio/pgstream/pkg/snapshot/store | 0.00% (ø) |  |
| github.com/xataio/pgstream/pkg/snapshot/store/mocks | 0.00% (ø) |  |
| github.com/xataio/pgstream/pkg/snapshot/store/postgres | 84.62% (**+3.22%**) | :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/errors.go | 0.00% (ø) | 55 (+55) | 0 | 55 (+55) |  |
| github.com/xataio/pgstream/pkg/snapshot/snapshot.go | 0.00% (ø) | 22 (+14) | 0 | 22 (+14) |  |
| github.com/xataio/pgstream/pkg/snapshot/store/mocks/mock_snapshot_store.go | 0.00% (ø) | 5 (+2) | 0 | 5 (+2) |  |
| github.com/xataio/pgstream/pkg/snapshot/store/postgres/pg_snapshot_store.go | 84.62% (**+3.22%**) | 52 (+9) | 44 (+9) | 8 | :thumbsup: |
| github.com/xataio/pgstream/pkg/snapshot/store/snapshot_store.go | 0.00% (ø) | 0 | 0 | 0 |  |

_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/pkg/snapshot/errors.go b/pkg/snapshot/errors.go
new file mode 100644
index 0000000..166ac45
--- /dev/null
+++ b/pkg/snapshot/errors.go
@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package snapshot
+
+import "errors"
+
+type Errors struct {
+	Snapshot error        `json:"snapshot,omitempty"`
+	Tables   []TableError `json:"tables,omitempty"`
+}
+
+type TableError struct {
+	Table    string `json:"table"`
+	ErrorMsg string `json:"error"`
+}
+
+func (e *Errors) Error() string {
+	var err error
+	if e.Snapshot != nil {
+		err = e.Snapshot
+	}
+
+	for _, table := range e.Tables {
+		if err == nil {
+			err = table
+			continue
+		}
+		err = errors.Join(err, table)
+	}
+	return err.Error()
+}
+
+func (e *Errors) IsTableError(table string) bool {
+	if e == nil {
+		return false
+	}
+
+	if e.Snapshot != nil {
+		return true
+	}
+
+	// treat wildcard table errors as true if there are any table errors.
+	if table == "*" && len(e.Tables) > 0 {
+		return true
+	}
+
+	for _, tableErr := range e.Tables {
+		if tableErr.Table == table {
+			return true
+		}
+	}
+	return false
+}
+
+func (e *Errors) GetFailedTables() []string {
+	if e == nil {
+		return nil
+	}
+	failedTables := make([]string, 0, len(e.Tables))
+	for _, table := range e.Tables {
+		failedTables = append(failedTables, table.Table)
+	}
+	return failedTables
+}
+
+func NewTableError(table string, err error) TableError {
+	errMsg := ""
+	if err != nil {
+		errMsg = err.Error()
+	}
+	return TableError{Table: table, ErrorMsg: errMsg}
+}
+
+func (e TableError) Error() string {
+	return "snapshot error for table " + e.Table + ": " + e.ErrorMsg
+}
diff --git a/pkg/snapshot/snapshot.go b/pkg/snapshot/snapshot.go
index f794e72..1ce0bb4 100644
--- a/pkg/snapshot/snapshot.go
+++ b/pkg/snapshot/snapshot.go
@@ -2,12 +2,19 @@
 
 package snapshot
 
+import (
+	"errors"
+)
+
 type Snapshot struct {
-	SchemaName          string
-	TableName           string
-	IdentityColumnNames []string
-	Status              Status
-	Error               error
+	SchemaName string
+	TableNames []string
+}
+
+type Request struct {
+	Snapshot Snapshot
+	Status   Status
+	Errors   *Errors
 }
 
 type Status string
@@ -19,14 +26,35 @@ const (
 )
 
 func (s *Snapshot) IsValid() bool {
-	return s != nil && s.SchemaName != "" && s.TableName != "" && len(s.IdentityColumnNames) > 0
+	return s != nil && s.SchemaName != "" && len(s.TableNames) > 0
+}
+
+func (r *Request) MarkCompleted(err error) {
+	r.Status = StatusCompleted
+	if err == nil {
+		return
+	}
+
+	var snapshotErrs *Errors
+	if errors.As(err, &snapshotErrs) {
+		r.Errors = snapshotErrs
+	} else {
+		r.Errors = &Errors{Snapshot: err}
+	}
+}
+
+func (r *Request) MarkInProgress() {
+	r.Status = StatusInProgress
+}
+
+func (r *Request) IsPending() bool {
+	return r.Status == StatusRequested
 }
 
-func (s *Snapshot) MarkCompleted(err error) {
-	s.Status = StatusCompleted
-	s.Error = err
+func (r *Request) HasFailed() bool {
+	return r.Status == StatusCompleted && r.Errors != nil
 }
 
-func (s *Snapshot) MarkInProgress() {
-	s.Status = StatusInProgress
+func (r *Request) HasFailedForTable(table string) bool {
+	return r.Status == StatusCompleted && r.Errors.IsTableError(table)
 }
diff --git a/pkg/snapshot/store/mocks/mock_snapshot_store.go b/pkg/snapshot/store/mocks/mock_snapshot_store.go
index 4e90130..35cfcdb 100644
--- a/pkg/snapshot/store/mocks/mock_snapshot_store.go
+++ b/pkg/snapshot/store/mocks/mock_snapshot_store.go
@@ -9,19 +9,28 @@ import (
 )
 
 type Store struct {
-	CreateSnapshotRequestFn func(context.Context, *snapshot.Snapshot) error
-	UpdateSnapshotRequestFn func(context.Context, *snapshot.Snapshot) error
-	GetSnapshotRequestsFn   func(ctx context.Context, status snapshot.Status) ([]*snapshot.Snapshot, error)
+	CreateSnapshotRequestFn       func(context.Context, *snapshot.Request) error
+	UpdateSnapshotRequestFn       func(context.Context, *snapshot.Request) error
+	GetSnapshotRequestsByStatusFn func(ctx context.Context, status snapshot.Status) ([]*snapshot.Request, error)
+	GetSnapshotRequestsBySchemaFn func(ctx context.Context, s string) ([]*snapshot.Request, error)
 }
 
-func (m *Store) CreateSnapshotRequest(ctx context.Context, s *snapshot.Snapshot) error {
+func (m *Store) CreateSnapshotRequest(ctx context.Context, s *snapshot.Request) error {
 	return m.CreateSnapshotRequestFn(ctx, s)
 }
 
-func (m *Store) UpdateSnapshotRequest(ctx context.Context, s *snapshot.Snapshot) error {
+func (m *Store) UpdateSnapshotRequest(ctx context.Context, s *snapshot.Request) error {
 	return m.UpdateSnapshotRequestFn(ctx, s)
 }
 
-func (m *Store) GetSnapshotRequests(ctx context.Context, status snapshot.Status) ([]*snapshot.Snapshot, error) {
-	return m.GetSnapshotRequestsFn(ctx, status)
+func (m *Store) GetSnapshotRequestsByStatus(ctx context.Context, status snapshot.Status) ([]*snapshot.Request, error) {
+	return m.GetSnapshotRequestsByStatusFn(ctx, status)
+}
+
+func (m *Store) GetSnapshotRequestsBySchema(ctx context.Context, s string) ([]*snapshot.Request, error) {
+	return m.GetSnapshotRequestsBySchemaFn(ctx, s)
+}
+
+func (m *Store) Close() error {
+	return nil
 }
diff --git a/pkg/snapshot/store/postgres/pg_snapshot_store.go b/pkg/snapshot/store/postgres/pg_snapshot_store.go
index 9adfe6c..fb92dce 100644
--- a/pkg/snapshot/store/postgres/pg_snapshot_store.go
+++ b/pkg/snapshot/store/postgres/pg_snapshot_store.go
@@ -18,7 +18,7 @@ type Store struct {
 
 const queryLimit = 1000
 
-func NewStore(ctx context.Context, url string) (*Store, error) {
+func New(ctx context.Context, url string) (*Store, error) {
 	conn, err := postgres.NewConnPool(ctx, url)
 	if err != nil {
 		return nil, err
@@ -40,69 +40,86 @@ func (s *Store) Close() error {
 	return s.conn.Close(context.Background())
 }
 
-func (s *Store) CreateSnapshotRequest(ctx context.Context, req *snapshot.Snapshot) error {
-	query := fmt.Sprintf(`INSERT INTO %s (schema_name, table_name, identity_column_names, created_at, updated_at, status)
-	VALUES($1, $2, $3,'now()','now()','requested')`, snapshotsTable())
-	_, err := s.conn.Exec(ctx, query, req.SchemaName, req.TableName, pq.StringArray(req.IdentityColumnNames))
+func (s *Store) CreateSnapshotRequest(ctx context.Context, req *snapshot.Request) error {
+	query := fmt.Sprintf(`INSERT INTO %s (schema_name, table_names, created_at, updated_at, status)
+	VALUES($1, $2,'now()','now()','requested')`, snapshotsTable())
+	_, err := s.conn.Exec(ctx, query, req.Snapshot.SchemaName, pq.StringArray(req.Snapshot.TableNames))
 	if err != nil {
 		return fmt.Errorf("error creating snapshot request: %w", err)
 	}
 	return nil
 }
 
-func (s *Store) UpdateSnapshotRequest(ctx context.Context, req *snapshot.Snapshot) error {
-	errStr := ""
-	if req.Error != nil {
-		errStr = req.Error.Error()
-	}
-	query := fmt.Sprintf(`UPDATE %s SET status = '%s', error = '%s', updated_at = 'now()'
-	WHERE schema_name = '%s' and table_name = '%s' and status != 'completed'`, snapshotsTable(), req.Status, errStr, req.SchemaName, req.TableName)
-	_, err := s.conn.Exec(ctx, query)
+func (s *Store) UpdateSnapshotRequest(ctx context.Context, req *snapshot.Request) error {
+	query := fmt.Sprintf(`UPDATE %s SET status = $1, errors = $2, updated_at = 'now()'
+	WHERE schema_name = $3 and table_names = $4 and status != 'completed'`, snapshotsTable())
+	_, err := s.conn.Exec(ctx, query, req.Status, req.Errors, req.Snapshot.SchemaName, pq.StringArray(req.Snapshot.TableNames))
 	if err != nil {
 		return fmt.Errorf("error updating snapshot request: %w", err)
 	}
 	return nil
 }
 
-func (s *Store) GetSnapshotRequests(ctx context.Context, status snapshot.Status) ([]*snapshot.Snapshot, error) {
-	query := fmt.Sprintf(`SELECT schema_name,table_name,identity_column_names,status FROM %s
+func (s *Store) GetSnapshotRequestsByStatus(ctx context.Context, status snapshot.Status) ([]*snapshot.Request, error) {
+	query := fmt.Sprintf(`SELECT schema_name,table_names,status,errors FROM %s
 	WHERE status = '%s' ORDER BY req_id ASC LIMIT %d`, snapshotsTable(), status, queryLimit)
 	rows, err := s.conn.Query(ctx, query)
 	if err != nil {
-		return nil, fmt.Errorf("error getting snapshot requests: %w", err)
+		return nil, fmt.Errorf("error getting snapshot requests by status: %w", err)
 	}
 	defer rows.Close()
 
-	snapshots := []*snapshot.Snapshot{}
+	snapshotRequests := []*snapshot.Request{}
 	for rows.Next() {
-		snapshot := &snapshot.Snapshot{}
-		if err := rows.Scan(&snapshot.SchemaName, &snapshot.TableName, &snapshot.IdentityColumnNames, &snapshot.Status); err != nil {
+		req := &snapshot.Request{}
+		if err := rows.Scan(&req.Snapshot.SchemaName, &req.Snapshot.TableNames, &req.Status, &req.Errors); err != nil {
 			return nil, fmt.Errorf("scanning snapshot row: %w", err)
 		}
 
-		snapshots = append(snapshots, snapshot)
+		snapshotRequests = append(snapshotRequests, req)
+	}
+
+	return snapshotRequests, nil
+}
+
+func (s *Store) GetSnapshotRequestsBySchema(ctx context.Context, schema string) ([]*snapshot.Request, error) {
+	query := fmt.Sprintf(`SELECT schema_name,table_names,status,errors FROM %s
+	WHERE schema_name = $1 ORDER BY req_id ASC LIMIT %d`, snapshotsTable(), queryLimit)
+	rows, err := s.conn.Query(ctx, query, schema)
+	if err != nil {
+		return nil, fmt.Errorf("error getting snapshot requests: %w", err)
+	}
+	defer rows.Close()
+
+	snapshotRequests := []*snapshot.Request{}
+	for rows.Next() {
+		req := &snapshot.Request{}
+		if err := rows.Scan(&req.Snapshot.SchemaName, &req.Snapshot.TableNames, &req.Status, &req.Errors); err != nil {
+			return nil, fmt.Errorf("scanning snapshot request row: %w", err)
+		}
+
+		snapshotRequests = append(snapshotRequests, req)
 	}
 
-	return snapshots, nil
+	return snapshotRequests, nil
 }
 
 func (s *Store) createTable(ctx context.Context) error {
 	createQuery := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s(
 	req_id SERIAL PRIMARY KEY,
 	schema_name TEXT,
-	table_name TEXT,
-	identity_column_names TEXT[],
+	table_names TEXT[],
 	created_at TIMESTAMP WITH TIME ZONE,
 	updated_at TIMESTAMP WITH TIME ZONE,
 	status TEXT CHECK (status IN ('requested', 'in progress', 'completed')),
-	error TEXT )`, snapshotsTable())
+	errors JSONB )`, snapshotsTable())
 	_, err := s.conn.Exec(ctx, createQuery)
 	if err != nil {
 		return fmt.Errorf("error creating snapshots postgres table: %w", err)
 	}
 
 	uniqueIndexQuery := fmt.Sprintf(`CREATE UNIQUE INDEX IF NOT EXISTS schema_table_status_unique_index
-	ON %s(schema_name,table_name) WHERE status != 'completed'`, snapshotsTable())
+	ON %s(schema_name,table_names) WHERE status != 'completed'`, snapshotsTable())
 	_, err = s.conn.Exec(ctx, uniqueIndexQuery)
 	if err != nil {
 		return fmt.Errorf("error creating unique index on snapshots postgres table: %w", err)
diff --git a/pkg/snapshot/store/snapshot_store.go b/pkg/snapshot/store/snapshot_store.go
index 28e5432..4f31bd1 100644
--- a/pkg/snapshot/store/snapshot_store.go
+++ b/pkg/snapshot/store/snapshot_store.go
@@ -9,9 +9,11 @@ import (
 )
 
 type Store interface {
-	CreateSnapshotRequest(context.Context, *snapshot.Snapshot) error
-	UpdateSnapshotRequest(context.Context, *snapshot.Snapshot) error
-	GetSnapshotRequests(ctx context.Context, status snapshot.Status) ([]*snapshot.Snapshot, error)
+	CreateSnapshotRequest(context.Context, *snapshot.Request) error
+	UpdateSnapshotRequest(context.Context, *snapshot.Request) error
+	GetSnapshotRequestsByStatus(ctx context.Context, status snapshot.Status) ([]*snapshot.Request, error)
+	GetSnapshotRequestsBySchema(ctx context.Context, schema string) ([]*snapshot.Request, error)
+	Close() error
 }
 
 const (


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 e771b3d..076e53d 100644
--- a/pkg/snapshot/store/postgres/pg_snapshot_store_test.go
+++ b/pkg/snapshot/store/postgres/pg_snapshot_store_test.go
@@ -18,10 +18,12 @@ import (
 func TestStore_CreateSnapshotRequest(t *testing.T) {
 	t.Parallel()
 
-	testSnapshot := &snapshot.Snapshot{
-		SchemaName:          "test-schema",
-		TableName:           "test-table",
-		IdentityColumnNames: []string{"id"},
+	testSnapshot := snapshot.Snapshot{
+		SchemaName: "test-schema",
+		TableNames: []string{"test-table-1", "test-table-2"},
+	}
+	testSnapshotRequest := &snapshot.Request{
+		Snapshot: testSnapshot,
 	}
 
 	errTest := errors.New("oh noes")
@@ -36,10 +38,10 @@ func TestStore_CreateSnapshotRequest(t *testing.T) {
 			name: "ok",
 			querier: &postgresmocks.Querier{
 				ExecFn: func(ctx context.Context, _ uint, s string, a ...any) (postgres.CommandTag, error) {
-					wantQuery := fmt.Sprintf(`INSERT INTO %s (schema_name, table_name, identity_column_names, created_at, updated_at, status)
-	VALUES($1, $2, $3,'now()','now()','requested')`, snapshotsTable())
+					wantQuery := fmt.Sprintf(`INSERT INTO %s (schema_name, table_names, created_at, updated_at, status)
+	VALUES($1, $2,'now()','now()','requested')`, snapshotsTable())
 					require.Equal(t, wantQuery, s)
-					wantAttr := []any{testSnapshot.SchemaName, testSnapshot.TableName, pq.StringArray(testSnapshot.IdentityColumnNames)}
+					wantAttr := []any{testSnapshot.SchemaName, pq.StringArray(testSnapshot.TableNames)}
 					require.Equal(t, wantAttr, a)
 					return postgres.CommandTag{}, nil
 				},
@@ -67,7 +69,7 @@ func TestStore_CreateSnapshotRequest(t *testing.T) {
 			store := &Store{
 				conn: tc.querier,
 			}
-			err := store.CreateSnapshotRequest(context.Background(), testSnapshot)
+			err := store.CreateSnapshotRequest(context.Background(), testSnapshotRequest)
 			require.ErrorIs(t, err, tc.wantErr)
 		})
 	}
@@ -77,52 +79,55 @@ func TestStore_UpdateSnapshotRequest(t *testing.T) {
 	t.Parallel()
 
 	testSnapshot := snapshot.Snapshot{
-		SchemaName:          "test-schema",
-		TableName:           "test-table",
-		IdentityColumnNames: []string{"id"},
-		Status:              snapshot.StatusInProgress,
+		SchemaName: "test-schema",
+		TableNames: []string{"test-table-1", "test-table-2"},
+	}
+	testSnapshotRequest := snapshot.Request{
+		Snapshot: testSnapshot,
+		Status:   snapshot.StatusInProgress,
 	}
 
 	errTest := errors.New("oh noes")
+	testSnapshotErr := &snapshot.Errors{Snapshot: errTest}
 
 	tests := []struct {
 		name    string
 		querier postgres.Querier
-		req     *snapshot.Snapshot
+		req     *snapshot.Request
 
 		wantErr error
 	}{
 		{
-			name: "ok - update without error string",
+			name: "ok - update without error",
 			querier: &postgresmocks.Querier{
 				ExecFn: func(ctx context.Context, _ uint, s string, a ...any) (postgres.CommandTag, error) {
-					wantQuery := fmt.Sprintf(`UPDATE %s SET status = '%s', error = '%s', updated_at = 'now()'
-	WHERE schema_name = '%s' and table_name = '%s' and status != 'completed'`,
-						snapshotsTable(), snapshot.StatusInProgress, "", testSnapshot.SchemaName, testSnapshot.TableName)
+					wantQuery := fmt.Sprintf(`UPDATE %s SET status = $1, errors = $2, updated_at = 'now()'
+	WHERE schema_name = $3 and table_names = $4 and status != 'completed'`,
+						snapshotsTable())
 					require.Equal(t, wantQuery, s)
-					require.Empty(t, a)
+					require.Equal(t, []any{snapshot.StatusInProgress, (*snapshot.Errors)(nil), testSnapshot.SchemaName, pq.StringArray(testSnapshot.TableNames)}, a)
 					return postgres.CommandTag{}, nil
 				},
 			},
-			req: &testSnapshot,
+			req: &testSnapshotRequest,
 
 			wantErr: nil,
 		},
 		{
-			name: "ok - update with error string",
+			name: "ok - update with error",
 			querier: &postgresmocks.Querier{
 				ExecFn: func(ctx context.Context, _ uint, s string, a ...any) (postgres.CommandTag, error) {
-					wantQuery := fmt.Sprintf(`UPDATE %s SET status = '%s', error = '%s', updated_at = 'now()'
-	WHERE schema_name = '%s' and table_name = '%s' and status != 'completed'`,
-						snapshotsTable(), snapshot.StatusInProgress, errTest.Error(), testSnapshot.SchemaName, testSnapshot.TableName)
+					wantQuery := fmt.Sprintf(`UPDATE %s SET status = $1, errors = $2, updated_at = 'now()'
+	WHERE schema_name = $3 and table_names = $4 and status != 'completed'`,
+						snapshotsTable())
 					require.Equal(t, wantQuery, s)
-					require.Empty(t, a)
+					require.Equal(t, []any{snapshot.StatusInProgress, testSnapshotErr, testSnapshot.SchemaName, pq.StringArray(testSnapshot.TableNames)}, a)
 					return postgres.CommandTag{}, nil
 				},
 			},
-			req: func() *snapshot.Snapshot {
-				s := testSnapshot
-				s.Error = errTest
+			req: func() *snapshot.Request {
+				s := testSnapshotRequest
+				s.Errors = testSnapshotErr
 				return &s
 			}(),
 
@@ -135,7 +140,7 @@ func TestStore_UpdateSnapshotRequest(t *testing.T) {
 					return postgres.CommandTag{}, errTest
 				},
 			},
-			req: &testSnapshot,
+			req: &testSnapshotRequest,
 
 			wantErr: errTest,
 		},
@@ -155,30 +160,159 @@ func TestStore_UpdateSnapshotRequest(t *testing.T) {
 	}
 }
 
-func TestStore_GetSnapshotRequests(t *testing.T) {
+func TestStore_GetSnapshotRequestsBySchema(t *testing.T) {
 	t.Parallel()
 
-	testSnapshot := &snapshot.Snapshot{
-		SchemaName:          "test-schema",
-		TableName:           "test-table",
-		IdentityColumnNames: []string{"id"},
-		Status:              snapshot.StatusInProgress,
+	testSnapshot := snapshot.Snapshot{
+		SchemaName: "test-schema",
+		TableNames: []string{"test-table-1", "test-table-2"},
 	}
+	errTest := errors.New("oh noes")
+	testSnapshotRequest := snapshot.Request{
+		Snapshot: testSnapshot,
+		Status:   snapshot.StatusInProgress,
+		Errors:   &snapshot.Errors{Snapshot: errTest},
+	}
+
+	tests := []struct {
+		name    string
+		querier postgres.Querier
+
+		wantRequests []*snapshot.Request
+		wantErr      error
+	}{
+		{
+			name: "ok - no results",
+			querier: &postgresmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (postgres.Rows, error) {
+					wantQuery := fmt.Sprintf(`SELECT schema_name,table_names,status,errors FROM %s
+	WHERE schema_name = $1 ORDER BY req_id ASC LIMIT %d`, snapshotsTable(), queryLimit)
+					require.Equal(t, []any{testSnapshot.SchemaName}, args)
+					require.Equal(t, wantQuery, query)
+					return &postgresmocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(_ uint) bool { return false },
+					}, nil
+				},
+			},
+
+			wantRequests: []*snapshot.Request{},
+			wantErr:      nil,
+		},
+		{
+			name: "ok",
+			querier: &postgresmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (postgres.Rows, error) {
+					wantQuery := fmt.Sprintf(`SELECT schema_name,table_names,status,errors FROM %s
+	WHERE schema_name = $1 ORDER BY req_id ASC LIMIT %d`, snapshotsTable(), queryLimit)
+					require.Equal(t, []any{testSnapshot.SchemaName}, args)
+					require.Equal(t, wantQuery, query)
+					return &postgresmocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							require.Len(t, dest, 4)
+							schemaName, ok := dest[0].(*string)
+							require.True(t, ok)
+							*schemaName = testSnapshot.SchemaName
+
+							tableNames, ok := dest[1].(*[]string)
+							require.True(t, ok)
+							*tableNames = testSnapshot.TableNames
+
+							status, ok := dest[2].(*snapshot.Status)
+							require.True(t, ok)
+							*status = testSnapshotRequest.Status
+
+							errs, ok := dest[3].(**snapshot.Errors)
+							require.True(t, ok)
+							*errs = testSnapshotRequest.Errors
+							return nil
+						},
+					}, nil
+				},
+			},
 
+			wantRequests: []*snapshot.Request{
+				&testSnapshotRequest,
+			},
+			wantErr: nil,
+		},
+		{
+			name: "error - querying",
+			querier: &postgresmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (postgres.Rows, error) {
+					return nil, errTest
+				},
+			},
+
+			wantRequests: nil,
+			wantErr:      errTest,
+		},
+		{
+			name: "error - scanning row",
+			querier: &postgresmocks.Querier{
+				QueryFn: func(ctx context.Context, query string, args ...any) (postgres.Rows, error) {
+					wantQuery := fmt.Sprintf(`SELECT schema_name,table_names,status,errors FROM %s
+	WHERE schema_name = $1 ORDER BY req_id ASC LIMIT %d`, snapshotsTable(), queryLimit)
+					require.Equal(t, []any{testSnapshot.SchemaName}, args)
+					require.Equal(t, wantQuery, query)
+					return &postgresmocks.Rows{
+						CloseFn: func() {},
+						NextFn:  func(i uint) bool { return i == 1 },
+						ScanFn: func(dest ...any) error {
+							return errTest
+						},
+					}, nil
+				},
+			},
+
+			wantRequests: nil,
+			wantErr:      errTest,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			store := &Store{
+				conn: tc.querier,
+			}
+			requests, err := store.GetSnapshotRequestsBySchema(context.Background(), testSnapshot.SchemaName)
+			require.ErrorIs(t, err, tc.wantErr)
+			require.Equal(t, tc.wantRequests, requests)
+		})
+	}
+}
+
+func TestStore_GetSnapshotRequestsByStatus(t *testing.T) {
+	t.Parallel()
+
+	testSnapshot := snapshot.Snapshot{
+		SchemaName: "test-schema",
+		TableNames: []string{"test-table-1", "test-table-2"},
+	}
 	errTest := errors.New("oh noes")
+	testSnapshotRequest := snapshot.Request{
+		Snapshot: testSnapshot,
+		Status:   snapshot.StatusInProgress,
+		Errors:   &snapshot.Errors{Snapshot: errTest},
+	}
 
 	tests := []struct {
 		name    string
 		querier postgres.Querier
 
-		wantSnapshots []*snapshot.Snapshot
-		wantErr       error
+		wantRequests []*snapshot.Request
+		wantErr      error
 	}{
 		{
 			name: "ok - no results",
 			querier: &postgresmocks.Querier{
 				QueryFn: func(ctx context.Context, query string, args ...any) (postgres.Rows, error) {
-					wantQuery := fmt.Sprintf(`SELECT schema_name,table_name,identity_column_names,status FROM %s
+					wantQuery := fmt.Sprintf(`SELECT schema_name,table_names,status,errors FROM %s
 	WHERE status = '%s' ORDER BY req_id ASC LIMIT %d`, snapshotsTable(), snapshot.StatusInProgress, queryLimit)
 					require.Equal(t, wantQuery, query)
 					return &postgresmocks.Rows{
@@ -188,14 +322,14 @@ func TestStore_GetSnapshotRequests(t *testing.T) {
 				},
 			},
 
-			wantSnapshots: []*snapshot.Snapshot{},
-			wantErr:       nil,
+			wantRequests: []*snapshot.Request{},
+			wantErr:      nil,
 		},
 		{
 			name: "ok",
 			querier: &postgresmocks.Querier{
 				QueryFn: func(ctx context.Context, query string, args ...any) (postgres.Rows, error) {
-					wantQuery := fmt.Sprintf(`SELECT schema_name,table_name,identity_column_names,status FROM %s
+					wantQuery := fmt.Sprintf(`SELECT schema_name,table_names,status,errors FROM %s
 	WHERE status = '%s' ORDER BY req_id ASC LIMIT %d`, snapshotsTable(), snapshot.StatusInProgress, queryLimit)
 					require.Equal(t, wantQuery, query)
 					return &postgresmocks.Rows{
@@ -207,25 +341,25 @@ func TestStore_GetSnapshotRequests(t *testing.T) {
 							require.True(t, ok)
 							*schemaName = testSnapshot.SchemaName
 
-							tableName, ok := dest[1].(*string)
+							tableNames, ok := dest[1].(*[]string)
 							require.True(t, ok)
-							*tableName = testSnapshot.TableName
+							*tableNames = testSnapshot.TableNames
 
-							idColumns, ok := dest[2].(*[]string)
+							status, ok := dest[2].(*snapshot.Status)
 							require.True(t, ok)
-							*idColumns = testSnapshot.IdentityColumnNames
+							*status = testSnapshotRequest.Status
 
-							status, ok := dest[3].(*snapshot.Status)
+							errs, ok := dest[3].(**snapshot.Errors)
 							require.True(t, ok)
-							*status = testSnapshot.Status
+							*errs = testSnapshotRequest.Errors
 							return nil
 						},
 					}, nil
 				},
 			},
 
-			wantSnapshots: []*snapshot.Snapshot{
-				testSnapshot,
+			wantRequests: []*snapshot.Request{
+				&testSnapshotRequest,
 			},
 			wantErr: nil,
 		},
@@ -237,14 +371,14 @@ func TestStore_GetSnapshotRequests(t *testing.T) {
 				},
 			},
 
-			wantSnapshots: nil,
-			wantErr:       errTest,
+			wantRequests: nil,
+			wantErr:      errTest,
 		},
 		{
 			name: "error - scanning row",
 			querier: &postgresmocks.Querier{
 				QueryFn: func(ctx context.Context, query string, args ...any) (postgres.Rows, error) {
-					wantQuery := fmt.Sprintf(`SELECT schema_name,table_name,identity_column_names,status FROM %s
+					wantQuery := fmt.Sprintf(`SELECT schema_name,table_names,status,errors FROM %s
 	WHERE status = '%s' ORDER BY req_id ASC LIMIT %d`, snapshotsTable(), snapshot.StatusInProgress, queryLimit)
 					require.Equal(t, wantQuery, query)
 					return &postgresmocks.Rows{
@@ -257,8 +391,8 @@ func TestStore_GetSnapshotRequests(t *testing.T) {
 				},
 			},
 
-			wantSnapshots: nil,
-			wantErr:       errTest,
+			wantRequests: nil,
+			wantErr:      errTest,
 		},
 	}
 
@@ -270,9 +404,9 @@ func TestStore_GetSnapshotRequests(t *testing.T) {
 			store := &Store{
 				conn: tc.querier,
 			}
-			snapshots, err := store.GetSnapshotRequests(context.Background(), snapshot.StatusInProgress)
+			requests, err := store.GetSnapshotRequestsByStatus(context.Background(), snapshot.StatusInProgress)
 			require.ErrorIs(t, err, tc.wantErr)
-			require.Equal(t, tc.wantSnapshots, snapshots)
+			require.Equal(t, tc.wantRequests, requests)
 		})
 	}
 }
@@ -297,17 +431,16 @@ func TestStore_createTable(t *testing.T) {
 						wantQuery := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s(
 	req_id SERIAL PRIMARY KEY,
 	schema_name TEXT,
-	table_name TEXT,
-	identity_column_names TEXT[],
+	table_names TEXT[],
 	created_at TIMESTAMP WITH TIME ZONE,
 	updated_at TIMESTAMP WITH TIME ZONE,
 	status TEXT CHECK (status IN ('requested', 'in progress', 'completed')),
-	error TEXT )`, snapshotsTable())
+	errors JSONB )`, snapshotsTable())
 						require.Equal(t, wantQuery, s)
 						require.Empty(t, a)
 					case 2:
 						wantQuery := fmt.Sprintf(`CREATE UNIQUE INDEX IF NOT EXISTS schema_table_status_unique_index
-	ON %s(schema_name,table_name) WHERE status != 'completed'`, snapshotsTable())
+	ON %s(schema_name,table_names) WHERE status != 'completed'`, snapshotsTable())
 						require.Equal(t, wantQuery, s)
 						require.Empty(t, a)
 					default: