syncing up to 5ab31378261546b692f5f1e224680d9cf31b59e7

Co-authored-by: Joey Greco <[email protected]>
superblocksteam · Nov 7, 2024 · 9e0ea94 · 9e0ea94
1 parent 71607cd
commit 9e0ea94
Show file tree

Hide file tree

Showing 6 changed files with 171 additions and 12 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Allow branch name to be given in workflow URLs as a query param: `fetch.branch_name`
 - Upgrade Google Cloud Secret Manager Go package to `v1.14.2` (fixes "certificate_config.json: permission denied" error connecting to GCP secret manager)
+- Added support for key-pair authentication in Snowflake Plugin
 
 ## v1.16.0
 

diff --git a/internal/registration/registration.go b/internal/registration/registration.go
@@ -48,7 +48,7 @@ var SUPERBLOCKS_PLUGIN_VERSIONS = map[string][]string{
 	"rockset":            {"0.0.7"},
 	"s3":                 {"0.0.11"},
 	"salesforce":         {"0.0.1"},
-	"snowflake":          {"0.0.7"},
+	"snowflake":          {"0.0.8"},
 	"superblocks-ocr":    {"0.0.1"},
 	"workflow":           {"0.0.4"},
 	"redis":              {"0.0.1"},

diff --git a/internal/registration/registration_test.go b/internal/registration/registration_test.go
@@ -68,5 +68,5 @@ func TestRegistrationRequest(t *testing.T) {
 	assert.Equal(t, "https://foo.bar.com/api/v1/agents/register", request.URL.String())
 	assert.Equal(t, http.MethodPost, request.Method)
 	assert.Equal(t, expectedRequestHeaders, request.Header)
-	assert.Equal(t, "{\"pluginVersions\":{\"athena\":[\"0.0.1\"],\"bigquery\":[\"0.0.7\"],\"cockroachdb\":[\"0.0.2\"],\"confluent\":[\"0.0.1\"],\"cosmosdb\":[\"0.0.1\"],\"couchbase\":[\"0.0.1\"],\"databricks\":[\"0.0.1\"],\"dynamodb\":[\"0.0.7\"],\"email\":[\"0.0.7\"],\"gcs\":[\"0.0.1\"],\"graphql\":[\"0.0.10\"],\"gsheets\":[\"0.0.18\"],\"javascript\":[\"0.0.9\"],\"kafka\":[\"0.0.1\"],\"kinesis\":[\"0.0.1\"],\"mariadb\":[\"0.0.11\"],\"mongodb\":[\"0.0.7\"],\"mssql\":[\"0.0.10\"],\"mysql\":[\"0.0.11\"],\"openai\":[\"0.0.3\"],\"oracledb\":[\"0.0.1\"],\"postgres\":[\"0.0.11\"],\"python\":[\"0.0.7\"],\"redis\":[\"0.0.1\"],\"redpanda\":[\"0.0.1\"],\"redshift\":[\"0.0.7\"],\"restapi\":[\"0.0.13\"],\"restapiintegration\":[\"0.0.14\"],\"rockset\":[\"0.0.7\"],\"s3\":[\"0.0.11\"],\"salesforce\":[\"0.0.1\"],\"smtp\":[\"0.0.1\"],\"snowflake\":[\"0.0.7\"],\"superblocks-ocr\":[\"0.0.1\"],\"workflow\":[\"0.0.4\"]},\"type\":2,\"tags\":{\"foo\":[\"bar\"]},\"signingKeyId\":\"123\",\"verificationKeyIds\":[\"123\"],\"verificationKeys\":{\"123\":{\"algorithm\":\"ALGORITHM_RSA\",\"key\":\"key1\"}}}\n", string(body))
+	assert.Equal(t, "{\"pluginVersions\":{\"athena\":[\"0.0.1\"],\"bigquery\":[\"0.0.7\"],\"cockroachdb\":[\"0.0.2\"],\"confluent\":[\"0.0.1\"],\"cosmosdb\":[\"0.0.1\"],\"couchbase\":[\"0.0.1\"],\"databricks\":[\"0.0.1\"],\"dynamodb\":[\"0.0.7\"],\"email\":[\"0.0.7\"],\"gcs\":[\"0.0.1\"],\"graphql\":[\"0.0.10\"],\"gsheets\":[\"0.0.18\"],\"javascript\":[\"0.0.9\"],\"kafka\":[\"0.0.1\"],\"kinesis\":[\"0.0.1\"],\"mariadb\":[\"0.0.11\"],\"mongodb\":[\"0.0.7\"],\"mssql\":[\"0.0.10\"],\"mysql\":[\"0.0.11\"],\"openai\":[\"0.0.3\"],\"oracledb\":[\"0.0.1\"],\"postgres\":[\"0.0.11\"],\"python\":[\"0.0.7\"],\"redis\":[\"0.0.1\"],\"redpanda\":[\"0.0.1\"],\"redshift\":[\"0.0.7\"],\"restapi\":[\"0.0.13\"],\"restapiintegration\":[\"0.0.14\"],\"rockset\":[\"0.0.7\"],\"s3\":[\"0.0.11\"],\"salesforce\":[\"0.0.1\"],\"smtp\":[\"0.0.1\"],\"snowflake\":[\"0.0.8\"],\"superblocks-ocr\":[\"0.0.1\"],\"workflow\":[\"0.0.4\"]},\"type\":2,\"tags\":{\"foo\":[\"bar\"]},\"signingKeyId\":\"123\",\"verificationKeyIds\":[\"123\"],\"verificationKeys\":{\"123\":{\"algorithm\":\"ALGORITHM_RSA\",\"key\":\"key1\"}}}\n", string(body))
 }
diff --git a/workers/javascript/packages/plugins/snowflake/src/util.test.ts b/workers/javascript/packages/plugins/snowflake/src/util.test.ts
@@ -82,6 +82,120 @@ describe('connectionOptionsFromDatasourceConfiguration', () => {
     });
   });
 
+  it('works for key-pair without password', async () => {
+    const datasourceConfiguration = {
+      connectionType: 'key-pair',
+      authentication: {
+        custom: {
+          account: { value: 'account' }
+        }
+      },
+      keyPair: {
+        privateKey: 'pk'
+      }
+    } as SnowflakeDatasourceConfiguration;
+    const connectionOptions = connectionOptionsFromDatasourceConfiguration(datasourceConfiguration);
+
+    expect(connectionOptions).toEqual({
+      account: 'account',
+      authenticator: 'SNOWFLAKE_JWT',
+      privateKey: 'pk'
+    });
+  });
+
+  // openssl genpkey -algorithm RSA -outform PEM -aes256 -pass pass:foo -pkeyopt rsa_keygen_bits:2048 | pbcopy
+  it('works for key-pair with password', async () => {
+    const datasourceConfiguration = {
+      connectionType: 'key-pair',
+      authentication: {
+        custom: {
+          account: { value: 'account' }
+        }
+      },
+      keyPair: {
+        privateKey: `-----BEGIN ENCRYPTED PRIVATE KEY-----
+MIIFNTBfBgkqhkiG9w0BBQ0wUjAxBgkqhkiG9w0BBQwwJAQQae5CUXLDXyT0NKTd
+nC9rFwICCAAwDAYIKoZIhvcNAgkFADAdBglghkgBZQMEASoEENiexhDUrSkVh8Lq
+CLysnhEEggTQDRGIgcyjfxVPgdsnBqXIVJwk6CbEw6u3rA3sS53si4m5wkZEnuOt
+BgC79tbGSCTpAjbb226D04DofI3vt/VaeCvwrS3bObCcY0Y1E1f+KkhIwVz4WG1I
+0OdnHdQ0GA5Zda0265tUoOPb2BTZdSj6RWlXNBjQZSSpekCJGfvkOtc0xA1tiFo1
+u1qFKHs8dXOjCYLY00U39EhAuC7oa12l8BMXSyblyWE8HFFz2qVwI2bdg+JmseFe
+oV41Zs0kTPwuWTzPSi+fa8Lo28v/lysTNozeY21AFUmoIH5Fnu4l58s9hzk45uHJ
+1eLI+k9l62tbA935B47AOnSDlppUzU83y2iMCvJduWBhuQ76y+AnLybDOpQDnfX9
+I9kmto22s/AUMohi2F8+V6B9cCWQ3rVDgoLqfGh3EzJfRtWWSynmLmbf5vgz8yZ7
+uZy/UN1zELWHciN0hRWRlZ56UxGgUj27nA88VsarT9RQP/FN+XR/S7p+6UPK8ul0
+BnSZu3HroBSj/hV3NrbaYE5i3SUhp8A/n3RQc1LaTk+ZhaXWEOYqUoRjB9AUEAjG
+/7MglilS/PCVt+ofwUUtuwbFp4+F6B1eo3bFTT310BNYoc7/pRgh2pUYvzFByqhQ
+RlLVzK66KYE+O6TAF5BJhlSrt/NvzYELYXJb91UoltqIC0rtMMADwrJKfn4Lr/Zl
+ejpHiUcHB9n53ej0DACJt1gxa57Uqum2S0ds1JZ9Mmb+aNy/lO/e9lbkQo4z1qrO
+LH6oRyVdHIOEPARBYuSGZ4LOAg0SnmRTEmhl4MG8fEifUXhcjrOyb0ynjns4s+6l
+wA3W/+HegjTSvPpAPB1jeDx9gGaAje5rGEbnqqqLj7F4SrakD5VxAU5wdzFpr26b
+026Nmdf66UcYcKAjTDOAQjtqjXA4J+RS73P41yt7YRAji4cPt29t96iiLFYRLQRo
+qw3qRz1EWUrLDNNfLs2HDD3wDuSCX+UaObgPpY11R2yTgl5S3FkXRanLdJ+0gqdg
+iWh01qfIo8/GTnc5WoonsHQLSp1oV6UwCcptjOVg35UwiABqBUpmLog6vSqRSgR6
+arakRDZ84kzxKtJIMYpS9UOSkBgUTKfMVikPj8eeuZec/uJWUGHsBclpdXle3Cye
+EcZnVOuC7/+bNMF6m6qV1vFomeBoq0wfhkAfNStwEq2BJC8ME84hGmLNpXF8tC46
+UnEmYjGIpuL+qBKVOuta1XMqbY1D0v6OUzYvE3RxKGHu7QyZCr3VzNMVMqQkeqhS
+KvnXi8st6gmyyXddE+VRmwWXmbW+czK8hhlAUuoybdYZA7T29ALk6QOxskZKbi3+
+BuJYdb583IAMQ4/I1UXbAPCFEpnDCcmhprIztRhuXofE6pg0KkR0XhqKce2Xypog
+RZyZfnsRP78B5OTE2EfFoNX9f5aHR79nHU7gh4nJRnCPb0xvrd0fRNy/aK/RVenk
+f2MHcntgA2HdZ8RUgrF7SoN1JpBQUsNIGlmQz2KCK4Nx78jrlHGjQu3TeRQ7FoLR
+H45q4UkZQdcKzSlLUgIVxu1l292muCV1+6g/SZl8rLJjfRNxV3RRFt0IYnm439n3
+y4BdZfATsd4Q1gjGoALvDK0i7pwG+40wxmKFDkj95QmskRMSX8q0DUg=
+-----END ENCRYPTED PRIVATE KEY-----`,
+        password: 'foo'
+      }
+    } as SnowflakeDatasourceConfiguration;
+    const connectionOptions = connectionOptionsFromDatasourceConfiguration(datasourceConfiguration);
+
+    expect(connectionOptions).toEqual({
+      account: 'account',
+      authenticator: 'SNOWFLAKE_JWT',
+      privateKey: `-----BEGIN PRIVATE KEY-----
+MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCXjMqyEL55ZO7b
+/20zZeOPOCeWT87Z2L+mt7fz0Jw9F9Qi7xjopxfadEf9+lAKzFA0WqpaflF0CkUh
+a8PVVW7yIx/fxgdgPmPdYeM8L0WO4ZuqqAU0/GXUsB9clucGvSPlb3LwN1SPeePy
+jp9QAVhuXuxxJYckcIDtDJqdotEH4CWRlkT7sfsVvLewHfLrZMrx2fT+vEPkG6iO
+xc2gGFY3qGSR4ol6gek2OrPul3LOonFP4J2LCq961V/oqEZVNmY8uJRIIVYGE9wJ
+FzYNcAbPmY45PNHbsRBxySkva4sqCHvdlTDyF4J0zrgGTUSn8zFy1GuUHPxU9jj1
+yA3H3KTvAgMBAAECggEACITdkAzo6CIIW02Mces0U4SwlP5/fjd07oY9TPNZFQYV
+4rDWEZ3y9VC8C0ba4QCDMa43DUc0hPVe1XYdWAypY+sdi5KyL3LSYktyXpu5qUJ9
+YaZ4RHY7sy0DVk0VR4eyUN9m8qHpmvcbHTNvOSwaMHodiG63nhDE7o5a0qmoXmzh
+BNstWaTxkAe0n9lSTDKACJWS2adGamlhhpqzp+EVVrVqODXslk8+OhMLmS2+FWuy
+mYzdqTdNbvMLkdQvdmoBXF8BkgfU/Kgwcn6Y/g1wZw5NEKGfzUYWhjvq+s3WZoeN
+5uuRTYII2QFwtj5bzKBySjDIXN9kJ3TD1KfbSIHUgQKBgQDOJx2Y93d4CLng7KBo
+vZrFjpF4CrgEk91vLmU649pyir+vK7CtM+XWHVXGBx69um6gagpY5AAH/AURnNT/
+7faoEDwJqR0s5SDVP8K5PwYZyxWVIMwgTzxWwbCsEb0P/Bbvqenp0zNAG8Vh5YI3
+eecOh6+BMiNMt/W2xmYvgYRy3wKBgQC8McEQ28Pg7udGiLS0UrPEr+NZtMWzqTQ4
+0lMS+RA4MHhZSCpM1TbcP3Eupyt8bvib/7lqT2rr/qQ2m+T5gY2WQrKbyFWrvD8c
+R94c3YhUqW9lS6z7AZ+LhjttXHGGRbZP/aSXYEILqBiBO9HbxhLh+Iz4/998jIne
+pYlXt6yf8QKBgEykdDeDgVIKBHkf3/8wxpK+D00Oxx1Ej+We3RnIzlUZSmxolNW/
+3qn82/+0c/RblHdlFRW5Jl3Rj1zd7r57jOEsr/VzfxpK0SsW+mD+klkSjKKVv+4f
+JzKl7fX63kxMD4bop8M7tukVqgtcVU4krwdS4KfqqP8DwYFDP4hX4ZMHAoGAd3dS
+KySHRQwDjvgLVolFiy9osLKb6kAYYZXKnLm0/SZvz6WLDLkxGUHA1K/UYCqF8Wm1
+x3Hg2y0MC4qNIYKHYgK3JUNYdyuKGKbarhJHkA77Ix+WEMVoBYdRxEux2V35rO/E
+A0BczM+JtshFoTEtHXvN6edsdME2aDtHY4K6t9ECgYBMvDlYOHcMu/NF07JmbQD0
+DQn5LpdqpVxSridY+uhvMQyByqJ7F9/kCg591jnIh0guaGBWIk+Xtqslgs1yuobM
+DJuQTfIrVUrtlBiKbK68eu8BWQmDkrry7lspDBBrAlL/Av23LtUHAhfaryHP91Hm
+oBHJaYtVCXd3VBWCVLcfnw==
+-----END PRIVATE KEY-----
+`
+    });
+  });
+
+  it('fails for key-pair when privateKey is not present', () => {
+    expect(() =>
+      connectionOptionsFromDatasourceConfiguration({
+        connectionType: 'key-pair',
+        authentication: {
+          custom: {
+            account: { value: 'account' }
+          }
+        }
+      })
+    ).toThrow('Missing required fields: privateKey');
+  });
+
   it('fails when authentication is not present', () => {
     expect(() => connectionOptionsFromDatasourceConfiguration({})).toThrow('authentication expected but not present');
   });
@@ -92,7 +206,7 @@ describe('connectionOptionsFromDatasourceConfiguration', () => {
       authentication: {}
     } as SnowflakeDatasourceConfiguration;
     expect(() => connectionOptionsFromDatasourceConfiguration(datasourceConfiguration)).toThrow(
-      'Missing required fields: username,password,account,databaseName'
+      'Missing required fields: account,username,password,databaseName'
     );
   });
 
@@ -102,7 +216,7 @@ describe('connectionOptionsFromDatasourceConfiguration', () => {
       authentication: {}
     } as SnowflakeDatasourceConfiguration;
     expect(() => connectionOptionsFromDatasourceConfiguration(datasourceConfiguration)).toThrow(
-      'Missing required fields: username,password,account,authenticatorUr'
+      'Missing required fields: account,username,password,authenticatorUrl'
     );
   });
 });
diff --git a/workers/javascript/packages/plugins/snowflake/src/util.ts b/workers/javascript/packages/plugins/snowflake/src/util.ts
@@ -1,6 +1,19 @@
+import * as crypto from 'crypto';
 import { SnowflakeDatasourceConfiguration, ErrorCode, IntegrationError } from '@superblocks/shared';
 import { ConnectionOptions } from 'snowflake-sdk';
 
+export function getPrivateKeyValue({ privateKey, password }: { privateKey: string; password: string }): string {
+  const privateKeyObject = crypto.createPrivateKey({
+    key: privateKey,
+    format: 'pem',
+    passphrase: password
+  });
+  return privateKeyObject.export({
+    format: 'pem',
+    type: 'pkcs8'
+  }) as string;
+}
+
 export function connectionOptionsFromDatasourceConfiguration(datasourceConfiguration: SnowflakeDatasourceConfiguration): ConnectionOptions {
   const auth = datasourceConfiguration.authentication;
   if (!auth) {
@@ -9,24 +22,43 @@ export function connectionOptionsFromDatasourceConfiguration(datasourceConfigura
 
   const missingFields: string[] = [];
 
-  // these fields are always required regardless of auth type
-  if (!auth.username) {
-    missingFields.push('username');
-  }
-  if (!auth.password) {
-    missingFields.push('password');
-  }
   if (!auth.custom?.account?.value) {
     missingFields.push('account');
   }
 
   switch (datasourceConfiguration.connectionType) {
+    case 'key-pair': {
+      // https://docs.snowflake.com/en/developer-guide/node-js/nodejs-driver-authenticate#label-nodejs-key-pair-authentication
+
+      if (!datasourceConfiguration?.keyPair?.privateKey) {
+        missingFields.push('privateKey');
+      }
+
+      handleMissingFields(missingFields);
+      let privateKey = datasourceConfiguration.keyPair.privateKey;
+      let password = datasourceConfiguration?.keyPair?.password;
+      if (password) {
+        privateKey = getPrivateKeyValue({ privateKey, password });
+      }
+      return {
+        account: auth.custom.account.value,
+        authenticator: 'SNOWFLAKE_JWT',
+        privateKey
+      };
+    }
     case 'okta': {
       // https://docs.snowflake.com/en/developer-guide/node-js/nodejs-driver-authenticate#using-native-sso-through-okta
 
+      if (!auth.username) {
+        missingFields.push('username');
+      }
+      if (!auth.password) {
+        missingFields.push('password');
+      }
       if (!datasourceConfiguration?.okta?.authenticatorUrl) {
         missingFields.push('authenticatorUrl');
       }
+
       handleMissingFields(missingFields);
       return {
         account: auth.custom.account.value,
@@ -36,6 +68,12 @@ export function connectionOptionsFromDatasourceConfiguration(datasourceConfigura
       };
     }
     default: {
+      if (!auth.username) {
+        missingFields.push('username');
+      }
+      if (!auth.password) {
+        missingFields.push('password');
+      }
       if (!auth.custom?.databaseName?.value) {
         missingFields.push('databaseName');
       }

diff --git a/workers/javascript/packages/shared/src/types/datasource/index.ts b/workers/javascript/packages/shared/src/types/datasource/index.ts
@@ -357,10 +357,16 @@ export type S3DatasourceConfiguration = AWSDatasourceConfiguration & { endpoint?
 
 export type SnowflakeDatasourceConfiguration = DBDatasourceConfiguration &
   DBConnection & {
-    connectionType?: 'fields' | 'okta';
+    // okta: https://docs.snowflake.com/en/developer-guide/node-js/nodejs-driver-authenticate#use-native-sso-through-okta
+    // key-pair: https://docs.snowflake.com/en/developer-guide/node-js/nodejs-driver-authenticate#label-nodejs-key-pair-authentication
+    connectionType?: 'fields' | 'okta' | 'key-pair';
     okta?: {
       authenticatorUrl?: string;
     };
+    keyPair?: {
+      privateKey: string;
+      password?: string;
+    };
     authentication?: {
       username?: string;
       password?: string;