소스 검색

install to db

Yeuoly 10 달 전
부모
커밋
f42bee7019
2개의 변경된 파일51개의 추가작업 그리고 1개의 파일을 삭제
  1. 4 1
      internal/db/pgsql.go
  2. 47 0
      internal/types/models/curd/atomic.go

+ 4 - 1
internal/db/pgsql.go

@@ -3,6 +3,7 @@ package db
 import (
 	"fmt"
 
+	"github.com/langgenius/dify-plugin-daemon/internal/utils/log"
 	"gorm.io/gorm"
 	"gorm.io/gorm/clause"
 )
@@ -306,7 +307,9 @@ func WithTransaction(fn func(tx *gorm.DB) error, ctx ...*gorm.DB) error {
 
 	err := fn(tx)
 	if err != nil {
-		tx.Rollback()
+		if err := tx.Rollback().Error; err != nil {
+			log.Error("failed to rollback tx: %v", err)
+		}
 		return err
 	}
 	tx.Commit()

+ 47 - 0
internal/types/models/curd/atomic.go

@@ -85,6 +85,53 @@ func CreatePlugin(
 
 		installation_to_be_returns = installation
 
+		// create tool installation
+		if declaration.Tool != nil {
+			tool_installation := &models.ToolInstallation{
+				PluginID:               plugin_to_be_returns.PluginID,
+				PluginUniqueIdentifier: plugin_to_be_returns.PluginUniqueIdentifier,
+				TenantID:               tenant_id,
+				Provider:               declaration.Tool.Identity.Name,
+				Declaration:            *declaration.Tool,
+			}
+
+			err := db.Create(tool_installation, tx)
+			if err != nil {
+				return err
+			}
+		}
+
+		// create endpoint installation
+		if declaration.Endpoint != nil {
+			endpoint_installation := &models.EndpointInstallation{
+				PluginID:               plugin_to_be_returns.PluginID,
+				PluginUniqueIdentifier: plugin_to_be_returns.PluginUniqueIdentifier,
+				TenantID:               tenant_id,
+				Declaration:            *declaration.Endpoint,
+			}
+
+			err := db.Create(endpoint_installation, tx)
+			if err != nil {
+				return err
+			}
+		}
+
+		// create model installation
+		if declaration.Model != nil {
+			model_installation := &models.AIModelInstallation{
+				PluginID:               plugin_to_be_returns.PluginID,
+				PluginUniqueIdentifier: plugin_to_be_returns.PluginUniqueIdentifier,
+				TenantID:               tenant_id,
+				Provider:               declaration.Model.Provider,
+				Declaration:            *declaration.Model,
+			}
+
+			err := db.Create(model_installation, tx)
+			if err != nil {
+				return err
+			}
+		}
+
 		return nil
 	})