Merge pull request #51 from keboola/roman-pst-2064

PST-2064: Preserve same dest. snflk workspace for snflk wr configs
keboola · Sep 25, 2024 · 6d6bacf · 6d6bacf
2 parents 32596e7 + 3e576a6
commit 6d6bacf
Show file tree

Hide file tree

Showing 3 changed files with 300 additions and 0 deletions.
diff --git a/src/Component.php b/src/Component.php
@@ -75,6 +75,7 @@ protected function run(): void
             $sourceJobRunner,
             $destJobRunner,
             $sourceProjectClient,
+            $destProjectClient,
             $migrationsClient,
             $destProjectClient->getApiUrl(),
             $destProjectClient->getTokenString(),

diff --git a/src/Migrate.php b/src/Migrate.php
@@ -12,6 +12,7 @@
 use Keboola\StorageApi\Client as StorageClient;
 use Keboola\StorageApi\Components;
 use Keboola\StorageApi\DevBranches;
+use Keboola\StorageApi\Options\Components\Configuration;
 use Keboola\Syrup\ClientException as SyrupClientException;
 use Psr\Log\LoggerInterface;
 
@@ -25,6 +26,8 @@ class Migrate
 
     private StorageClient $sourceProjectStorageClient;
 
+    private StorageClient $destProjectStorageClient;
+
     private MigrationsClient $migrationsClient;
 
     private string $sourceProjectUrl;
@@ -52,15 +55,25 @@ class Migrate
         'gooddata-writer',
     ];
 
+    public const SNOWFLAKE_WRITER_COMPONENT_IDS = [
+        'keboola.wr-db-snowflake', // aws
+        'keboola.wr-snowflake-blob-storage', // azure
+        'keboola.wr-db-snowflake-gcs', // gcp
+        'keboola.wr-db-snowflake-gcs-s3', // gcp with s3
+    ];
+
     private string $migrateDataMode;
 
     private array $db;
 
+    private array $migratedSnowflakeWorkspaces = [];
+
     public function __construct(
         Config $config,
         JobRunner $sourceJobRunner,
         JobRunner $destJobRunner,
         StorageClient $sourceProjectStorageClient,
+        StorageClient $destProjectStorageClient,
         MigrationsClient $migrationsClient,
         string $destinationProjectUrl,
         string $destinationProjectToken,
@@ -69,6 +82,7 @@ public function __construct(
         $this->sourceJobRunner = $sourceJobRunner;
         $this->destJobRunner = $destJobRunner;
         $this->sourceProjectStorageClient = $sourceProjectStorageClient;
+        $this->destProjectStorageClient = $destProjectStorageClient;
         $this->migrationsClient = $migrationsClient;
         $this->sourceProjectUrl = $config->getSourceProjectUrl();
         $this->sourceProjectToken = $config->getSourceProjectToken();
@@ -203,6 +217,15 @@ private function migrateSecrets(): void
                         $this->dryRun
                     );
 
+                if (in_array($component['id'], self::SNOWFLAKE_WRITER_COMPONENT_IDS, true)) {
+                    $this->preserveProperSnowflakeWorkspace(
+                        $component['id'],
+                        $config['id'],
+                        $response['data']['componentId'],
+                        $response['data']['configId']
+                    );
+                }
+
                 $message = $response['message'];
                 if ($this->dryRun) {
                     $message = '[dry-run] ' . $message;
@@ -316,4 +339,55 @@ private function getRestoreConfigData(array $restoreCredentials): array
             throw new UserException('Unrecognized restore credentials.');
         }
     }
+
+    private function preserveProperSnowflakeWorkspace(
+        string $sourceComponentId,
+        string $sourceConfigurationId,
+        string $destinationComponentId,
+        string $destinationConfigurationId
+    ): void {
+        if ($this->dryRun) {
+            return;
+        }
+        $sourceComponentsApi = new Components($this->sourceProjectStorageClient);
+        $sourceConfigurationData = (array) $sourceComponentsApi
+            ->getConfiguration($sourceComponentId, $sourceConfigurationId);
+
+        $destinationComponentsApi = new Components($this->destProjectStorageClient);
+        $destinationConfigurationData = (array) $destinationComponentsApi
+            ->getConfiguration($destinationComponentId, $destinationConfigurationId);
+
+        $snowflakeUser = $sourceConfigurationData['configuration']['parameters']['db']['user'];
+        $migratedWorkspaceParameters = $this->migratedSnowflakeWorkspaces[$snowflakeUser] ?? null;
+
+        if ($migratedWorkspaceParameters) {
+            // Use the existing Snowflake workspace from a previous configuration that has the same source workspace
+            $destinationConfigurationData['configuration']['parameters']['db'] = $migratedWorkspaceParameters;
+
+            $destinationConfiguration = (new Configuration())
+                ->setConfigurationId($destinationConfigurationId)
+                ->setComponentId($destinationComponentId)
+                ->setName($destinationConfigurationData['name'])
+                ->setDescription($destinationConfigurationData['description'])
+                ->setIsDisabled($destinationConfigurationData['isDisabled'])
+                ->setConfiguration($destinationConfigurationData['configuration']);
+
+            $destinationComponentsApi->updateConfiguration($destinationConfiguration);
+
+            $this->logger->info(
+                sprintf(
+                    "Used existing Snowflake workspace '%s' for configuration with ID '%s' (%s).",
+                    $migratedWorkspaceParameters['user'],
+                    $destinationConfigurationId,
+                    $destinationComponentId,
+                ),
+                ['secrets']
+            );
+            return;
+        }
+
+        // Store Snowflake workspace for next configurations
+        $workspaceParameters = $destinationConfigurationData['configuration']['parameters']['db'];
+        $this->migratedSnowflakeWorkspaces[$snowflakeUser] = $workspaceParameters;
+    }
 }