Repository: amp-labs/connectors
Pull request #1444 (2025-03-05 21:43:16)

[CON-1591]feat(Salesforce): GetRecordsWithIds
# Changes

* Extended existing SOQL builder to accept list of identifiers.
  * Added mock test to ensure query is built as expected.
* Added test script to query Accounts (see below).

# Tests
`salesforce/record/main.go`:
![image](https://github.com/user-attachments/assets/77ce3eb2-3ed5-4502-b823-ddeaed9a5307)
![image](https://github.com/user-attachments/assets/9ec66f77-d491-404e-a330-f216420911c5)

# Notes
Associations argument at `GetRecordsWithIds` is preserved to **match method signature** of other deep connectors. However, it is not used by the Salesforce implementation.

Cobalt0s: * **#1444** <a href="https://app.graphite.dev/github/pr/amp-labs/connectors/1444?utm_source=stack-comment-icon" target="_blank"><img src="https://static.graphite.dev/graphite-32x32-black.png" alt="Graphite" width="10px" height="10px"/></a> 👈 <a href="https://app.graphite.dev/github/pr/amp-labs/connectors/1444?utm_source=stack-comment-view-in-graphite" target="_blank">(View in Graphite)</a>
* `main`


<h2></h2>

This stack of pull requests is managed by <a href="https://graphite.dev?utm-source=stack-comment"><b>Graphite</b></a>. Learn more about <a href="https://stacking.dev/?utm_source=stack-comment">stacking</a>.
<!-- Current dependencies on/for this PR: -->

Patch:
diff --git a/providers/salesforce/record.go b/providers/salesforce/record.go
new file mode 100644
index 00000000..389e91eb
--- /dev/null
+++ b/providers/salesforce/record.go
@@ -0,0 +1,76 @@
+package salesforce
+
+import (
+	"context"
+
+	"github.com/amp-labs/connectors/common"
+	"github.com/amp-labs/connectors/common/urlbuilder"
+	"github.com/amp-labs/connectors/internal/datautils"
+)
+
+// GetRecordsWithIds returns records matching identifiers.
+func (c *Connector) GetRecordsWithIds( // nolint:revive
+	ctx context.Context,
+	objectName string,
+	ids []string,
+	fields []string,
+	associations []string, // no-op, preserved to match other deep connectors
+) ([]common.ReadResultRow, error) {
+	// Sanitize method arguments.
+	config := recordsByIDsParams{
+		ReadParams: common.ReadParams{
+			ObjectName:        objectName,
+			Fields:            datautils.NewSetFromList(fields),
+			AssociatedObjects: associations,
+		},
+		RecordIdentifiers: datautils.NewSetFromList(ids),
+	}
+
+	if err := config.ValidateParams(true); err != nil {
+		return nil, err
+	}
+
+	url, err := c.buildReadByIdentifierURL(config)
+	if err != nil {
+		return nil, err
+	}
+
+	rsp, err := c.Client.Get(ctx, url.String())
+	if err != nil {
+		return nil, err
+	}
+
+	readResult, err := common.ParseResult(
+		rsp,
+		getRecords,
+		getNextRecordsURL,
+		common.GetMarshaledData,
+		config.Fields,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return readResult.Data, nil
+}
+
+type recordsByIDsParams struct {
+	common.ReadParams
+	RecordIdentifiers datautils.Set[string]
+}
+
+func (c *Connector) buildReadByIdentifierURL(config recordsByIDsParams) (*urlbuilder.URL, error) {
+	// Requesting record identifiers using SOQL query.
+	url, err := c.getRestApiURL("query")
+	if err != nil {
+		return nil, err
+	}
+
+	query := makeSOQL(config.ReadParams).
+		WithIDs(config.RecordIdentifiers.List()).
+		String()
+
+	url.WithQueryParam("q", query)
+
+	return url, nil
+}
diff --git a/providers/salesforce/soql.go b/providers/salesforce/soql.go
index 78389aa9..c1d27ef0 100644
--- a/providers/salesforce/soql.go
+++ b/providers/salesforce/soql.go
@@ -5,6 +5,13 @@ import (
 	"strings"
 )
 
+// nolint:lll
+const (
+	// See `Limiting Result Rows` section on this webpage:
+	// https://developer.salesforce.com/docs/atlas.en-us.soql_sosl.meta/soql_sosl/sforce_api_calls_soql_select_fields.htm
+	identifiersLimitStr = "200"
+)
+
 // soqlBuilder builder of Salesforce Object Query Language.
 // It constructs query dynamically.
 type soqlBuilder struct {
@@ -20,7 +27,7 @@ func (s *soqlBuilder) SelectFields(fields []string) *soqlBuilder {
 			s.fields = "FIELDS(ALL)"
 			// if all fields are to be returned then we must limit to avoid error.
 			// Error example: `The SOQL FIELDS function must have a LIMIT of at most 200`
-			s.limit = "200"
+			s.limit = identifiersLimitStr
 
 			return s
 		}
@@ -47,6 +54,19 @@ func (s *soqlBuilder) Where(condition string) *soqlBuilder {
 	return s
 }
 
+func (s *soqlBuilder) WithIDs(identifiers []string) *soqlBuilder {
+	// Decorate each id with quotes.
+	for index, id := range identifiers {
+		identifiers[index] = fmt.Sprintf("'%v'", id)
+	}
+
+	identifiersList := strings.Join(identifiers, ",")
+
+	// nolint:lll
+	// https://developer.salesforce.com/docs/atlas.en-us.soql_sosl.meta/soql_sosl/sforce_api_calls_soql_select_fields.htm
+	return s.Where(fmt.Sprintf("Id IN (%v)", identifiersList))
+}
+
 func (s *soqlBuilder) String() string {
 	query := fmt.Sprintf("SELECT %s FROM %s", s.fields, s.from)
 


Test patch:
diff --git a/providers/salesforce/soql_test.go b/providers/salesforce/soql_test.go
new file mode 100644
index 00000000..50ade0ad
--- /dev/null
+++ b/providers/salesforce/soql_test.go
@@ -0,0 +1,44 @@
+package salesforce
+
+import (
+	"testing"
+
+	"github.com/amp-labs/connectors/common"
+	"github.com/amp-labs/connectors/internal/datautils"
+	"gotest.tools/v3/assert"
+)
+
+func TestSoqlBuilderWithIDs(t *testing.T) {
+	t.Parallel()
+
+	soql := makeSOQL(common.ReadParams{
+		ObjectName: "Account",
+		// Note: fields doesn't preserve order of elements.
+		// To simplify test only one element is included.
+		Fields: datautils.NewSet("shippingstreet"),
+	})
+
+	{
+		// SOQL builder must produce query matching documentation.
+		// nolint:lll
+		// https://developer.salesforce.com/docs/atlas.en-us.soql_sosl.meta/soql_sosl/sforce_api_calls_soql_select_fields.htm
+		soql.WithIDs(
+			[]string{
+				"001ak00000OQ4RxAAL",
+				"001ak00000OQ4RyAAL",
+				"001ak00000OQ4TZAA1",
+				"001ak00000OQ4TbAAL",
+				"001ak00000OQ4VBAA1",
+				"001ak00000OQ4VCAA1",
+			})
+
+		output := soql.String()
+		assert.Equal(t, output, "SELECT shippingstreet FROM Account WHERE Id IN ("+
+			"'001ak00000OQ4RxAAL',"+
+			"'001ak00000OQ4RyAAL',"+
+			"'001ak00000OQ4TZAA1',"+
+			"'001ak00000OQ4TbAAL',"+
+			"'001ak00000OQ4VBAA1',"+
+			"'001ak00000OQ4VCAA1')", "mismatching SOQL query string")
+	}
+}
diff --git a/test/salesforce/record/main.go b/test/salesforce/record/main.go
new file mode 100644
index 00000000..1d1bf514
--- /dev/null
+++ b/test/salesforce/record/main.go
@@ -0,0 +1,39 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+
+	connTest "github.com/amp-labs/connectors/test/salesforce"
+	"github.com/amp-labs/connectors/test/utils"
+)
+
+func main() {
+	// Handle Ctrl-C gracefully.
+	ctx, done := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer done()
+
+	// Set up slog logging.
+	utils.SetupLogging()
+
+	conn := connTest.GetSalesforceConnector(ctx)
+
+	res, err := conn.GetRecordsWithIds(ctx,
+		"Account",
+		[]string{
+			"001ak00000OQ4RxAAL",
+			"001ak00000OQ4RyAAL",
+			"001ak00000OQ4TZAA1",
+		},
+		[]string{"id", "name", "shippingstreet"},
+		nil)
+	if err != nil {
+		utils.Fail("error reading from connector", "error", err)
+	}
+
+	fmt.Println("Reading..")
+	utils.DumpJSON(res, os.Stdout)
+}