feat: 支持同时连接多个mongo TencentBlueKing#2591

yaoxuwan · Sep 24, 2024 · 9e3ecbc · 9e3ecbc
1 parent 24e0bf7
commit 9e3ecbc
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 3 deletions.
diff --git a/...-analyst/src/main/kotlin/com/tencent/bkrepo/analyst/configuration/ScannerConfiguration.kt b/...-analyst/src/main/kotlin/com/tencent/bkrepo/analyst/configuration/ScannerConfiguration.kt
@@ -27,6 +27,9 @@
 
 package com.tencent.bkrepo.analyst.configuration
 
+import com.mongodb.ConnectionString
+import com.mongodb.MongoClientSettings
+import com.mongodb.client.MongoClients
 import com.tencent.bkrepo.analysis.executor.api.ExecutorClient
 import com.tencent.bkrepo.analyst.dispatcher.SubtaskDispatcherFactory
 import com.tencent.bkrepo.analyst.dispatcher.SubtaskPoller
@@ -42,10 +45,18 @@ import com.tencent.bkrepo.common.service.condition.ConditionalOnNotAssembly
 import com.tencent.bkrepo.repository.api.OperateLogClient
 import com.tencent.bkrepo.repository.api.ProjectUsageStatisticsClient
 import org.springframework.beans.factory.ObjectProvider
+import org.springframework.beans.factory.annotation.Qualifier
+import org.springframework.boot.autoconfigure.mongo.MongoProperties
+import org.springframework.boot.context.properties.ConfigurationProperties
 import org.springframework.boot.context.properties.EnableConfigurationProperties
 import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClients
 import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.Configuration
+import org.springframework.data.mongodb.MongoDatabaseFactory
+import org.springframework.data.mongodb.SpringDataMongoDB
+import org.springframework.data.mongodb.core.MongoTemplate
+import org.springframework.data.mongodb.core.SimpleMongoClientDatabaseFactory
+import org.springframework.data.mongodb.core.convert.MappingMongoConverter
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor
 import java.util.function.Consumer
 
@@ -97,4 +108,36 @@ class ScannerConfiguration {
             analystScanEventConsumer.accept(it)
         }
     }
+
+    @Bean
+    @ConfigurationProperties("spring.data.mongodb.analyst")
+    fun analystMongoProperties(): MongoProperties {
+        return MongoProperties()
+    }
+
+    @Bean
+    fun analystMongoTemplate(
+        mongoTemplate: MongoTemplate,
+        @Qualifier("analystMongoProperties") analystMongoProperties: MongoProperties,
+        converter: MappingMongoConverter?
+    ): MongoTemplate? {
+        // 没有设置独立数据库时，使用统一的数据库
+        return if (analystMongoProperties.database.isNullOrBlank()) {
+            mongoTemplate
+        } else {
+            val settings =
+                MongoClientSettings.builder()
+                    .applyConnectionString(ConnectionString(analystMongoProperties.determineUri()))
+                    .uuidRepresentation(analystMongoProperties.uuidRepresentation)
+                    .build()
+
+            val databaseFactory: MongoDatabaseFactory =
+                SimpleMongoClientDatabaseFactory(
+                    MongoClients.create(settings, SpringDataMongoDB.driverInformation()),
+                    analystMongoProperties.database
+                )
+
+            MongoTemplate(databaseFactory, converter)
+        }
+    }
 }
diff --git a/...alyst/biz-analyst/src/main/kotlin/com/tencent/bkrepo/analyst/dao/ScannerSimpleMongoDao.kt b/...alyst/biz-analyst/src/main/kotlin/com/tencent/bkrepo/analyst/dao/ScannerSimpleMongoDao.kt
@@ -29,10 +29,20 @@ package com.tencent.bkrepo.analyst.dao
 
 import com.tencent.bkrepo.common.api.pojo.Page
 import com.tencent.bkrepo.common.mongo.dao.simple.SimpleMongoDao
+import org.springframework.beans.factory.annotation.Autowired
 import org.springframework.data.domain.PageRequest
+import org.springframework.data.mongodb.core.MongoTemplate
 import org.springframework.data.mongodb.core.query.Query
 
 abstract class ScannerSimpleMongoDao<T> : SimpleMongoDao<T>() {
+
+    @Autowired
+    private lateinit var analystMongoTemplate: MongoTemplate
+
+    override fun determineMongoTemplate(): MongoTemplate {
+        return analystMongoTemplate
+    }
+
     fun page(query: Query, pageRequest: PageRequest): Page<T> {
         val count = count(query)
         val records = find(query.with(pageRequest))

diff --git a/...kend/analyst/biz-analyst/src/main/kotlin/com/tencent/bkrepo/analyst/dao/SpdxLicenseDao.kt b/...kend/analyst/biz-analyst/src/main/kotlin/com/tencent/bkrepo/analyst/dao/SpdxLicenseDao.kt
@@ -27,7 +27,6 @@
 
 package com.tencent.bkrepo.analyst.dao
 
-import com.tencent.bkrepo.common.mongo.dao.simple.SimpleMongoDao
 import com.tencent.bkrepo.analyst.model.TSpdxLicense
 import com.tencent.bkrepo.common.api.util.EscapeUtils
 import org.bson.BsonRegularExpression
@@ -37,7 +36,7 @@ import org.springframework.data.mongodb.core.query.isEqualTo
 import org.springframework.stereotype.Repository
 
 @Repository
-class SpdxLicenseDao : SimpleMongoDao<TSpdxLicense>() {
+class SpdxLicenseDao : ScannerSimpleMongoDao<TSpdxLicense>() {
     fun findByLicenseId(licenseId: String): TSpdxLicense? {
         return this.findOne(Query(TSpdxLicense::licenseId.isEqualTo(licenseId)))
     }

diff --git a/...on/common-mongo/src/main/kotlin/com/tencent/bkrepo/common/mongo/MongoAutoConfiguration.kt b/...on/common-mongo/src/main/kotlin/com/tencent/bkrepo/common/mongo/MongoAutoConfiguration.kt
@@ -35,12 +35,15 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty
 import org.springframework.boot.autoconfigure.mongo.MongoProperties
 import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.Configuration
+import org.springframework.context.annotation.Primary
 import org.springframework.context.annotation.PropertySource
 import org.springframework.data.mongodb.MongoDatabaseFactory
 import org.springframework.data.mongodb.MongoTransactionManager
+import org.springframework.data.mongodb.core.MongoTemplate
 import org.springframework.data.mongodb.core.convert.DefaultDbRefResolver
 import org.springframework.data.mongodb.core.convert.DefaultMongoTypeMapper
 import org.springframework.data.mongodb.core.convert.MappingMongoConverter
+import org.springframework.data.mongodb.core.convert.MongoConverter
 import org.springframework.data.mongodb.core.convert.MongoCustomConversions
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext
 
@@ -75,9 +78,21 @@ class MongoAutoConfiguration {
         mappingContext.isAutoIndexCreation = mongoProperties.isAutoIndexCreation
 
         val converter = MappingMongoConverter(dbRefResolver, mappingContext)
-        converter.typeMapper = DefaultMongoTypeMapper(null)
+        converter.setTypeMapper(DefaultMongoTypeMapper(null))
         converter.afterPropertiesSet()
         converter.setMapKeyDotReplacement("#dot#")
         return converter
     }
+
+    @Bean
+    @Primary
+    fun mongoProperties(): MongoProperties {
+        return MongoProperties()
+    }
+
+    @Bean
+    @Primary
+    fun mongoTemplate(factory: MongoDatabaseFactory, converter: MongoConverter?): MongoTemplate {
+        return MongoTemplate(factory, converter)
+    }
 }