From 3300a4091a4730052ce3ffb1507396b63613d359 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Jodas?= <ondra.jodas@gmail.com>
Date: Tue, 4 Feb 2025 23:09:20 +0100
Subject: [PATCH 1/7] Refactor Cleanup class: optimize code, add logging,
 improve type safety

---
 src/Cleanup.php | 266 ++++++++++++++++++++++++++++++------------------
 1 file changed, 168 insertions(+), 98 deletions(-)

diff --git a/src/Cleanup.php b/src/Cleanup.php
index 29ac404..309d6cb 100644
--- a/src/Cleanup.php
+++ b/src/Cleanup.php
@@ -17,6 +17,10 @@
 
 class Cleanup
 {
+    private const OWNERSHIP = 'OWNERSHIP';
+    private const USER = 'USER';
+    private const ROLE = 'ROLE';
+
     public function __construct(
         readonly Config $config,
         readonly Connection $sourceConnection,
@@ -87,13 +91,12 @@ public function sourceAccount(): void
 
     public function preMigration(string $mainRoleName): void
     {
+        $this->logger->info(sprintf('Starting pre-migration cleanup with main role: %s', $mainRoleName));
         $sqls = [];
-        $currentRole = $this->config->getTargetSnowflakeRole();
+        $currentRole = null;
 
-        $mainRole = $this->destinationConnection->fetchAll(sprintf(
-            'SHOW ROLES LIKE %s',
-            QueryBuilder::quote($mainRoleName),
-        ));
+        // Check if main role is assigned to target user
+        $this->logger->info('Checking main role assignment and ownership');
         $grantsToTargetUser = $this->destinationConnection->fetchAll(sprintf(
             'SHOW GRANTS TO USER %s',
             QueryBuilder::quoteIdentifier($this->config->getTargetSnowflakeUser()),
@@ -104,6 +107,11 @@ public function preMigration(string $mainRoleName): void
             false,
         );
 
+        // Check if target role has ownership of main role
+        $mainRole = $this->destinationConnection->fetchAll(sprintf(
+            'SHOW ROLES LIKE %s',
+            QueryBuilder::quote($mainRoleName),
+        ));
         $hasMainRoleOwnership = array_reduce(
             $mainRole,
             fn ($found, $v) => $found || $v['owner'] === $this->config->getTargetSnowflakeRole(),
@@ -118,116 +126,105 @@ public function preMigration(string $mainRoleName): void
             }
             $this->destinationConnection->grantRoleToUser($this->config->getTargetSnowflakeUser(), $mainRoleName);
         }
+
+        // Process each database
         foreach ($this->config->getDatabases() as $database) {
+            $this->logger->info(sprintf('Processing database: %s', $database));
             $this->destinationConnection->useRole($this->config->getTargetSnowflakeRole());
 
+            // Check if database exists
             $dbExists = $this->destinationConnection->fetchAll(sprintf(
                 'SHOW DATABASES LIKE %s;',
                 QueryBuilder::quote($database)
             ));
+
             if (!$dbExists) {
-                continue;
-            }
-            $databaseRole = $this->destinationConnection->getOwnershipRoleOnDatabase($database);
-            $data = $this->getDataToRemove($this->destinationConnection, $databaseRole);
-
-            $currentUser = $this->destinationConnection->fetchAll('SELECT CURRENT_USER() AS "user";');
-            foreach ($data['USER'] ?? [] as $user) {
-                if ($user['granted_by'] !== $currentRole) {
-                    $currentRole = $user['granted_by'];
-                    if ($currentRole === $mainRoleName && !$mainRoleExistsOnTargetUser) {
-                        $sqls[] = sprintf(
-                            'GRANT ROLE %s TO USER %s;',
-                            Helper::quoteIdentifier($currentRole),
-                            Helper::quoteIdentifier($this->config->getTargetSnowflakeUser())
-                        );
-                    }
-                    $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($currentRole));
-                }
+                $this->logger->info(sprintf('Database %s does not exist, checking for role with same name', $database));
+                // If database doesn't exist, check if there's a role with the same name
+                $roleExists = $this->destinationConnection->fetchAll(sprintf(
+                    'SHOW ROLES LIKE %s;',
+                    QueryBuilder::quote($database)
+                ));
 
-                if ($currentUser[0]['user'] !== $user['name']) {
-                    $sqls[] = sprintf('DROP USER IF EXISTS %s;', Helper::quoteIdentifier($user['name']));
+                if (!$roleExists) {
+                    continue;
                 }
+                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $database);
+                $roleToRemove = $database;
+            } else {
+                $this->logger->info(sprintf('Database %s exists, getting ownership role', $database));
+                $databaseRole = $this->destinationConnection->getOwnershipRoleOnDatabase($database);
+                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $databaseRole);
+                $roleToRemove = $databaseRole;
             }
 
-            foreach ($data['ROLE'] ?? [] as $role) {
-                if ($role['granted_by'] !== $currentRole) {
-                    $currentRole = $role['granted_by'];
-                    $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($currentRole));
-                }
-                $this->destinationConnection->useRole($mainRoleName);
-                if ($currentRole === $mainRoleName && !$mainRoleExistsOnTargetUser) {
-                    $this->destinationConnection->grantRoleToUser(
-                        $this->config->getTargetSnowflakeUser(),
-                        $currentRole,
-                    );
-                }
-                try {
-                    $this->destinationConnection->useRole($role['granted_by']);
-                } catch (RuntimeException $e) {
-                    $this->destinationConnection->grantRoleToUser(
-                        $this->config->getTargetSnowflakeUser(),
-                        $role['granted_by'],
-                    );
-                    $this->destinationConnection->useRole($role['granted_by']);
-                }
+            // First revoke all future grants from roles
+            foreach ($dataToRemove['ROLE'] as $role) {
+                $this->switchRole($role['granted_by'], $mainRoleName, $sqls, $currentRole);
 
-                /** @var FutureGrantToRole[] $futureGrants */
                 $futureGrants = array_map(
                     fn(array $v) => FutureGrantToRole::fromArray($v),
                     $this->destinationConnection->fetchAll(sprintf(
                         'SHOW FUTURE GRANTS TO ROLE %s',
-                        $role['name']
+                        Helper::quoteIdentifier($role['name'])
                     ))
                 );
+
                 foreach ($futureGrants as $futureGrant) {
                     $sqls[] = sprintf(
                         'REVOKE %s ON FUTURE TABLES IN SCHEMA %s FROM ROLE %s;',
                         $futureGrant->getPrivilege(),
                         $futureGrant->getName(),
-                        Helper::quoteIdentifier($futureGrant->getGranteeName()),
+                        Helper::quoteIdentifier($futureGrant->getGranteeName())
                     );
                 }
-                $sqls[] = sprintf('DROP ROLE IF EXISTS %s;', Helper::quoteIdentifier($role['name']));
             }
 
-            $projectUser = $this->getProjectUser($databaseRole);
-
-            if ($projectUser->getGrantedBy() !== $currentRole) {
-                $currentRole = $projectUser->getGrantedBy();
-                if ($currentRole === $mainRoleName && !$mainRoleExistsOnTargetUser) {
-                    $sqls[] = sprintf(
-                        'USE ROLE %s;',
-                        Helper::quoteIdentifier($this->config->getTargetSnowflakeRole())
-                    );
-                    $sqls[] = sprintf(
-                        'GRANT ROLE %s TO USER %s;',
-                        Helper::quoteIdentifier($currentRole),
-                        Helper::quoteIdentifier($this->config->getTargetSnowflakeUser())
-                    );
-                }
-                $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($currentRole));
+            // Drop users owned by roles
+            if (!empty($dataToRemove['USER'])) {
+                $this->logger->info(sprintf('Dropping %d users', count($dataToRemove['USER'])));
+            }
+            foreach ($dataToRemove['USER'] as $user) {
+                $this->switchRole($user['granted_by'], $mainRoleName, $sqls, $currentRole);
+                $sqls[] = sprintf(
+                    'DROP USER IF EXISTS %s;',
+                    Helper::quoteIdentifier($user['name'])
+                );
             }
-            $sqls[] = sprintf(
-                'DROP USER IF EXISTS %s;',
-                Helper::quoteIdentifier($projectUser->getGranteeName())
-            );
 
-            $sqls[] = sprintf('DROP ROLE IF EXISTS %s;', Helper::quoteIdentifier($databaseRole));
+            // Drop roles
+            if (!empty($dataToRemove['ROLE'])) {
+                $this->logger->info(sprintf('Dropping %d roles', count($dataToRemove['ROLE'])));
+            }
+            foreach ($dataToRemove['ROLE'] as $role) {
+                $this->switchRole($role['granted_by'], $mainRoleName, $sqls, $currentRole);
+                $sqls[] = sprintf(
+                    'DROP ROLE IF EXISTS %s;',
+                    Helper::quoteIdentifier($role['name'])
+                );
+            }
 
+            // Drop database role and handle database if exists
+            $this->switchRole($this->config->getTargetSnowflakeRole(), $mainRoleName, $sqls, $currentRole);
             $sqls[] = sprintf(
-                'DROP DATABASE IF EXISTS %s;',
-                Helper::quoteIdentifier($database . '_OLD')
+                'DROP ROLE IF EXISTS %s;',
+                Helper::quoteIdentifier($roleToRemove)
             );
 
-            $sqls[] = sprintf(
-                'ALTER DATABASE IF EXISTS %s RENAME TO %s;',
-                Helper::quoteIdentifier($database),
-                Helper::quoteIdentifier($database . '_OLD'),
-            );
+            if ($dbExists) {
+                $sqls[] = sprintf(
+                    'DROP DATABASE IF EXISTS %s;',
+                    Helper::quoteIdentifier($database . '_OLD')
+                );
+                $sqls[] = sprintf(
+                    'ALTER DATABASE IF EXISTS %s RENAME TO %s;',
+                    Helper::quoteIdentifier($database),
+                    Helper::quoteIdentifier($database . '_OLD')
+                );
+            }
         }
 
-        $this->destinationConnection->useRole($mainRoleName);
+        // Execute all SQL commands
         foreach ($sqls as $sql) {
             if ($this->config->getSynchronizeDryPremigrationCleanupRun()) {
                 $this->logger->info($sql);
@@ -235,6 +232,7 @@ public function preMigration(string $mainRoleName): void
                 $this->destinationConnection->query($sql);
             }
         }
+
         if ($this->config->getSynchronizeDryPremigrationCleanupRun() && $sqls) {
             throw new UserException('!!! PLEASE RUN SQLS ON TARGET SNOWFLAKE ACCOUNT !!!');
         }
@@ -297,36 +295,108 @@ public function postMigration(): void
         }
     }
 
-    private function getDataToRemove(Connection $connection, string $role): array
+    private function switchRole(string $role, string $mainRoleName, array &$sqls, ?string &$currentRole): void
+    {
+        if ($currentRole === $role) {
+            return;
+        }
+
+        try {
+            $this->logger->debug(sprintf('Switching to role: %s', $role));
+            $this->destinationConnection->useRole($role);
+            $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($role));
+            $currentRole = $role;
+        } catch (RuntimeException $e) {
+            $this->logger->info(sprintf(
+                'Cannot switch directly to role %s, trying through main role %s',
+                $role,
+                $mainRoleName
+            ));
+            $this->destinationConnection->useRole($mainRoleName);
+            if ($currentRole !== $mainRoleName) {
+                $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($mainRoleName));
+                $currentRole = $mainRoleName;
+            }
+            $this->destinationConnection->grantRoleToUser(
+                $this->config->getTargetSnowflakeUser(),
+                $role
+            );
+            $this->logger->debug(sprintf('Granted role %s to user, switching to it', $role));
+            $this->destinationConnection->useRole($role);
+            $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($role));
+            $currentRole = $role;
+        }
+    }
+
+    private function getDataToRemove(Connection $connection, string $name): array
     {
+        $this->logger->debug(sprintf('Getting data to remove for: %s', $name));
+        $result = [
+            self::ROLE => [],
+            self::USER => [],
+        ];
+
+        // Get all grants for the name (both as role and user)
         $grants = $connection->fetchAll(sprintf(
             'SHOW GRANTS TO ROLE %s',
-            Helper::quoteIdentifier($role)
+            Helper::quoteIdentifier($name)
         ));
 
-        $filteredGrants = array_filter($grants, function ($v) {
-            $usageWarehouse = $v['privilege'] === 'USAGE' && $v['granted_on'] === 'WAREHOUSE';
-            $ownership = $v['privilege'] === 'OWNERSHIP' && (in_array($v['granted_on'], ['USER', 'ROLE']));
+        if (empty($grants)) {
+            $this->logger->debug('No grants found, checking if it is a user');
+            // Check if it's a user
+            $userExists = $connection->fetchAll(sprintf(
+                'SHOW USERS LIKE %s',
+                QueryBuilder::quote($name)
+            ));
+            if ($userExists) {
+                $this->logger->debug('User found, adding to removal list');
+                $userExists[0]['granted_by'] = $this->config->getTargetSnowflakeRole();
+                $result[self::USER][$name] = $userExists[0];
+            }
+            return $result;
+        }
 
-            return $ownership || $usageWarehouse;
-        });
+        $this->logger->debug('Processing grants');
+        // Filter grants by ownership
+        $ownershipGrants = array_filter(
+            $grants,
+            fn($v) => $v['privilege'] === self::OWNERSHIP
+        );
 
-        $mapGrants = [];
-        foreach ($filteredGrants as $filteredGrant) {
-            $mapGrants[$filteredGrant['granted_on']][$filteredGrant['name']] = $filteredGrant;
-        }
+        // Process owned roles and their users
+        foreach ($ownershipGrants as $grant) {
+            if ($grant['granted_on'] === self::ROLE) {
+                $this->logger->debug(sprintf('Found owned role: %s', $grant['name']));
+                $result[self::ROLE][$grant['name']] = $grant;
 
-        if (isset($mapGrants['ROLE'])) {
-            $roleGrants = $mapGrants['ROLE'];
-            foreach ($roleGrants as $roleGrant) {
-                $mapGrants = array_merge_recursive(
-                    $this->getDataToRemove($connection, $roleGrant['name']),
-                    $mapGrants,
-                );
+                // Get users owned by this role
+                $roleGrants = $connection->fetchAll(sprintf(
+                    'SHOW GRANTS TO ROLE %s',
+                    Helper::quoteIdentifier($grant['name'])
+                ));
+
+                foreach (array_filter(
+                    $roleGrants,
+                    fn($v) => $v['privilege'] === self::OWNERSHIP && $v['granted_on'] === self::USER
+                ) as $userGrant) {
+                    $this->logger->debug(
+                        sprintf('Found user owned by role %s: %s', $grant['name'], $userGrant['name'])
+                    );
+                    $result[self::USER][$userGrant['name']] = $userGrant;
+                }
+            } elseif ($grant['granted_on'] === self::USER) {
+                $this->logger->debug(sprintf('Found directly owned user: %s', $grant['name']));
+                $result[self::USER][$grant['name']] = $grant;
             }
         }
 
-        return $mapGrants;
+        $this->logger->debug(sprintf(
+            'Found %d roles and %d users to remove',
+            count($result[self::ROLE]),
+            count($result[self::USER])
+        ));
+        return $result;
     }
 
     private function getProjectUser(string $databaseRole): GrantToUser

From 9d1399df9731adae6d776c8eb4fe1b44698f08de Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Jodas?= <ondra.jodas@gmail.com>
Date: Tue, 4 Feb 2025 23:21:03 +0100
Subject: [PATCH 2/7] Add lowercase role name support

---
 src/Cleanup.php | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/Cleanup.php b/src/Cleanup.php
index 309d6cb..55f04cb 100644
--- a/src/Cleanup.php
+++ b/src/Cleanup.php
@@ -140,17 +140,25 @@ public function preMigration(string $mainRoleName): void
 
             if (!$dbExists) {
                 $this->logger->info(sprintf('Database %s does not exist, checking for role with same name', $database));
-                // If database doesn't exist, check if there's a role with the same name
-                $roleExists = $this->destinationConnection->fetchAll(sprintf(
-                    'SHOW ROLES LIKE %s;',
-                    QueryBuilder::quote($database)
-                ));
+                // Check if role exists with exact or lowercase name
+                $roleName = null;
+                foreach ([$database, strtolower($database)] as $nameVariant) {
+                    $roleExists = $this->destinationConnection->fetchAll(sprintf(
+                        'SHOW ROLES LIKE %s',
+                        QueryBuilder::quote($nameVariant)
+                    ));
+                    if ($roleExists) {
+                        $roleName = $nameVariant;
+                        break;
+                    }
+                }
 
-                if (!$roleExists) {
+                if ($roleName === null) {
                     continue;
                 }
-                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $database);
-                $roleToRemove = $database;
+
+                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $roleName);
+                $roleToRemove = $roleName;
             } else {
                 $this->logger->info(sprintf('Database %s exists, getting ownership role', $database));
                 $databaseRole = $this->destinationConnection->getOwnershipRoleOnDatabase($database);

From 0b4945894c54c4705b98bb95df70233074c77efb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Jodas?= <ondra.jodas@gmail.com>
Date: Wed, 5 Feb 2025 23:47:02 +0100
Subject: [PATCH 3/7] Add mainRoleName retrieval from database grants

---
 src/Cleanup.php | 70 ++++++++++++++++++++++++-------------------------
 1 file changed, 35 insertions(+), 35 deletions(-)

diff --git a/src/Cleanup.php b/src/Cleanup.php
index 55f04cb..fa05f68 100644
--- a/src/Cleanup.php
+++ b/src/Cleanup.php
@@ -37,7 +37,14 @@ public function sourceAccount(): void
             $databaseRole = $this->sourceConnection->getOwnershipRoleOnDatabase($database);
             $projectUser = $this->getProjectUser($databaseRole);
 
-            $data = $this->getDataToRemove($this->sourceConnection, $databaseRole);
+            $grantedOnDatabaseRole = $this->sourceConnection->fetchAll(sprintf(
+                'SHOW GRANTS ON ROLE %s',
+                Helper::quoteIdentifier($databaseRole)
+            ));
+            assert(count($grantedOnDatabaseRole) === 1);
+            $mainRoleName = current($grantedOnDatabaseRole)['granted_by'];
+
+            $data = $this->getDataToRemove($this->sourceConnection, $databaseRole, $mainRoleName);
 
             // drop roles
             $sqls[] = sprintf('DROP ROLE %s;', Helper::quoteIdentifier($databaseRole));
@@ -157,19 +164,16 @@ public function preMigration(string $mainRoleName): void
                     continue;
                 }
 
-                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $roleName);
-                $roleToRemove = $roleName;
+                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $roleName, $mainRoleName);
             } else {
                 $this->logger->info(sprintf('Database %s exists, getting ownership role', $database));
                 $databaseRole = $this->destinationConnection->getOwnershipRoleOnDatabase($database);
-                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $databaseRole);
-                $roleToRemove = $databaseRole;
+                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $databaseRole, $mainRoleName);
             }
 
             // First revoke all future grants from roles
             foreach ($dataToRemove['ROLE'] as $role) {
-                $this->switchRole($role['granted_by'], $mainRoleName, $sqls, $currentRole);
-
+                $this->destinationConnection->useRole($role['granted_by']);
                 $futureGrants = array_map(
                     fn(array $v) => FutureGrantToRole::fromArray($v),
                     $this->destinationConnection->fetchAll(sprintf(
@@ -178,6 +182,9 @@ public function preMigration(string $mainRoleName): void
                     ))
                 );
 
+                if ($futureGrants) {
+                    $this->switchRole($role['granted_by'], $mainRoleName, $sqls, $currentRole);
+                }
                 foreach ($futureGrants as $futureGrant) {
                     $sqls[] = sprintf(
                         'REVOKE %s ON FUTURE TABLES IN SCHEMA %s FROM ROLE %s;',
@@ -213,11 +220,7 @@ public function preMigration(string $mainRoleName): void
             }
 
             // Drop database role and handle database if exists
-            $this->switchRole($this->config->getTargetSnowflakeRole(), $mainRoleName, $sqls, $currentRole);
-            $sqls[] = sprintf(
-                'DROP ROLE IF EXISTS %s;',
-                Helper::quoteIdentifier($roleToRemove)
-            );
+            $this->switchRole($mainRoleName, $mainRoleName, $sqls, $currentRole);
 
             if ($dbExists) {
                 $sqls[] = sprintf(
@@ -336,7 +339,7 @@ private function switchRole(string $role, string $mainRoleName, array &$sqls, ?s
         }
     }
 
-    private function getDataToRemove(Connection $connection, string $name): array
+    private function getDataToRemove(Connection $connection, string $name, string $mainRoleName): array
     {
         $this->logger->debug(sprintf('Getting data to remove for: %s', $name));
         $result = [
@@ -350,21 +353,6 @@ private function getDataToRemove(Connection $connection, string $name): array
             Helper::quoteIdentifier($name)
         ));
 
-        if (empty($grants)) {
-            $this->logger->debug('No grants found, checking if it is a user');
-            // Check if it's a user
-            $userExists = $connection->fetchAll(sprintf(
-                'SHOW USERS LIKE %s',
-                QueryBuilder::quote($name)
-            ));
-            if ($userExists) {
-                $this->logger->debug('User found, adding to removal list');
-                $userExists[0]['granted_by'] = $this->config->getTargetSnowflakeRole();
-                $result[self::USER][$name] = $userExists[0];
-            }
-            return $result;
-        }
-
         $this->logger->debug('Processing grants');
         // Filter grants by ownership
         $ownershipGrants = array_filter(
@@ -376,7 +364,7 @@ private function getDataToRemove(Connection $connection, string $name): array
         foreach ($ownershipGrants as $grant) {
             if ($grant['granted_on'] === self::ROLE) {
                 $this->logger->debug(sprintf('Found owned role: %s', $grant['name']));
-                $result[self::ROLE][$grant['name']] = $grant;
+                $result[self::ROLE][] = $grant;
 
                 // Get users owned by this role
                 $roleGrants = $connection->fetchAll(sprintf(
@@ -391,19 +379,31 @@ private function getDataToRemove(Connection $connection, string $name): array
                     $this->logger->debug(
                         sprintf('Found user owned by role %s: %s', $grant['name'], $userGrant['name'])
                     );
-                    $result[self::USER][$userGrant['name']] = $userGrant;
+                    $result[self::USER][] = $userGrant;
                 }
             } elseif ($grant['granted_on'] === self::USER) {
                 $this->logger->debug(sprintf('Found directly owned user: %s', $grant['name']));
-                $result[self::USER][$grant['name']] = $grant;
+                $result[self::USER][] = $grant;
             }
         }
 
-        $this->logger->debug(sprintf(
-            'Found %d roles and %d users to remove',
-            count($result[self::ROLE]),
-            count($result[self::USER])
+        // Check if user with same name exists
+        $userExists = $connection->fetchAll(sprintf(
+            'SHOW USERS LIKE %s',
+            QueryBuilder::quote($name)
         ));
+        if ($userExists) {
+            $this->logger->debug(sprintf('Found user with same name: %s', $name));
+            $userExists[0]['granted_by'] = $mainRoleName;
+            $result[self::USER][] = $userExists[0];
+        }
+
+        // Add the role itself to be removed
+        $result[self::ROLE][] = [
+            'name' => $name,
+            'granted_by' => $mainRoleName,
+        ];
+
         return $result;
     }
 

From 93c064a9e76d9b5f171e446e7f7f0b2bbe547aed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Jodas?= <ondra.jodas@gmail.com>
Date: Thu, 6 Feb 2025 00:06:14 +0100
Subject: [PATCH 4/7] Simplify getDataToRemove method and improve code
 readability

---
 src/Cleanup.php | 115 +++++++++++++++++++++++-------------------------
 1 file changed, 55 insertions(+), 60 deletions(-)

diff --git a/src/Cleanup.php b/src/Cleanup.php
index fa05f68..1a8135c 100644
--- a/src/Cleanup.php
+++ b/src/Cleanup.php
@@ -183,7 +183,7 @@ public function preMigration(string $mainRoleName): void
                 );
 
                 if ($futureGrants) {
-                    $this->switchRole($role['granted_by'], $mainRoleName, $sqls, $currentRole);
+                    [$sqls, $currentRole] = $this->switchRole($role['granted_by'], $mainRoleName, $sqls, $currentRole);
                 }
                 foreach ($futureGrants as $futureGrant) {
                     $sqls[] = sprintf(
@@ -200,7 +200,7 @@ public function preMigration(string $mainRoleName): void
                 $this->logger->info(sprintf('Dropping %d users', count($dataToRemove['USER'])));
             }
             foreach ($dataToRemove['USER'] as $user) {
-                $this->switchRole($user['granted_by'], $mainRoleName, $sqls, $currentRole);
+                [$sqls, $currentRole] = $this->switchRole($user['granted_by'], $mainRoleName, $sqls, $currentRole);
                 $sqls[] = sprintf(
                     'DROP USER IF EXISTS %s;',
                     Helper::quoteIdentifier($user['name'])
@@ -212,7 +212,7 @@ public function preMigration(string $mainRoleName): void
                 $this->logger->info(sprintf('Dropping %d roles', count($dataToRemove['ROLE'])));
             }
             foreach ($dataToRemove['ROLE'] as $role) {
-                $this->switchRole($role['granted_by'], $mainRoleName, $sqls, $currentRole);
+                [$sqls, $currentRole] = $this->switchRole($role['granted_by'], $mainRoleName, $sqls, $currentRole);
                 $sqls[] = sprintf(
                     'DROP ROLE IF EXISTS %s;',
                     Helper::quoteIdentifier($role['name'])
@@ -220,7 +220,7 @@ public function preMigration(string $mainRoleName): void
             }
 
             // Drop database role and handle database if exists
-            $this->switchRole($mainRoleName, $mainRoleName, $sqls, $currentRole);
+            [$sqls, $currentRole] = $this->switchRole($mainRoleName, $mainRoleName, $sqls, $currentRole);
 
             if ($dbExists) {
                 $sqls[] = sprintf(
@@ -306,39 +306,6 @@ public function postMigration(): void
         }
     }
 
-    private function switchRole(string $role, string $mainRoleName, array &$sqls, ?string &$currentRole): void
-    {
-        if ($currentRole === $role) {
-            return;
-        }
-
-        try {
-            $this->logger->debug(sprintf('Switching to role: %s', $role));
-            $this->destinationConnection->useRole($role);
-            $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($role));
-            $currentRole = $role;
-        } catch (RuntimeException $e) {
-            $this->logger->info(sprintf(
-                'Cannot switch directly to role %s, trying through main role %s',
-                $role,
-                $mainRoleName
-            ));
-            $this->destinationConnection->useRole($mainRoleName);
-            if ($currentRole !== $mainRoleName) {
-                $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($mainRoleName));
-                $currentRole = $mainRoleName;
-            }
-            $this->destinationConnection->grantRoleToUser(
-                $this->config->getTargetSnowflakeUser(),
-                $role
-            );
-            $this->logger->debug(sprintf('Granted role %s to user, switching to it', $role));
-            $this->destinationConnection->useRole($role);
-            $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($role));
-            $currentRole = $role;
-        }
-    }
-
     private function getDataToRemove(Connection $connection, string $name, string $mainRoleName): array
     {
         $this->logger->debug(sprintf('Getting data to remove for: %s', $name));
@@ -347,47 +314,39 @@ private function getDataToRemove(Connection $connection, string $name, string $m
             self::USER => [],
         ];
 
-        // Get all grants for the name (both as role and user)
         $grants = $connection->fetchAll(sprintf(
             'SHOW GRANTS TO ROLE %s',
             Helper::quoteIdentifier($name)
         ));
 
-        $this->logger->debug('Processing grants');
-        // Filter grants by ownership
         $ownershipGrants = array_filter(
             $grants,
             fn($v) => $v['privilege'] === self::OWNERSHIP
         );
 
-        // Process owned roles and their users
+        $result = $this->processOwnershipGrants($ownershipGrants, $result);
+        $result = $this->addUserIfExists($connection, $name, $mainRoleName, $result);
+        $result = $this->addRoleToRemove($name, $mainRoleName, $result);
+
+        return $result;
+    }
+
+    private function processOwnershipGrants(array $ownershipGrants, array $result): array
+    {
         foreach ($ownershipGrants as $grant) {
             if ($grant['granted_on'] === self::ROLE) {
                 $this->logger->debug(sprintf('Found owned role: %s', $grant['name']));
                 $result[self::ROLE][] = $grant;
-
-                // Get users owned by this role
-                $roleGrants = $connection->fetchAll(sprintf(
-                    'SHOW GRANTS TO ROLE %s',
-                    Helper::quoteIdentifier($grant['name'])
-                ));
-
-                foreach (array_filter(
-                    $roleGrants,
-                    fn($v) => $v['privilege'] === self::OWNERSHIP && $v['granted_on'] === self::USER
-                ) as $userGrant) {
-                    $this->logger->debug(
-                        sprintf('Found user owned by role %s: %s', $grant['name'], $userGrant['name'])
-                    );
-                    $result[self::USER][] = $userGrant;
-                }
             } elseif ($grant['granted_on'] === self::USER) {
                 $this->logger->debug(sprintf('Found directly owned user: %s', $grant['name']));
                 $result[self::USER][] = $grant;
             }
         }
+        return $result;
+    }
 
-        // Check if user with same name exists
+    private function addUserIfExists(Connection $connection, string $name, string $mainRoleName, array $result): array
+    {
         $userExists = $connection->fetchAll(sprintf(
             'SHOW USERS LIKE %s',
             QueryBuilder::quote($name)
@@ -397,16 +356,52 @@ private function getDataToRemove(Connection $connection, string $name, string $m
             $userExists[0]['granted_by'] = $mainRoleName;
             $result[self::USER][] = $userExists[0];
         }
+        return $result;
+    }
 
-        // Add the role itself to be removed
+    private function addRoleToRemove(string $name, string $mainRoleName, array $result): array
+    {
         $result[self::ROLE][] = [
             'name' => $name,
             'granted_by' => $mainRoleName,
         ];
-
         return $result;
     }
 
+    private function switchRole(string $role, string $mainRoleName, array $sqls, ?string $currentRole): array
+    {
+        if ($currentRole === $role) {
+            return [$sqls, $currentRole];
+        }
+
+        try {
+            $this->logger->debug(sprintf('Switching to role: %s', $role));
+            $this->destinationConnection->useRole($role);
+            $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($role));
+            $currentRole = $role;
+        } catch (RuntimeException $e) {
+            $this->logger->info(sprintf(
+                'Cannot switch directly to role %s, trying through main role %s',
+                $role,
+                $mainRoleName
+            ));
+            $this->destinationConnection->useRole($mainRoleName);
+            if ($currentRole !== $mainRoleName) {
+                $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($mainRoleName));
+                $currentRole = $mainRoleName;
+            }
+            $this->destinationConnection->grantRoleToUser(
+                $this->config->getTargetSnowflakeUser(),
+                $role
+            );
+            $this->logger->debug(sprintf('Granted role %s to user, switching to it', $role));
+            $this->destinationConnection->useRole($role);
+            $sqls[] = sprintf('USE ROLE %s;', Helper::quoteIdentifier($role));
+            $currentRole = $role;
+        }
+        return [$sqls, $currentRole];
+    }
+
     private function getProjectUser(string $databaseRole): GrantToUser
     {
         $grantsOfRole = $this->destinationConnection->fetchAll(sprintf('SHOW GRANTS OF ROLE %s', $databaseRole));

From 3af4bb31ea1da55b77afd6e2c2ec5f1c560b951b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Jodas?= <ondra.jodas@gmail.com>
Date: Mon, 10 Feb 2025 22:47:09 +0100
Subject: [PATCH 5/7] changes after review

---
 src/Cleanup.php | 45 +++++++++++++++++----------------------------
 1 file changed, 17 insertions(+), 28 deletions(-)

diff --git a/src/Cleanup.php b/src/Cleanup.php
index 1a8135c..1b65862 100644
--- a/src/Cleanup.php
+++ b/src/Cleanup.php
@@ -49,7 +49,7 @@ public function sourceAccount(): void
             // drop roles
             $sqls[] = sprintf('DROP ROLE %s;', Helper::quoteIdentifier($databaseRole));
 
-            $roles = array_map(fn(array $v) => GrantToRole::fromArray($v), $data['ROLE'] ?? []);
+            $roles = array_map(fn(array $v) => GrantToRole::fromArray($v), $data[self::ROLE] ?? []);
             foreach ($roles as $role) {
                 $futureGrants = array_map(
                     fn(array $v) => FutureGrantToRole::fromArray($v),
@@ -77,7 +77,7 @@ public function sourceAccount(): void
             // drop users
             $sqls[] = sprintf('DROP USER %s;', Helper::quoteIdentifier($projectUser->getGranteeName()));
 
-            $users = array_map(fn(array $v) => GrantToRole::fromArray($v), $data['USER'] ?? []);
+            $users = array_map(fn(array $v) => GrantToRole::fromArray($v), $data[self::USER] ?? []);
             foreach ($users as $user) {
                 $sqls[] = sprintf(
                     'DROP USER %s;',
@@ -110,7 +110,7 @@ public function preMigration(string $mainRoleName): void
         ));
         $mainRoleExistsOnTargetUser = array_reduce(
             $grantsToTargetUser,
-            fn ($found, $v) => $found || $v['role'] === $mainRoleName,
+            fn ($found, $v) => $found || $v[self::ROLE] === $mainRoleName,
             false,
         );
 
@@ -172,7 +172,7 @@ public function preMigration(string $mainRoleName): void
             }
 
             // First revoke all future grants from roles
-            foreach ($dataToRemove['ROLE'] as $role) {
+            foreach ($dataToRemove[self::ROLE] as $role) {
                 $this->destinationConnection->useRole($role['granted_by']);
                 $futureGrants = array_map(
                     fn(array $v) => FutureGrantToRole::fromArray($v),
@@ -196,10 +196,10 @@ public function preMigration(string $mainRoleName): void
             }
 
             // Drop users owned by roles
-            if (!empty($dataToRemove['USER'])) {
-                $this->logger->info(sprintf('Dropping %d users', count($dataToRemove['USER'])));
+            if (!empty($dataToRemove[self::USER])) {
+                $this->logger->info(sprintf('Dropping %d users', count($dataToRemove[self::USER])));
             }
-            foreach ($dataToRemove['USER'] as $user) {
+            foreach ($dataToRemove[self::USER] as $user) {
                 [$sqls, $currentRole] = $this->switchRole($user['granted_by'], $mainRoleName, $sqls, $currentRole);
                 $sqls[] = sprintf(
                     'DROP USER IF EXISTS %s;',
@@ -208,10 +208,10 @@ public function preMigration(string $mainRoleName): void
             }
 
             // Drop roles
-            if (!empty($dataToRemove['ROLE'])) {
-                $this->logger->info(sprintf('Dropping %d roles', count($dataToRemove['ROLE'])));
+            if (!empty($dataToRemove[self::ROLE])) {
+                $this->logger->info(sprintf('Dropping %d roles', count($dataToRemove[self::ROLE])));
             }
-            foreach ($dataToRemove['ROLE'] as $role) {
+            foreach ($dataToRemove[self::ROLE] as $role) {
                 [$sqls, $currentRole] = $this->switchRole($role['granted_by'], $mainRoleName, $sqls, $currentRole);
                 $sqls[] = sprintf(
                     'DROP ROLE IF EXISTS %s;',
@@ -308,7 +308,7 @@ public function postMigration(): void
 
     private function getDataToRemove(Connection $connection, string $name, string $mainRoleName): array
     {
-        $this->logger->debug(sprintf('Getting data to remove for: %s', $name));
+        $this->logger->debug(sprintf('Getting data to remove for role: %s', $name));
         $result = [
             self::ROLE => [],
             self::USER => [],
@@ -325,8 +325,7 @@ private function getDataToRemove(Connection $connection, string $name, string $m
         );
 
         $result = $this->processOwnershipGrants($ownershipGrants, $result);
-        $result = $this->addUserIfExists($connection, $name, $mainRoleName, $result);
-        $result = $this->addRoleToRemove($name, $mainRoleName, $result);
+        $result = $this->addProjectRoleAndUserToRemove($name, $mainRoleName, $result);
 
         return $result;
     }
@@ -345,26 +344,16 @@ private function processOwnershipGrants(array $ownershipGrants, array $result):
         return $result;
     }
 
-    private function addUserIfExists(Connection $connection, string $name, string $mainRoleName, array $result): array
-    {
-        $userExists = $connection->fetchAll(sprintf(
-            'SHOW USERS LIKE %s',
-            QueryBuilder::quote($name)
-        ));
-        if ($userExists) {
-            $this->logger->debug(sprintf('Found user with same name: %s', $name));
-            $userExists[0]['granted_by'] = $mainRoleName;
-            $result[self::USER][] = $userExists[0];
-        }
-        return $result;
-    }
-
-    private function addRoleToRemove(string $name, string $mainRoleName, array $result): array
+    private function addProjectRoleAndUserToRemove(string $name, string $mainRoleName, array $result): array
     {
         $result[self::ROLE][] = [
             'name' => $name,
             'granted_by' => $mainRoleName,
         ];
+        $result[self::USER][] = [
+            'name' => $name,
+            'granted_by' => $mainRoleName,
+        ];
         return $result;
     }
 

From 2f2cc3b2605c9e0d89a20ce4e6bb7e4bb021e04c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Jodas?= <ondra.jodas@gmail.com>
Date: Mon, 10 Feb 2025 22:49:00 +0100
Subject: [PATCH 6/7] rename getDataToRemove to getDataToRemoveForRole

---
 src/Cleanup.php | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/Cleanup.php b/src/Cleanup.php
index 1b65862..95ba5c9 100644
--- a/src/Cleanup.php
+++ b/src/Cleanup.php
@@ -44,7 +44,7 @@ public function sourceAccount(): void
             assert(count($grantedOnDatabaseRole) === 1);
             $mainRoleName = current($grantedOnDatabaseRole)['granted_by'];
 
-            $data = $this->getDataToRemove($this->sourceConnection, $databaseRole, $mainRoleName);
+            $data = $this->getDataToRemoveForRole($this->sourceConnection, $databaseRole, $mainRoleName);
 
             // drop roles
             $sqls[] = sprintf('DROP ROLE %s;', Helper::quoteIdentifier($databaseRole));
@@ -110,7 +110,7 @@ public function preMigration(string $mainRoleName): void
         ));
         $mainRoleExistsOnTargetUser = array_reduce(
             $grantsToTargetUser,
-            fn ($found, $v) => $found || $v[self::ROLE] === $mainRoleName,
+            fn ($found, $v) => $found || $v['role'] === $mainRoleName,
             false,
         );
 
@@ -164,11 +164,11 @@ public function preMigration(string $mainRoleName): void
                     continue;
                 }
 
-                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $roleName, $mainRoleName);
+                $dataToRemove = $this->getDataToRemoveForRole($this->destinationConnection, $roleName, $mainRoleName);
             } else {
                 $this->logger->info(sprintf('Database %s exists, getting ownership role', $database));
                 $databaseRole = $this->destinationConnection->getOwnershipRoleOnDatabase($database);
-                $dataToRemove = $this->getDataToRemove($this->destinationConnection, $databaseRole, $mainRoleName);
+                $dataToRemove = $this->getDataToRemoveForRole($this->destinationConnection, $databaseRole, $mainRoleName);
             }
 
             // First revoke all future grants from roles
@@ -306,7 +306,7 @@ public function postMigration(): void
         }
     }
 
-    private function getDataToRemove(Connection $connection, string $name, string $mainRoleName): array
+    private function getDataToRemoveForRole(Connection $connection, string $name, string $mainRoleName): array
     {
         $this->logger->debug(sprintf('Getting data to remove for role: %s', $name));
         $result = [

From 6cdf4d6df67060d0ad5bb88112354b4fcc05133c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Jodas?= <ondra.jodas@gmail.com>
Date: Mon, 10 Feb 2025 22:55:00 +0100
Subject: [PATCH 7/7] fix phpcs

---
 src/Cleanup.php | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/Cleanup.php b/src/Cleanup.php
index 95ba5c9..ce16e09 100644
--- a/src/Cleanup.php
+++ b/src/Cleanup.php
@@ -168,7 +168,11 @@ public function preMigration(string $mainRoleName): void
             } else {
                 $this->logger->info(sprintf('Database %s exists, getting ownership role', $database));
                 $databaseRole = $this->destinationConnection->getOwnershipRoleOnDatabase($database);
-                $dataToRemove = $this->getDataToRemoveForRole($this->destinationConnection, $databaseRole, $mainRoleName);
+                $dataToRemove = $this->getDataToRemoveForRole(
+                    $this->destinationConnection,
+                    $databaseRole,
+                    $mainRoleName,
+                );
             }
 
             // First revoke all future grants from roles