fix: add schema mode verification to UnifiedSchemaHash test

- Add SchemaModeForTesting() method to all datastores returning (SchemaMode, error)
- Add SchemaModeForTesting() to indexcheck and relationshipintegrity proxies
- Update UnifiedSchemaHashTest to explicitly check schema mode from datastore
- Test now correctly validates:
  - Hash is NOT written in ReadLegacyWriteLegacy mode
  - Hash IS written and correct in ReadLegacyWriteBoth, ReadNewWriteBoth, ReadNewWriteNew
- Test fails explicitly if datastore doesn't implement SchemaModeForTesting()
- Fixes test failures in default mode for MySQL, CRDB, and Spanner

Author: josephschorr

internal/datastore/crdb/crdb.go

@@ -550,6 +550,11 @@ func (cds *crdbDatastore) SchemaHashWatcherForTesting() datastore.SingleStoreSch
 	return newCRDBSchemaHashWatcher(cds.readPool, schemaChangefeedFormat)
 }
 
+// SchemaModeForTesting returns the current schema mode for testing purposes.
+func (cds *crdbDatastore) SchemaModeForTesting() (options.SchemaMode, error) {
+	return cds.schemaMode, nil
+}
+
 type crdbSchemaHashReaderForTesting struct {
 	query *pool.RetryPool
 }

internal/datastore/memdb/memdb.go

@@ -426,6 +426,12 @@ func (mdb *memdbDatastore) SchemaHashWatcherForTesting() datastore.SingleStoreSc
 	return newMemdbSchemaHashWatcher(mdb)
 }
 
+// SchemaModeForTesting returns the current schema mode for testing purposes.
+// MemDB always operates in a mode equivalent to ReadNewWriteNew.
+func (mdb *memdbDatastore) SchemaModeForTesting() (options.SchemaMode, error) {
+	return options.SchemaModeReadNewWriteNew, nil
+}
+
 type memdbSchemaHashReaderForTesting struct {
 	db *memdbDatastore
 }

internal/datastore/mysql/datastore.go

@@ -597,6 +597,11 @@ func (mds *mysqlDatastore) SchemaHashWatcherForTesting() datastore.SingleStoreSc
 	return newMySQLSchemaHashWatcher(mds.db, mds.driver.SchemaRevision())
 }
 
+// SchemaModeForTesting returns the current schema mode for testing purposes.
+func (mds *mysqlDatastore) SchemaModeForTesting() (dsoptions.SchemaMode, error) {
+	return mds.schemaMode, nil
+}
+
 type mysqlSchemaHashReaderForTesting struct {
 	db                      *sql.DB
 	schemaRevisionTableName string

internal/datastore/postgres/postgres.go

@@ -705,6 +705,11 @@ func (pgd *pgDatastore) SchemaHashWatcherForTesting() datastore.SingleStoreSchem
 	return newPGSchemaHashWatcher(queryFuncs)
 }
 
+// SchemaModeForTesting returns the current schema mode for testing purposes.
+func (pgd *pgDatastore) SchemaModeForTesting() (dsoptions.SchemaMode, error) {
+	return pgd.schemaMode, nil
+}
+
 type pgSchemaHashReaderForTesting struct {
 	query pgxcommon.DBFuncQuerier
 }

internal/datastore/proxy/indexcheck/indexcheck.go

@@ -129,6 +129,19 @@ func (p *indexcheckingProxy) SchemaHashWatcherForTesting() datastore.SingleStore
 	return nil
 }
 
+// SchemaModeForTesting returns the current schema mode for testing purposes.
+// This delegates to the underlying datastore if it supports the test interface.
+func (p *indexcheckingProxy) SchemaModeForTesting() (options.SchemaMode, error) {
+	type schemaModeProvider interface {
+		SchemaModeForTesting() (options.SchemaMode, error)
+	}
+
+	if provider, ok := p.delegate.(schemaModeProvider); ok {
+		return provider.SchemaModeForTesting()
+	}
+	return options.SchemaModeReadLegacyWriteLegacy, fmt.Errorf("delegate datastore does not implement SchemaModeForTesting()")
+}
+
 type indexcheckingReader struct {
 	parent   datastore.SQLDatastore
 	delegate datastore.Reader

internal/datastore/proxy/relationshipintegrity.go

@@ -357,6 +357,31 @@ func (r *relationshipIntegrityProxy) SchemaHashWatcherForTesting() datastore.Sin
 	return nil
 }
 
+// SchemaModeForTesting returns the current schema mode for testing purposes.
+// This delegates to the underlying datastore if it supports the test interface.
+func (r *relationshipIntegrityProxy) SchemaModeForTesting() (options.SchemaMode, error) {
+	type provider interface {
+		SchemaModeForTesting() (options.SchemaMode, error)
+	}
+
+	// Check delegate directly
+	if p, ok := r.ds.(provider); ok {
+		return p.SchemaModeForTesting()
+	}
+
+	// Try unwrapping if delegate is itself a proxy
+	type unwrapper interface {
+		Unwrap() datastore.Datastore
+	}
+	if u, ok := r.ds.(unwrapper); ok {
+		if p, ok := u.Unwrap().(provider); ok {
+			return p.SchemaModeForTesting()
+		}
+	}
+
+	return options.SchemaModeReadLegacyWriteLegacy, fmt.Errorf("delegate datastore does not implement SchemaModeForTesting()")
+}
+
 func (r *relationshipIntegrityProxy) Unwrap() datastore.Datastore {
 	return r.ds
 }

internal/datastore/spanner/spanner.go

@@ -491,6 +491,11 @@ func (sd *spannerDatastore) SchemaHashWatcherForTesting() datastore.SingleStoreS
 	return newSpannerSchemaHashWatcher(sd.client)
 }
 
+// SchemaModeForTesting returns the current schema mode for testing purposes.
+func (sd *spannerDatastore) SchemaModeForTesting() (dsoptions.SchemaMode, error) {
+	return sd.schemaMode, nil
+}
+
 type spannerSchemaHashReaderForTesting struct {
 	client *spanner.Client
 }

pkg/datastore/test/unifiedschema.go

@@ -954,15 +954,34 @@ func UnifiedSchemaHashTest(t *testing.T, tester DatastoreTester) {
 	})
 	require.NoError(err)
 
+	// Get the schema mode from the datastore to determine expected behavior
+	schemaModeProvider, ok := ds.(interface {
+		SchemaModeForTesting() (options.SchemaMode, error)
+	})
+	require.True(ok, "datastore must implement SchemaModeForTesting() for this test")
+	schemaMode, err := schemaModeProvider.SchemaModeForTesting()
+	require.NoError(err, "failed to get schema mode from datastore")
+
+	// Determine if this mode should write the unified schema hash
+	// The hash is written in all modes that write to the unified schema
+	shouldWriteHash := schemaMode == options.SchemaModeReadLegacyWriteBoth ||
+		schemaMode == options.SchemaModeReadNewWriteBoth ||
+		schemaMode == options.SchemaModeReadNewWriteNew
+
 	// Read the schema hash from schema_revision table
-	hash, err := readerImpl.ReadSchemaHash(ctx)
-	require.NoError(err)
-	require.NotEmpty(hash, "schema hash should not be empty")
+	hash, hashErr := readerImpl.ReadSchemaHash(ctx)
 
-	// Verify the hash is correct by computing the expected hash
-	// The hash is computed from sorted definitions (matching datastore behavior)
-	expectedHash := computeExpectedSchemaHash(t, testSchemaDefinitions)
-	require.Equal(expectedHash, hash, "schema hash should match computed hash of sorted schema text")
+	if shouldWriteHash {
+		// Hash MUST be present and correct
+		require.NoError(hashErr, "schema hash should be present in mode %s", schemaMode)
+		require.NotEmpty(hash, "schema hash should not be empty")
+		expectedHash := computeExpectedSchemaHash(t, testSchemaDefinitions)
+		require.Equal(expectedHash, hash, "schema hash should match computed hash of sorted schema text")
+	} else {
+		// Hash should NOT be present in ReadLegacyWriteLegacy mode
+		require.True(errors.Is(hashErr, datastore.ErrSchemaNotFound),
+			"expected ErrSchemaNotFound in mode %s, got: %v", schemaMode, hashErr)
+	}
 
 	// Update the schema
 	updatedSchemaText, _, err := generator.GenerateSchema(updatedSchemaDefinitions)
@@ -982,15 +1001,23 @@ func UnifiedSchemaHashTest(t *testing.T, tester DatastoreTester) {
 	})
 	require.NoError(err)
 
-	// Read the updated schema hash
-	updatedHash, err := hashReader.SchemaHashReaderForTesting().ReadSchemaHash(ctx)
-	require.NoError(err)
-	require.NotEmpty(updatedHash, "updated schema hash should not be empty")
-	require.NotEqual(hash, updatedHash, "schema hash should change after update")
+	// Read the updated schema hash and verify based on mode
+	updatedHash, updatedHashErr := hashReader.SchemaHashReaderForTesting().ReadSchemaHash(ctx)
 
-	// Verify the updated hash is correct by computing the expected hash
-	expectedUpdatedHash := computeExpectedSchemaHash(t, updatedSchemaDefinitions)
-	require.Equal(expectedUpdatedHash, updatedHash, "updated schema hash should match computed hash of sorted updated schema text")
+	if shouldWriteHash {
+		// Hash MUST be present, correct, and different from the first hash
+		require.NoError(updatedHashErr, "updated schema hash should be present in mode %s", schemaMode)
+		require.NotEmpty(updatedHash, "updated schema hash should not be empty")
+		require.NotEqual(hash, updatedHash, "schema hash should change after update")
+
+		// Verify the updated hash is correct
+		expectedUpdatedHash := computeExpectedSchemaHash(t, updatedSchemaDefinitions)
+		require.Equal(expectedUpdatedHash, updatedHash, "updated schema hash should match computed hash of sorted updated schema text")
+	} else {
+		// Hash should still NOT be present in ReadLegacyWriteLegacy mode
+		require.True(errors.Is(updatedHashErr, datastore.ErrSchemaNotFound),
+			"expected ErrSchemaNotFound after update in mode %s, got: %v", schemaMode, updatedHashErr)
+	}
 }
 
 // UnifiedSchemaHashWatchTest tests the schema hash watcher functionality