Repository: gofr-dev/gofr

Issue:
Support Migrations for DGraph


Discussion:

Solution patch:
diff --git a/docs/datasources/getting-started/page.md b/docs/datasources/getting-started/page.md
index cdb531d5..ae96b652 100644
--- a/docs/datasources/getting-started/page.md
+++ b/docs/datasources/getting-started/page.md
@@ -93,7 +93,7 @@ as unnecessary database drivers are not being compiled and added to the build.
 - ✅
 - ✅
 - ✅
--
+- ✅
 
 ---
 
diff --git a/pkg/gofr/container/datasources.go b/pkg/gofr/container/datasources.go
index ad4b3466..344a7516 100644
--- a/pkg/gofr/container/datasources.go
+++ b/pkg/gofr/container/datasources.go
@@ -384,6 +384,32 @@ type SolrProvider interface {
 
 // Dgraph defines the methods for interacting with a Dgraph database.
 type Dgraph interface {
+	// ApplySchema applies or updates the complete database schema.
+	// Parameters:
+	// - ctx: Context for request cancellation and timeouts
+	// - schema: Schema definition in Dgraph Schema Definition Language (SDL) format
+	// Returns:
+	// - error: An error if the schema application fails
+	ApplySchema(ctx context.Context, schema string) error
+
+	// AddOrUpdateField atomically creates or updates a single field definition.
+	// Parameters:
+	// - ctx: Context for request cancellation and timeouts
+	// - fieldName: Name of the field/predicate to create or update
+	// - fieldType: Dgraph data type (e.g., string, int, datetime)
+	// - directives: Space-separated Dgraph directives (e.g., "@index(hash) @upsert")
+	// Returns:
+	// - error: An error if the field operation fails
+	AddOrUpdateField(ctx context.Context, fieldName, fieldType, directives string) error
+
+	// DropField permanently removes a field/predicate and all its associated data.
+	// Parameters:
+	// - ctx: Context for request cancellation and timeouts
+	// - fieldName: Name of the field/predicate to remove
+	// Returns:
+	// - error: An error if the field removal fails
+	DropField(ctx context.Context, fieldName string) error
+
 	// Query executes a read-only query in the Dgraph database and returns the result.
 	// Parameters:
 	// - ctx: The context for the query, used for controlling timeouts, cancellation, etc.
diff --git a/pkg/gofr/container/mock_datasources.go b/pkg/gofr/container/mock_datasources.go
index f36fe2fa..8cf39023 100644
--- a/pkg/gofr/container/mock_datasources.go
+++ b/pkg/gofr/container/mock_datasources.go
@@ -11040,6 +11040,20 @@ func (m *MockDgraph) EXPECT() *MockDgraphMockRecorder {
 	return m.recorder
 }
 
+// AddOrUpdateField mocks base method.
+func (m *MockDgraph) AddOrUpdateField(ctx context.Context, fieldName, fieldType, directives string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "AddOrUpdateField", ctx, fieldName, fieldType, directives)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// AddOrUpdateField indicates an expected call of AddOrUpdateField.
+func (mr *MockDgraphMockRecorder) AddOrUpdateField(ctx, fieldName, fieldType, directives any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddOrUpdateField", reflect.TypeOf((*MockDgraph)(nil).AddOrUpdateField), ctx, fieldName, fieldType, directives)
+}
+
 // Alter mocks base method.
 func (m *MockDgraph) Alter(ctx context.Context, op any) error {
 	m.ctrl.T.Helper()
@@ -11054,6 +11068,34 @@ func (mr *MockDgraphMockRecorder) Alter(ctx, op any) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Alter", reflect.TypeOf((*MockDgraph)(nil).Alter), ctx, op)
 }
 
+// ApplySchema mocks base method.
+func (m *MockDgraph) ApplySchema(ctx context.Context, schema string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ApplySchema", ctx, schema)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// ApplySchema indicates an expected call of ApplySchema.
+func (mr *MockDgraphMockRecorder) ApplySchema(ctx, schema any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ApplySchema", reflect.TypeOf((*MockDgraph)(nil).ApplySchema), ctx, schema)
+}
+
+// DropField mocks base method.
+func (m *MockDgraph) DropField(ctx context.Context, fieldName string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DropField", ctx, fieldName)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// DropField indicates an expected call of DropField.
+func (mr *MockDgraphMockRecorder) DropField(ctx, fieldName any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropField", reflect.TypeOf((*MockDgraph)(nil).DropField), ctx, fieldName)
+}
+
 // HealthCheck mocks base method.
 func (m *MockDgraph) HealthCheck(arg0 context.Context) (any, error) {
 	m.ctrl.T.Helper()
@@ -11166,6 +11208,20 @@ func (m *MockDgraphProvider) EXPECT() *MockDgraphProviderMockRecorder {
 	return m.recorder
 }
 
+// AddOrUpdateField mocks base method.
+func (m *MockDgraphProvider) AddOrUpdateField(ctx context.Context, fieldName, fieldType, directives string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "AddOrUpdateField", ctx, fieldName, fieldType, directives)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// AddOrUpdateField indicates an expected call of AddOrUpdateField.
+func (mr *MockDgraphProviderMockRecorder) AddOrUpdateField(ctx, fieldName, fieldType, directives any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddOrUpdateField", reflect.TypeOf((*MockDgraphProvider)(nil).AddOrUpdateField), ctx, fieldName, fieldType, directives)
+}
+
 // Alter mocks base method.
 func (m *MockDgraphProvider) Alter(ctx context.Context, op any) error {
 	m.ctrl.T.Helper()
@@ -11180,6 +11236,20 @@ func (mr *MockDgraphProviderMockRecorder) Alter(ctx, op any) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Alter", reflect.TypeOf((*MockDgraphProvider)(nil).Alter), ctx, op)
 }
 
+// ApplySchema mocks base method.
+func (m *MockDgraphProvider) ApplySchema(ctx context.Context, schema string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ApplySchema", ctx, schema)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// ApplySchema indicates an expected call of ApplySchema.
+func (mr *MockDgraphProviderMockRecorder) ApplySchema(ctx, schema any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ApplySchema", reflect.TypeOf((*MockDgraphProvider)(nil).ApplySchema), ctx, schema)
+}
+
 // Connect mocks base method.
 func (m *MockDgraphProvider) Connect() {
 	m.ctrl.T.Helper()
@@ -11192,6 +11262,20 @@ func (mr *MockDgraphProviderMockRecorder) Connect() *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Connect", reflect.TypeOf((*MockDgraphProvider)(nil).Connect))
 }
 
+// DropField mocks base method.
+func (m *MockDgraphProvider) DropField(ctx context.Context, fieldName string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DropField", ctx, fieldName)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// DropField indicates an expected call of DropField.
+func (mr *MockDgraphProviderMockRecorder) DropField(ctx, fieldName any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DropField", reflect.TypeOf((*MockDgraphProvider)(nil).DropField), ctx, fieldName)
+}
+
 // HealthCheck mocks base method.
 func (m *MockDgraphProvider) HealthCheck(arg0 context.Context) (any, error) {
 	m.ctrl.T.Helper()
diff --git a/pkg/gofr/datasource/dgraph/dgraph.go b/pkg/gofr/datasource/dgraph/dgraph.go
index 33906550..af375686 100644
--- a/pkg/gofr/datasource/dgraph/dgraph.go
+++ b/pkg/gofr/datasource/dgraph/dgraph.go
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/dgraph-io/dgo/v210"
@@ -40,6 +41,8 @@ var (
 	errInvalidMutation   = errors.New("invalid mutation type")
 	errInvalidOperation  = errors.New("invalid operation type")
 	errHealthCheckFailed = errors.New("dgraph health check failed")
+	errEmptySchema       = errors.New("schema cannot be empty")
+	errEmptyField        = errors.New("field name cannot be empty")
 )
 
 // New creates a new Dgraph client with the given configuration.
@@ -102,6 +105,39 @@ func (d *Client) UseTracer(tracer any) {
 	}
 }
 
+// ApplySchema applies a validated Dgraph schema.
+// The provided schema must be non-empty; otherwise, an error is returned.
+func (d *Client) ApplySchema(ctx context.Context, schema string) error {
+	if schema == "" {
+		return errEmptySchema
+	}
+
+	op := &api.Operation{Schema: schema}
+
+	return d.Alter(ctx, op)
+}
+
+// AddOrUpdateField creates or updates a field (predicate) definition in the Dgraph schema.
+// It adds a new field if it does not exist, or updates the field's type and/or directives if it does.
+func (d *Client) AddOrUpdateField(ctx context.Context, fieldName, fieldType, directives string) error {
+	if strings.TrimSpace(fieldName) == "" {
+		return errEmptyField
+	}
+
+	// Construct the field definition in the format: "<fieldName>: <fieldType> <directives>."
+	newDef := fmt.Sprintf("%s: %s %s.", fieldName, fieldType, directives)
+
+	return d.ApplySchema(ctx, newDef)
+}
+
+// DropField removes a specific field (predicate) and its associated data from the Dgraph schema.
+// It does so by creating an operation that instructs Dgraph to drop the given attribute.
+// An error is returned if the operation fails.
+func (d *Client) DropField(ctx context.Context, fieldName string) error {
+	op := &api.Operation{DropAttr: fieldName}
+	return d.Alter(ctx, op)
+}
+
 // Query executes a read-only query in the Dgraph database and returns the result.
 func (d *Client) Query(ctx context.Context, query string) (any, error) {
 	start := time.Now()
diff --git a/pkg/gofr/migration/datasource.go b/pkg/gofr/migration/datasource.go
index 858355b0..c71c96cd 100644
--- a/pkg/gofr/migration/datasource.go
+++ b/pkg/gofr/migration/datasource.go
@@ -15,6 +15,7 @@ type Datasource struct {
 	Mongo      Mongo
 	ArangoDB   ArangoDB
 	SurrealDB  SurrealDB
+	DGraph     DGraph
 }
 
 // It is a base implementation for migration manager, on this other database drivers have been wrapped.
diff --git a/pkg/gofr/migration/dgraph.go b/pkg/gofr/migration/dgraph.go
new file mode 100644
index 00000000..aa84da54
--- /dev/null
+++ b/pkg/gofr/migration/dgraph.go
@@ -0,0 +1,180 @@
+package migration
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	"github.com/dgraph-io/dgo/v210/protos/api"
+
+	"gofr.dev/pkg/gofr/container"
+)
+
+// dgraphDS is the adapter struct that implements migration operations.
+type dgraphDS struct {
+	client DGraph
+}
+
+// dgraphMigrator struct implements the migrator interface.
+type dgraphMigrator struct {
+	dgraphDS
+	migrator
+}
+
+const (
+	// dgraphSchema defines the migration schema with fully qualified predicate names.
+	dgraphSchema = `
+		migrations.version: int @index(int) .
+		migrations.method: string .
+		migrations.start_time: datetime .
+		migrations.duration: int .
+		type Migration {
+			migrations.version
+			migrations.method
+			migrations.start_time
+			migrations.duration
+		}
+	`
+
+	// getLastMigrationQuery fetches the most recent migration version.
+	getLastMigrationQuery = `
+		{
+			migrations(func: type(Migration), orderdesc: migrations.version, first: 1) {
+				migrations.version
+			}
+		}
+	`
+)
+
+// apply creates a new dgraphMigrator.
+func (ds dgraphDS) apply(m migrator) migrator {
+	return dgraphMigrator{
+		dgraphDS: ds,
+		migrator: m,
+	}
+}
+
+// ApplySchema applies the given schema to DGraph. It takes a context and schema string as parameters
+// and returns an error if the schema application fails.
+func (ds dgraphDS) ApplySchema(ctx context.Context, schema string) error {
+	return ds.client.ApplySchema(ctx, schema)
+}
+
+// AddOrUpdateField adds a new field or updates an existing field in DGraph schema.
+// Parameters:
+//   - ctx: The context for the operation
+//   - fieldName: Name of the field to add or update
+//   - fieldType: Data type of the field
+//   - directives: Additional DGraph directives for the field
+//
+// Returns an error if the operation fails.
+func (ds dgraphDS) AddOrUpdateField(ctx context.Context, fieldName, fieldType, directives string) error {
+	return ds.client.AddOrUpdateField(ctx, fieldName, fieldType, directives)
+}
+
+// DropField removes a field from DGraph schema.
+// Parameters:
+//   - ctx: The context for the operation
+//   - fieldName: Name of the field to remove
+//
+// Returns an error if the field deletion fails.
+func (ds dgraphDS) DropField(ctx context.Context, fieldName string) error {
+	return ds.client.DropField(ctx, fieldName)
+}
+
+// checkAndCreateMigrationTable ensures migration schema exists.
+func (dm dgraphMigrator) checkAndCreateMigrationTable(c *container.Container) error {
+	err := dm.ApplySchema(context.Background(), dgraphSchema)
+	if err != nil {
+		c.Debug("Migration schema might already exist:", err)
+	}
+
+	return dm.migrator.checkAndCreateMigrationTable(c)
+}
+
+// getLastMigration retrieves the last applied migration version.
+func (dm dgraphMigrator) getLastMigration(c *container.Container) int64 {
+	var response struct {
+		Migrations []struct {
+			Version int64 `json:"version"`
+		} `json:"migrations"`
+	}
+
+	resp, err := c.DGraph.Query(context.Background(), getLastMigrationQuery)
+	if err != nil {
+		c.Debug("Error fetching last migration:", err)
+		return 0
+	}
+
+	// If a response is returned, marshal it to JSON bytes then unmarshal into our response struct.
+	if resp != nil {
+		b, err := json.Marshal(resp)
+		if err != nil {
+			c.Debug("Error marshaling response:", err)
+			return 0
+		}
+
+		if err := json.Unmarshal(b, &response); err != nil {
+			c.Debug("Error unmarshalling migration response:", err)
+			return 0
+		}
+
+		if len(response.Migrations) > 0 {
+			return response.Migrations[0].Version
+		}
+	}
+
+	lm2 := dm.migrator.getLastMigration(c)
+	if lm2 > 0 {
+		return lm2
+	}
+
+	return 0
+}
+
+// beginTransaction starts a new migration transaction.
+func (dm dgraphMigrator) beginTransaction(c *container.Container) transactionData {
+	data := dm.migrator.beginTransaction(c)
+
+	c.Debug("Dgraph migrator begin successfully")
+
+	return data
+}
+
+// commitMigration commits the migration and records its metadata.
+func (dm dgraphMigrator) commitMigration(c *container.Container, data transactionData) error {
+	// Build the JSON payload for the migration record.
+	payload := map[string]any{
+		"migrations": []map[string]any{
+			{
+				"migrations.version":    data.MigrationNumber,
+				"migrations.method":     "UP",
+				"migrations.start_time": data.StartTime.Format(time.RFC3339),
+				"migrations.duration":   time.Since(data.StartTime).Milliseconds(),
+			},
+		},
+	}
+
+	jsonPayload, err := json.Marshal(payload)
+	if err != nil {
+		return err
+	}
+
+	_, err = c.DGraph.Mutate(context.Background(), &api.Mutation{
+		SetJson: jsonPayload,
+	})
+	if err != nil {
+		return err
+	}
+
+	c.Debugf("Inserted record for migration %v in Dgraph migrations", data.MigrationNumber)
+
+	return dm.migrator.commitMigration(c, data)
+}
+
+// rollback handles migration failure and rollback.
+func (dm dgraphMigrator) rollback(c *container.Container, data transactionData) {
+	dm.migrator.rollback(c, data)
+
+	c.Fatalf("Migration %v failed and rolled back", data.MigrationNumber)
+}
diff --git a/pkg/gofr/migration/interface.go b/pkg/gofr/migration/interface.go
index 773bbe9f..8747bf3b 100644
--- a/pkg/gofr/migration/interface.go
+++ b/pkg/gofr/migration/interface.go
@@ -99,6 +99,34 @@ type SurrealDB interface {
 	DropDatabase(ctx context.Context, database string) error
 }
 
+type DGraph interface {
+	// ApplySchema applies or updates the complete database schema.
+	// Parameters:
+	// - ctx: Context for request cancellation and timeouts
+	// - schema: Schema definition in Dgraph Schema Definition Language (SDL) format
+	// Returns:
+	// - error: An error if the schema application fails
+	ApplySchema(ctx context.Context, schema string) error
+
+	// AddOrUpdateField atomically creates or updates a single field definition.
+	// Parameters:
+	// - ctx: Context for request cancellation and timeouts
+	// - fieldName: Name of the field/predicate to create or update
+	// - fieldType: Dgraph data type (e.g., string, int, datetime)
+	// - directives: Space-separated Dgraph directives (e.g., "@index(hash) @upsert")
+	// Returns:
+	// - error: An error if the field operation fails
+	AddOrUpdateField(ctx context.Context, fieldName, fieldType, directives string) error
+
+	// DropField permanently removes a field/predicate and all its associated data.
+	// Parameters:
+	// - ctx: Context for request cancellation and timeouts
+	// - fieldName: Name of the field/predicate to remove
+	// Returns:
+	// - error: An error if the field removal fails
+	DropField(ctx context.Context, fieldName string) error
+}
+
 // keeping the migrator interface unexported as, right now it is not being implemented directly, by the externalDB drivers.
 // keeping the implementations for externalDB at one place such that if any change in migration logic, we would change directly here.
 type migrator interface {
diff --git a/pkg/gofr/migration/migration.go b/pkg/gofr/migration/migration.go
index 8f546657..ceef67e7 100644
--- a/pkg/gofr/migration/migration.go
+++ b/pkg/gofr/migration/migration.go
@@ -194,6 +194,16 @@ func getMigrator(c *container.Container) (Datasource, migrator, bool) {
 		c.Debug("initialized data source for surrealDB")
 	}
 
+	if !isNil(c.DGraph) {
+		ok = true
+
+		ds.DGraph = dgraphDS{c.DGraph}
+
+		mg = dgraphDS{c.DGraph}.apply(mg)
+
+		c.Debug("initialized data source for dgraph")
+	}
+
 	return ds, mg, ok
 }
 



Test patch:
diff --git a/pkg/gofr/datasource/dgraph/dgraph_test.go b/pkg/gofr/datasource/dgraph/dgraph_test.go
index f9d30809..0e9bfef5 100644
--- a/pkg/gofr/datasource/dgraph/dgraph_test.go
+++ b/pkg/gofr/datasource/dgraph/dgraph_test.go
@@ -30,6 +30,7 @@ func setupDB(t *testing.T) (*Client, *MockDgraphClient, *MockLogger, *MockMetric
 	client := New(config)
 	client.UseLogger(mockLogger)
 	client.UseMetrics(mockMetrics)
+	client.UseTracer(otel.GetTracerProvider().Tracer("gofr-dgraph"))
 
 	mockDgraphClient := NewMockDgraphClient(ctrl)
 	client.client = mockDgraphClient
@@ -290,3 +291,95 @@ func Test_HealthCheck_Error(t *testing.T) {
 
 	require.EqualError(t, err, errHealthCheckFailed.Error(), "Test_HealthCheck_Error Failed!")
 }
+
+func Test_ApplySchema_Success(t *testing.T) {
+	client, mockDgraphClient, mockLogger, mockMetrics := setupDB(t)
+
+	schema := "name: string @index(exact) ."
+	expectedOp := &api.Operation{Schema: schema}
+
+	mockDgraphClient.EXPECT().Alter(gomock.Any(), expectedOp).Return(nil)
+
+	mockLogger.EXPECT().Debug(gomock.Any())
+	mockLogger.EXPECT().Log(gomock.Any()).Times(1)
+	mockLogger.EXPECT().Debugf("dgraph alter succeeded in %dµs", gomock.Any())
+	mockMetrics.EXPECT().RecordHistogram(gomock.Any(), "dgraph_alter_duration", gomock.Any())
+
+	err := client.ApplySchema(context.Background(), schema)
+
+	require.NoError(t, err, "Test_ApplySchema_Success Failed!")
+}
+
+func Test_ApplySchema_EmptySchema(t *testing.T) {
+	config := Config{Host: "localhost", Port: "9080"}
+	client := New(config)
+
+	err := client.ApplySchema(context.Background(), "")
+
+	require.Equal(t, err, errEmptySchema, "Test_ApplySchema_EmptySchema Failed!")
+}
+
+func Test_AddOrUpdateField_Success(t *testing.T) {
+	client, mockDgraphClient, mockLogger, mockMetrics := setupDB(t)
+
+	fieldName := "email"
+	fieldType := "string"
+	directives := "@index(hash)"
+	expectedSchema := "email: string @index(hash)."
+	expectedOp := &api.Operation{Schema: expectedSchema}
+
+	mockDgraphClient.EXPECT().Alter(gomock.Any(), expectedOp).Return(nil)
+
+	mockLogger.EXPECT().Debug(gomock.Any())
+	mockLogger.EXPECT().Log(gomock.Any()).Times(1)
+	mockLogger.EXPECT().Debugf("dgraph alter succeeded in %dµs", gomock.Any())
+	mockMetrics.EXPECT().RecordHistogram(gomock.Any(), "dgraph_alter_duration", gomock.Any())
+
+	err := client.AddOrUpdateField(context.Background(), fieldName, fieldType, directives)
+
+	require.NoError(t, err, "Test_AddOrUpdateField_Success Failed!")
+}
+
+func Test_AddOrUpdateField_EmptyFieldName(t *testing.T) {
+	config := Config{Host: "localhost", Port: "9080"}
+	client := New(config)
+
+	err := client.AddOrUpdateField(context.Background(), "", "string", "@index(hash)")
+
+	require.Equal(t, err, errEmptyField, "Test_AddOrUpdateField_EmptyFieldName Failed!")
+}
+
+func Test_DropField_Success(t *testing.T) {
+	client, mockDgraphClient, mockLogger, mockMetrics := setupDB(t)
+
+	fieldName := "email"
+	expectedOp := &api.Operation{DropAttr: fieldName}
+
+	mockDgraphClient.EXPECT().Alter(gomock.Any(), expectedOp).Return(nil)
+
+	mockLogger.EXPECT().Debug(gomock.Any())
+	mockLogger.EXPECT().Log(gomock.Any()).Times(1)
+	mockLogger.EXPECT().Debugf("dgraph alter succeeded in %dµs", gomock.Any())
+	mockMetrics.EXPECT().RecordHistogram(gomock.Any(), "dgraph_alter_duration", gomock.Any())
+
+	err := client.DropField(context.Background(), fieldName)
+
+	require.NoError(t, err, "Test_DropField_Success Failed!")
+}
+
+func Test_DropField_Error(t *testing.T) {
+	client, mockDgraphClient, mockLogger, _ := setupDB(t)
+
+	fieldName := "email"
+	expectedOp := &api.Operation{DropAttr: fieldName}
+
+	mockDgraphClient.EXPECT().Alter(gomock.Any(), expectedOp).Return(errAlterFailed)
+
+	mockLogger.EXPECT().Debug(gomock.Any())
+	mockLogger.EXPECT().Log(gomock.Any()).Times(1)
+	mockLogger.EXPECT().Error("dgraph alter failed: ", errAlterFailed)
+
+	err := client.DropField(context.Background(), fieldName)
+
+	require.ErrorIs(t, err, errAlterFailed, "Test_DropField_Error Failed!")
+}
diff --git a/pkg/gofr/migration/dgraph_test.go b/pkg/gofr/migration/dgraph_test.go
new file mode 100644
index 00000000..9423a2d2
--- /dev/null
+++ b/pkg/gofr/migration/dgraph_test.go
@@ -0,0 +1,198 @@
+package migration
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+
+	"gofr.dev/pkg/gofr/container"
+	"gofr.dev/pkg/gofr/testutil"
+)
+
+func dgraphSetup(t *testing.T) (migrator, *container.MockDgraph, *container.Container) {
+	t.Helper()
+
+	mockContainer, mocks := container.NewMockContainer(t)
+	mockDGraph := mocks.DGraph
+
+	ds := Datasource{DGraph: mockContainer.DGraph}
+
+	dgraphDB := dgraphDS{client: mockDGraph}
+	migratorWithDGraph := dgraphDB.apply(&ds)
+
+	mockContainer.DGraph = mockDGraph
+
+	return migratorWithDGraph, mockDGraph, mockContainer
+}
+
+func Test_DGraphCheckAndCreateMigrationTable(t *testing.T) {
+	migratorWithDGraph, mockDGraph, mockContainer := dgraphSetup(t)
+
+	mockDGraph.EXPECT().ApplySchema(context.Background(), dgraphSchema).Return(nil)
+
+	err := migratorWithDGraph.checkAndCreateMigrationTable(mockContainer)
+
+	require.NoError(t, err, "Test_DGraphCheckAndCreateMigrationTable Failed!")
+}
+
+func Test_DGraphGetLastMigration(t *testing.T) {
+	migratorWithDGraph, mockDGraph, mockContainer := dgraphSetup(t)
+
+	testCases := []struct {
+		desc     string
+		err      error
+		mockResp map[string]any
+		expected int64
+	}{
+		{
+			desc: "success",
+			err:  nil,
+			mockResp: map[string]any{
+				"migrations": []map[string]any{
+					{
+						"version": float64(10),
+					},
+				},
+			},
+			expected: 10,
+		},
+		{
+			desc:     "query error",
+			err:      context.DeadlineExceeded,
+			mockResp: nil,
+			expected: 0,
+		},
+		{
+			desc:     "empty response",
+			err:      nil,
+			mockResp: map[string]any{},
+			expected: 0,
+		},
+	}
+
+	for i, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			// Set up mock expectation for the main query
+			mockDGraph.EXPECT().Query(context.Background(), getLastMigrationQuery).
+				Return(tc.mockResp, tc.err)
+
+			resp := migratorWithDGraph.getLastMigration(mockContainer)
+
+			assert.Equal(t, tc.expected, resp, "TEST[%v] Failed!", i)
+		})
+	}
+}
+
+func Test_DGraphCommitMigration(t *testing.T) {
+	migratorWithDGraph, mockDGraph, mockContainer := dgraphSetup(t)
+
+	timeNow := time.Now()
+
+	testCases := []struct {
+		desc string
+		err  error
+	}{
+		{"success", nil},
+		{"mutation failed", context.DeadlineExceeded},
+	}
+
+	td := transactionData{
+		StartTime:       timeNow,
+		MigrationNumber: 10,
+	}
+
+	for i, tc := range testCases {
+		mockDGraph.EXPECT().Mutate(context.Background(), gomock.Any()).Return(nil, tc.err)
+
+		err := migratorWithDGraph.commitMigration(mockContainer, td)
+
+		assert.Equal(t, tc.err, err, "TEST[%v]\n %v Failed!", i, tc.desc)
+	}
+}
+
+func Test_DGraphBeginTransaction(t *testing.T) {
+	logs := testutil.StdoutOutputForFunc(func() {
+		migratorWithDGraph, _, mockContainer := dgraphSetup(t)
+		migratorWithDGraph.beginTransaction(mockContainer)
+	})
+
+	assert.Contains(t, logs, "Dgraph migrator begin successfully")
+}
+
+func Test_DGraphDS_ApplySchema(t *testing.T) {
+	_, mockDGraph, _ := dgraphSetup(t)
+
+	ds := dgraphDS{client: mockDGraph}
+	ctx := context.Background()
+	schema := "test schema"
+
+	testCases := []struct {
+		desc string
+		err  error
+	}{
+		{"success", nil},
+		{"schema error", context.DeadlineExceeded},
+	}
+
+	for i, tc := range testCases {
+		mockDGraph.EXPECT().ApplySchema(ctx, schema).Return(tc.err)
+
+		err := ds.ApplySchema(ctx, schema)
+
+		assert.Equal(t, tc.err, err, "TEST[%v]\n %v Failed!", i, tc.desc)
+	}
+}
+
+func Test_DGraphDS_AddOrUpdateField(t *testing.T) {
+	_, mockDGraph, _ := dgraphSetup(t)
+
+	ds := dgraphDS{client: mockDGraph}
+	ctx := context.Background()
+	fieldName := "test"
+	fieldType := "string"
+	directives := "@index(exact)"
+
+	testCases := []struct {
+		desc string
+		err  error
+	}{
+		{"success", nil},
+		{"field error", context.DeadlineExceeded},
+	}
+
+	for i, tc := range testCases {
+		mockDGraph.EXPECT().AddOrUpdateField(ctx, fieldName, fieldType, directives).Return(tc.err)
+
+		err := ds.AddOrUpdateField(ctx, fieldName, fieldType, directives)
+
+		assert.Equal(t, tc.err, err, "TEST[%v]\n %v Failed!", i, tc.desc)
+	}
+}
+
+func Test_DGraphDS_DropField(t *testing.T) {
+	_, mockDGraph, _ := dgraphSetup(t)
+
+	ds := dgraphDS{client: mockDGraph}
+	ctx := context.Background()
+	fieldName := "test"
+
+	testCases := []struct {
+		desc string
+		err  error
+	}{
+		{"success", nil},
+		{"drop error", context.DeadlineExceeded},
+	}
+
+	for i, tc := range testCases {
+		mockDGraph.EXPECT().DropField(ctx, fieldName).Return(tc.err)
+
+		err := ds.DropField(ctx, fieldName)
+
+		assert.Equal(t, tc.err, err, "TEST[%v]\n %v Failed!", i, tc.desc)
+	}
+}
diff --git a/pkg/gofr/migration/migration_test.go b/pkg/gofr/migration/migration_test.go
index a1acec3d..6ed1063f 100644
--- a/pkg/gofr/migration/migration_test.go
+++ b/pkg/gofr/migration/migration_test.go
@@ -68,6 +68,7 @@ func initialiseClickHouseRunMocks(t *testing.T) (*MockClickhouse, *container.Con
 	mockContainer.PubSub = nil
 	mockContainer.ArangoDB = nil
 	mockContainer.SurrealDB = nil
+	mockContainer.DGraph = nil
 	mockContainer.Logger = logging.NewMockLogger(logging.DEBUG)
 	mockContainer.Clickhouse = mockClickHouse