cxp-26 add user provisioning

.github/workflows/capabilities_and_config.yaml

@@ -27,15 +27,18 @@ jobs:
       - name: Run and save config output
         run: ./connector config > config_schema.json
 
+      - name: Setup private key
+        run: |
+          echo "${{ secrets.BATON_PRIVATE_KEY }}" | base64 -d > /tmp/snowflake_key.p8
+          chmod 600 /tmp/snowflake_key.p8
+
       - name: Run and save capabilities output
         env:
           BATON_ACCOUNT_IDENTIFIER: example
           BATON_USER_IDENTIFIER: example
           BATON_ACCOUNT_URL: https://example.snowflakecomputing.com
-          BATON_PRIVATE_KEY_PATH: ${{ runner.temp }}/baton_private_key.pem
-        run: |
-          openssl genrsa -out "$BATON_PRIVATE_KEY_PATH" 2048
-          ./connector --sync-secrets capabilities > baton_capabilities.json
+          BATON_PRIVATE_KEY_PATH: /tmp/snowflake_key.p8
+        run: ./connector --sync-secrets capabilities > baton_capabilities.json
 
       - name: Commit changes
         uses: EndBug/add-and-commit@v9

.github/workflows/ci.yaml

@@ -43,19 +43,13 @@ jobs:
 
   test:
     runs-on: ubuntu-latest
-    # Define any services needed for the test suite (or delete this section)
-    # services:
-    #   postgres:
-    #     image: postgres:16
-    #     ports:
-    #       - "5432:5432"
-    #     env:
-    #       POSTGRES_PASSWORD: secretpassword
+    if: github.event_name == 'pull_request' || github.ref == 'refs/heads/main'
     env:
       BATON_LOG_LEVEL: debug
-      # Add any environment variables needed to run baton-snowflake
-      # BATON_BASE_URL: 'http://localhost:8080'
-      # BATON_ACCESS_TOKEN: 'secret_token'
+      # Snowflake connection configuration
+      BATON_ACCOUNT_IDENTIFIER: ${{ secrets.BATON_ACCOUNT_IDENTIFIER }}
+      BATON_USER_IDENTIFIER: ${{ secrets.BATON_USER_IDENTIFIER }}
+      BATON_ACCOUNT_URL: ${{ secrets.BATON_ACCOUNT_URL }}
       # The following parameters are passed to grant/revoke commands
       # Change these to the correct IDs for your test data
       CONNECTOR_GRANT: 'grant:entitlement:group:1234:member:user:9876'
@@ -69,37 +63,29 @@ jobs:
         uses: actions/setup-go@v5
         with:
           go-version-file: 'go.mod'
-      # Install any dependencies here (or delete this)
-      # - name: Install postgres client
-      #   run: sudo apt install postgresql-client
-      # Run any fixture setup here (or delete this)
-      # - name: Import sql into postgres
-      #   run: psql -h localhost --user postgres -f environment.sql
-      #   env:
-      #     PGPASSWORD: secretpassword
+
       - name: Build baton-snowflake
-        run: go build ./cmd/baton-snowflake
+        run: go build -o baton-snowflake ./cmd/baton-snowflake
       # - name: Run baton-snowflake
       #   run: ./baton-snowflake
 
       - name: Install baton
         run: ./scripts/get-baton.sh && mv baton /usr/local/bin
 
-      # - name: Check for grant before revoking
-      #   run:
-      #     baton grants --entitlement="${{ env.CONNECTOR_ENTITLEMENT }}" --output-format=json | jq --exit-status ".grants[].principal.id.resource == \"${{ env.CONNECTOR_PRINCIPAL }}\""
-
-      # - name: Revoke grants
-      #   run: ./baton-snowflake --revoke-grant="${{ env.CONNECTOR_GRANT }}"
-
-      # - name: Check grant was revoked
-      #   run: ./baton-snowflake && baton grants --entitlement="${{ env.CONNECTOR_ENTITLEMENT }}" --output-format=json | jq --exit-status "if .grants then .grants[]?.principal.id.resource != \"${{ env.CONNECTOR_PRINCIPAL }}\" else . end"
+      - name: Setup private key
+        run: |
+          echo "${{ secrets.BATON_PRIVATE_KEY }}" | base64 -d > /tmp/snowflake_key.p8
+          chmod 600 /tmp/snowflake_key.p8
 
-      # - name: Grant entitlement
-      #   # Change the grant arguments to the correct IDs for your test data
-      #   run: ./baton-snowflake --grant-entitlement="${{ env.CONNECTOR_ENTITLEMENT }}" --grant-principal="${{ env.CONNECTOR_PRINCIPAL }}" --grant-principal-type="${{ env.CONNECTOR_PRINCIPAL_TYPE }}"
-
-      # - name: Check grant was re-granted
-      #   run:
-      #     baton grants --entitlement="${{ env.CONNECTOR_ENTITLEMENT }}" --output-format=json | jq --exit-status ".grants[].principal.id.resource == \"${{ env.CONNECTOR_PRINCIPAL }}\""
+      - name: Test Account Provisioning
+        uses: ConductorOne/github-workflows/actions/account-provisioning@v4
+        with:
+          connector: './baton-snowflake'
+          account-email: 'testProvisioningUser@example.com'
+          account-login: 'testProvisioningUser'
+          account-profile: '{"first_name": "Test", "last_name": "User", "name": "testProvisioningUser", "email": "testProvisioningUser@example.com"}'
+          account-type: 'user'
+          search-method: 'email'
+        env:
+          BATON_PRIVATE_KEY_PATH: /tmp/snowflake_key.p8
 

pkg/connector/account_roles.go

@@ -46,7 +46,7 @@ func (o *accountRoleBuilder) List(ctx context.Context, parentResourceID *v2.Reso
 		return nil, nil, wrapError(err, "failed to get next page offset")
 	}
 
-	accountRoles, _, err := o.client.ListAccountRoles(ctx, cursor, resourcePageSize)
+	accountRoles, err := o.client.ListAccountRoles(ctx, cursor, resourcePageSize)
 	if err != nil {
 		return nil, nil, wrapError(err, "failed to list account roles")
 	}
@@ -95,7 +95,7 @@ func (o *accountRoleBuilder) Entitlements(_ context.Context, resource *v2.Resour
 }
 
 func (o *accountRoleBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) {
-	accountRoleGrantees, _, err := o.client.ListAccountRoleGrantees(ctx, resource.DisplayName)
+	accountRoleGrantees, err := o.client.ListAccountRoleGrantees(ctx, resource.DisplayName)
 	if err != nil {
 		return nil, nil, wrapError(err, "failed to list account role grantees")
 	}
@@ -138,7 +138,7 @@ func (o *accountRoleBuilder) Grant(ctx context.Context, principal *v2.Resource,
 		return nil, err
 	}
 
-	_, err := o.client.GrantAccountRole(ctx, entitlement.Resource.Id.Resource, principal.Id.Resource)
+	err := o.client.GrantAccountRole(ctx, entitlement.Resource.Id.Resource, principal.Id.Resource)
 	if err != nil {
 		err = wrapError(err, "failed to grant account role")
 
@@ -167,7 +167,7 @@ func (o *accountRoleBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annot
 		return nil, err
 	}
 
-	_, err := o.client.RevokeAccountRole(ctx, grant.Entitlement.Resource.Id.Resource, grant.Principal.Id.Resource)
+	err := o.client.RevokeAccountRole(ctx, grant.Entitlement.Resource.Id.Resource, grant.Principal.Id.Resource)
 	if err != nil {
 		err = wrapError(err, "failed to revoke account role")
 

pkg/connector/connector.go

@@ -49,13 +49,136 @@ func (d *Connector) Metadata(ctx context.Context) (*v2.ConnectorMetadata, error)
 	return &v2.ConnectorMetadata{
 		DisplayName: "Baton Snowflake",
 		Description: "Connector syncing users, databases, tables, and account roles from Snowflake.",
+		AccountCreationSchema: &v2.ConnectorAccountCreationSchema{
+			FieldMap: map[string]*v2.ConnectorAccountCreationSchema_Field{
+				"name": {
+					DisplayName: "User Name",
+					Required:    true,
+					Description: "The name of the user (required - case-sensitive)",
+					Placeholder: "username",
+					Order:       0,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"login": {
+					DisplayName: "Login Name",
+					Required:    false,
+					Description: "The login name for the user (defaults to email if not provided)",
+					Placeholder: "user@example.com",
+					Order:       1,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"display_name": {
+					DisplayName: "Display Name",
+					Required:    false,
+					Description: "The display name for the user",
+					Placeholder: "John Doe",
+					Order:       2,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"first_name": {
+					DisplayName: "First Name",
+					Required:    false,
+					Description: "The first name of the user",
+					Placeholder: "John",
+					Order:       3,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"last_name": {
+					DisplayName: "Last Name",
+					Required:    false,
+					Description: "The last name of the user",
+					Placeholder: "Doe",
+					Order:       4,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"email": {
+					DisplayName: "Email",
+					Required:    false,
+					Description: "The email address for the user",
+					Placeholder: "user@example.com",
+					Order:       5,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"comment": {
+					DisplayName: "Comment",
+					Required:    false,
+					Description: "A comment or description for the user",
+					Placeholder: "User description",
+					Order:       6,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"disabled": {
+					DisplayName: "Disabled",
+					Required:    false,
+					Description: "Whether the user account should be disabled",
+					Order:       7,
+					Field: &v2.ConnectorAccountCreationSchema_Field_BoolField{
+						BoolField: &v2.ConnectorAccountCreationSchema_BoolField{},
+					},
+				},
+				"default_warehouse": {
+					DisplayName: "Default Warehouse",
+					Required:    false,
+					Description: "The default warehouse to use when this user starts a session",
+					Placeholder: "COMPUTE_WH",
+					Order:       8,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"default_namespace": {
+					DisplayName: "Default Namespace",
+					Required:    false,
+					Description: "The default namespace to use when this user starts a session",
+					Placeholder: "DATABASE.SCHEMA",
+					Order:       9,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"default_role": {
+					DisplayName: "Default Role",
+					Required:    false,
+					Description: "The default role to use when this user starts a session",
+					Placeholder: "PUBLIC",
+					Order:       10,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+				"default_secondary_roles": {
+					DisplayName: "Default Secondary Roles",
+					Required:    false,
+					Description: "The default secondary roles of this user to use when starting a session. Valid values: ALL or NONE. Default is ALL.",
+					Placeholder: "ALL",
+					Order:       11,
+					Field: &v2.ConnectorAccountCreationSchema_Field_StringField{
+						StringField: &v2.ConnectorAccountCreationSchema_StringField{},
+					},
+				},
+			},
+		},
 	}, nil
 }
 
 // Validate is called to ensure that the connector is properly configured. It should exercise any API credentials
 // to be sure that they are valid.
 func (d *Connector) Validate(ctx context.Context) (annotations.Annotations, error) {
-	users, _, err := d.Client.ListUsers(ctx, "", 1)
+	users, err := d.Client.ListUsers(ctx, "", 1)
 	if err != nil {
 		return nil, err
 	}

pkg/connector/databases.go

@@ -61,7 +61,7 @@ func (o *databaseBuilder) List(ctx context.Context, parentResourceID *v2.Resourc
 		return nil, nil, wrapError(err, "failed to get next page offset")
 	}
 
-	databases, _, err := o.client.ListDatabases(ctx, cursor, resourcePageSize)
+	databases, err := o.client.ListDatabases(ctx, cursor, resourcePageSize)
 	if err != nil {
 		return nil, nil, wrapError(err, "failed to list databases")
 	}
@@ -111,9 +111,9 @@ func (o *databaseBuilder) Grants(ctx context.Context, resource *v2.Resource, _ r
 		return nil, nil, nil
 	}
 
-	owner, ownerResp, err := o.client.GetAccountRole(ctx, database.Owner)
+	owner, ownerStatusCode, err := o.client.GetAccountRole(ctx, database.Owner)
 	if err != nil {
-		if snowflake.IsUnprocessableEntity(ownerResp, err) {
+		if snowflake.IsUnprocessableEntity(ownerStatusCode, err) {
 			wrappedErr := fmt.Errorf("baton-snowflake: insufficient privileges for database owner role %q (database %q): %w", database.Owner, resource.Id.Resource, err)
 			return nil, nil, status.Error(codes.PermissionDenied, wrappedErr.Error())
 		}

pkg/connector/helpers.go

@@ -1,9 +1,22 @@
 package connector
 
-import "fmt"
+import (
+	"fmt"
+	"strings"
+)
 
 func wrapError(err error, message string) error {
 	return fmt.Errorf("snowflake-connector: %s: %w", message, err)
 }
 
 const resourcePageSize = 50
+
+// quoteSnowflakeIdentifier properly escapes and quotes a Snowflake identifier.
+// In Snowflake, double quotes inside identifiers must be escaped by doubling them.
+// Example: o"donnel becomes "o""donnel".
+func quoteSnowflakeIdentifier(identifier string) string {
+	// Escape double quotes by doubling them
+	escaped := strings.ReplaceAll(identifier, `"`, `""`)
+	// Wrap in double quotes
+	return fmt.Sprintf(`"%s"`, escaped)
+}

pkg/connector/tables.go

@@ -69,8 +69,8 @@ func (o *tableBuilder) isDBSharedOrSystem(ctx context.Context, resource *v2.Reso
 			return val == "true" || val == "1", nil
 		}
 	}
-	db, resp, err := o.client.GetDatabase(ctx, databaseName)
-	if snowflake.IsUnprocessableEntity(resp, err) {
+	db, statusCode, err := o.client.GetDatabase(ctx, databaseName)
+	if snowflake.IsUnprocessableEntity(statusCode, err) {
 		return true, nil
 	}
 	if err != nil {
@@ -148,7 +148,7 @@ func (o *tableBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId
 	}
 
 	const accountPageSize = 200
-	tables, nextCursor, _, err := o.client.ListTablesInAccount(ctx, cursor, accountPageSize)
+	tables, nextCursor, err := o.client.ListTablesInAccount(ctx, cursor, accountPageSize)
 	if err != nil {
 		return nil, nil, wrapError(err, "failed to list tables in account")
 	}
@@ -280,9 +280,9 @@ func (o *tableBuilder) Grants(ctx context.Context, resource *v2.Resource, opts r
 
 		switch tg.GrantedTo {
 		case grantedToRole:
-			role, resp, err := o.client.GetAccountRole(ctx, tg.GranteeName)
+			role, statusCode, err := o.client.GetAccountRole(ctx, tg.GranteeName)
 			if err != nil {
-				if snowflake.IsUnprocessableEntity(resp, err) {
+				if snowflake.IsUnprocessableEntity(statusCode, err) {
 					principalId, idErr := rs.NewResourceID(accountRoleResourceType, tg.GranteeName)
 					if idErr != nil {
 						continue
@@ -335,14 +335,14 @@ func (o *tableBuilder) Grants(ctx context.Context, resource *v2.Resource, opts r
 	}
 
 	if ownerPrincipalID == nil {
-		table, _, err := o.client.GetTable(ctx, databaseName, schemaName, tableName)
+		table, err := o.client.GetTable(ctx, databaseName, schemaName, tableName)
 		if err != nil {
 			return nil, nil, wrapError(err, "failed to get table for owner fallback")
 		}
 		if table != nil && table.Owner != "" && table.Owner != "SNOWFLAKE" {
-			owner, ownerResp, err := o.client.GetAccountRole(ctx, table.Owner)
+			owner, ownerStatusCode, err := o.client.GetAccountRole(ctx, table.Owner)
 			switch {
-			case snowflake.IsUnprocessableEntity(ownerResp, err):
+			case snowflake.IsUnprocessableEntity(ownerStatusCode, err):
 				// system role, skip
 			case err != nil:
 				return nil, nil, wrapError(err, fmt.Sprintf("failed to get account role for table owner %q", table.Owner))