Merge branch 'develop' of github.com:openlattice/mechanic into develop

openlattice · Mar 4, 2022 · f334e01 · f334e01
2 parents a73ee6d + 5f1dd21
commit f334e01
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 0 deletions.
diff --git a/src/main/kotlin/com/openlattice/mechanic/pods/MechanicUpgradePod.kt b/src/main/kotlin/com/openlattice/mechanic/pods/MechanicUpgradePod.kt
@@ -72,6 +72,7 @@ import com.openlattice.linking.graph.PostgresLinkingQueryService
 import com.openlattice.mechanic.MechanicCli.Companion.UPGRADE
 import com.openlattice.mechanic.Toolbox
 import com.openlattice.mechanic.upgrades.DeleteOrgMetadataEntitySets
+import com.openlattice.mechanic.upgrades.ExportOrganizationMembers
 import com.openlattice.mechanic.upgrades.V3StudyMigrationUpgrade
 import com.openlattice.organizations.roles.HazelcastPrincipalService
 import com.openlattice.organizations.roles.SecurePrincipalsManager
@@ -379,6 +380,11 @@ class MechanicUpgradePod {
         )
     }
 
+    @Bean
+    fun exportOrganizationMembers() : ExportOrganizationMembers {
+        return ExportOrganizationMembers(toolbox, hikariDataSource, principalService(), hazelcastInstance)
+    }
+
     @PostConstruct
     fun post() {
         Principals.init(principalService(), hazelcastInstance)

diff --git a/src/main/kotlin/com/openlattice/mechanic/upgrades/ExportOrganizationMembers.kt b/src/main/kotlin/com/openlattice/mechanic/upgrades/ExportOrganizationMembers.kt
@@ -0,0 +1,68 @@
+package com.openlattice.mechanic.upgrades
+
+import com.hazelcast.core.HazelcastInstance
+import com.openlattice.hazelcast.HazelcastMap
+import com.openlattice.mechanic.Toolbox
+import com.openlattice.organizations.roles.SecurePrincipalsManager
+import com.zaxxer.hikari.HikariDataSource
+import org.slf4j.LoggerFactory
+
+/**
+ *
+ * @author Matthew Tamayo-Rios &lt;[email protected]&gt;
+ */
+class ExportOrganizationMembers(
+    private val toolbox: Toolbox,
+    private val hds: HikariDataSource,
+    private val principalService: SecurePrincipalsManager,
+    hazelcast :HazelcastInstance
+) : Upgrade {
+    companion object {
+        private const val USERS_EXPORT_TABLE_NAME = "users_export"
+        const val USERS_EXPORT_TABLE = """
+            CREATE TABLE $USERS_EXPORT_TABLE_NAME ( 
+                organization_id uuid,
+                principal_id text NOT NULL,
+                principal_email text, 
+                principal_username text,
+                PRIMARY KEY( organization_id, principal_id ) );
+        """
+        const val INSERT_USER_SQL = """
+            INSERT INTO $USERS_EXPORT_TABLE_NAME (organization_id, principal_id, principal_email, principal_username) VALUES (?,?,?,?)
+        """
+        private val logger = LoggerFactory.getLogger(ExportOrganizationMembers::class.java)
+    }
+
+    private val organizations = HazelcastMap.ORGANIZATIONS.getMap(hazelcast)
+    private val users = HazelcastMap.USERS.getMap(hazelcast)
+
+    override fun upgrade(): Boolean {
+        val organizationsIds = organizations.keys.toMutableSet()
+        val members = principalService.getOrganizationMembers(organizationsIds)
+
+        hds.connection.use { connection ->
+            connection.createStatement().use { stmt -> stmt.execute(USERS_EXPORT_TABLE) }
+            connection.prepareStatement(INSERT_USER_SQL).use { ps ->
+            members.forEach { (orgId, orgMembers) ->
+                orgMembers.forEach { orgMember ->
+                    val user = users[ orgMember.name ]
+                        ps.setObject(1, orgId)
+                        ps.setString(2, orgMember.name)
+                        ps.setString(3, user?.email)
+                        ps.setString(4, user?.username)
+                    ps.addBatch()
+                    }
+                }
+                val userCount = ps.executeBatch().sum()
+                logger.info("Exported $userCount users across ${organizationsIds.size} organizations.")
+            }
+        }
+
+        return true
+    }
+
+    override fun getSupportedVersion(): Long {
+        return Version.V2021_07_23.value
+    }
+
+}