Skip to content

Commit

Permalink
Add support for batching requests over a certain size
Browse files Browse the repository at this point in the history
  • Loading branch information
@marcus-crane
marcus-crane committed Jan 4, 2023
1 parent 6370bc0 commit ce28dbf
Show file tree
Hide file tree
Showing 5 changed files with 70 additions and 39 deletions.
64 changes: 41 additions & 23 deletions backend/readwise.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,14 @@ import (

"github.com/go-resty/resty/v2"

"github.com/sirupsen/logrus"
log "github.com/sirupsen/logrus"
)

const (
HIGHLIGHT_REQUEST_BATCH_MAX = 2000
)

type Response struct {
Highlights []Highlight `json:"highlights"`
}
Expand Down Expand Up @@ -69,22 +74,27 @@ func (r *Readwise) CheckTokenValidity(token string) error {
return nil
}

func (r *Readwise) SendBookmarks(payload Response, token string) (int, error) {
client := resty.New()
resp, err := client.R().
SetHeader("Authorization", fmt.Sprintf("Token %s", token)).
SetHeader("User-Agent", UserAgent).
SetBody(payload).
Post(HighlightsEndpoint)
if err != nil {
return 0, fmt.Errorf("failed to send request to Readwise: code %d", resp.StatusCode())
}
if resp.StatusCode() != 200 {
log.WithFields(log.Fields{"status_code": resp.StatusCode(), "response": string(resp.Body())}).Error("Received a non-200 response from Readwise")
return 0, fmt.Errorf("received a non-200 status code from Readwise: code %d", resp.StatusCode())
func (r *Readwise) SendBookmarks(payloads []Response, token string) (int, error) {
// TODO: This is dumb, we count stuff that this function doesn't need to know about + we already know the size from earlier
submittedHighlights := 0
for _, payload := range payloads {
client := resty.New()
resp, err := client.R().
SetHeader("Authorization", fmt.Sprintf("Token %s", token)).
SetHeader("User-Agent", UserAgent).
SetBody(payload).
Post(HighlightsEndpoint)
if err != nil {
return 0, fmt.Errorf("failed to send request to Readwise: code %d", resp.StatusCode())
}
if resp.StatusCode() != 200 {
log.WithFields(log.Fields{"status_code": resp.StatusCode(), "response": string(resp.Body())}).Error("Received a non-200 response from Readwise")
return 0, fmt.Errorf("received a non-200 status code from Readwise: code %d", resp.StatusCode())
}
submittedHighlights += len(payload.Highlights)
}
log.WithField("highlight_count", len(payload.Highlights)).Info("Successfully sent bookmarks to Readwise")
return len(payload.Highlights), nil
log.WithField("batch_count", len(payloads)).Info("Successfully sent bookmarks to Readwise")
return submittedHighlights, nil
}

func (r *Readwise) RetrieveUploadedBooks(token string) (BookListResponse, error) {
Expand Down Expand Up @@ -157,9 +167,16 @@ func (r *Readwise) UploadCover(encodedCover string, bookId int, token string) er
return nil
}

func BuildPayload(bookmarks []Bookmark, contentIndex map[string]Content) (Response, error) {
var payload Response
for _, entry := range bookmarks {
func BuildPayload(bookmarks []Bookmark, contentIndex map[string]Content) ([]Response, error) {
var payloads []Response
var currentBatch Response
for count, entry := range bookmarks {
// If max payload size is reached, start building another batch which will be sent separately
if count > 0 && (count%HIGHLIGHT_REQUEST_BATCH_MAX == 0) {
fmt.Println(count / HIGHLIGHT_REQUEST_BATCH_MAX)
payloads = append(payloads, currentBatch)
currentBatch = Response{}
}
source := contentIndex[entry.VolumeID]
log.WithField("title", source.Title).Debug("Parsing highlight")
var createdAt string
Expand All @@ -172,15 +189,15 @@ func BuildPayload(bookmarks []Bookmark, contentIndex map[string]Content) (Respon
t, err := time.Parse("2006-01-02T15:04:05Z", entry.DateModified)
if err != nil {
log.WithError(err).WithFields(log.Fields{"title": source.Title, "volume_id": entry.VolumeID, "date_modified": entry.DateModified}).Error("Failed to parse a valid timestamp from date modified field")
return Response{}, err
return []Response{}, err
}
createdAt = t.Format("2006-01-02T15:04:05-07:00")
}
} else {
t, err := time.Parse("2006-01-02T15:04:05.000", entry.DateCreated)
if err != nil {
log.WithError(err).WithFields(log.Fields{"title": source.Title, "volume_id": entry.VolumeID, "date_modified": entry.DateModified}).Error("Failed to parse a valid timestamp from date created field")
return Response{}, err
return []Response{}, err
}
createdAt = t.Format("2006-01-02T15:04:05-07:00")
}
Expand Down Expand Up @@ -226,12 +243,13 @@ func BuildPayload(bookmarks []Bookmark, contentIndex map[string]Content) (Respon
Note: entry.Annotation,
HighlightedAt: createdAt,
}
payload.Highlights = append(payload.Highlights, highlight)
currentBatch.Highlights = append(currentBatch.Highlights, highlight)
}
log.WithFields(log.Fields{"title": source.Title, "volume_id": entry.VolumeID, "chunks": len(highlightChunks)}).Debug("Successfully compiled highlights for book")
}
log.WithField("highlight_count", len(payload.Highlights)).Info("Successfully parsed highlights")
return payload, nil
payloads = append(payloads, currentBatch)
log.WithFields(logrus.Fields{"highlight_count": len(currentBatch.Highlights), "batch_count": len(payloads)}).Info("Successfully parsed highlights")
return payloads, nil
}

func NormaliseText(s string) string {
Expand Down
39 changes: 26 additions & 13 deletions backend/readwise_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ func TestBuildPayload_NoBookmarks(t *testing.T) {
var contentIndex map[string]Content
var bookmarks []Bookmark
var actual, _ = BuildPayload(bookmarks, contentIndex)
assert.Equal(t, expected, actual)
assert.Equal(t, expected, actual[0])
}

func TestBuildPayload_BookmarksPresent(t *testing.T) {
Expand All @@ -33,7 +33,7 @@ func TestBuildPayload_BookmarksPresent(t *testing.T) {
{VolumeID: "mnt://kobo/blah/Good Book - An Author.epub", Text: "Hello World", DateCreated: "2006-01-02T15:04:05.000", Annotation: "Making a note here"},
}
var actual, _ = BuildPayload(bookmarks, contentIndex)
assert.Equal(t, expected, actual)
assert.Equal(t, expected, actual[0])
}

func TestBuildPayload_HandleAnnotationOnly(t *testing.T) {
Expand All @@ -53,7 +53,7 @@ func TestBuildPayload_HandleAnnotationOnly(t *testing.T) {
{VolumeID: "mnt://kobo/blah/Good Book - An Author.epub", DateCreated: "2006-01-02T15:04:05.000", Annotation: "Making a note here"},
}
var actual, _ = BuildPayload(bookmarks, contentIndex)
assert.Equal(t, expected, actual)
assert.Equal(t, expected, actual[0])
}

func TestBuildPayload_TitleFallback(t *testing.T) {
Expand All @@ -73,7 +73,7 @@ func TestBuildPayload_TitleFallback(t *testing.T) {
{VolumeID: "mnt://kobo/blah/Good Book - An Author.epub", Text: "Hello World", DateCreated: "2006-01-02T15:04:05.000", Annotation: "Making a note here"},
}
var actual, _ = BuildPayload(bookmarks, contentIndex)
assert.Equal(t, expected, actual)
assert.Equal(t, expected, actual[0])
}

func TestBuildPayload_TitleFallbackFailure(t *testing.T) {
Expand All @@ -91,7 +91,7 @@ func TestBuildPayload_TitleFallbackFailure(t *testing.T) {
{VolumeID: "\t", Text: "Hello World", DateCreated: "2006-01-02T15:04:05.000", Annotation: "Making a note here"},
}
var actual, _ = BuildPayload(bookmarks, contentIndex)
assert.Equal(t, expected, actual)
assert.Equal(t, expected, actual[0])
}

func TestBuildPayload_NoHighlightDateCreated(t *testing.T) {
Expand All @@ -109,7 +109,7 @@ func TestBuildPayload_NoHighlightDateCreated(t *testing.T) {
{VolumeID: "\t", Text: "Hello World", DateCreated: "", Annotation: "Making a note here", DateModified: "2006-01-02T15:04:05Z"},
}
var actual, _ = BuildPayload(bookmarks, contentIndex)
assert.Equal(t, expected, actual)
assert.Equal(t, expected, actual[0])
}

func TestBuildPayload_NoHighlightDateAtAll(t *testing.T) {
Expand All @@ -118,10 +118,10 @@ func TestBuildPayload_NoHighlightDateAtAll(t *testing.T) {
{VolumeID: "abc123", Text: "Hello World", Annotation: "Making a note here"},
}
var actual, _ = BuildPayload(bookmarks, contentIndex)
assert.Equal(t, actual.Highlights[0].SourceURL, bookmarks[0].VolumeID)
assert.Equal(t, actual.Highlights[0].Text, bookmarks[0].Text)
assert.Equal(t, actual.Highlights[0].Note, bookmarks[0].Annotation)
assert.NotEmpty(t, actual.Highlights[0].HighlightedAt)
assert.Equal(t, actual[0].Highlights[0].SourceURL, bookmarks[0].VolumeID)
assert.Equal(t, actual[0].Highlights[0].Text, bookmarks[0].Text)
assert.Equal(t, actual[0].Highlights[0].Note, bookmarks[0].Annotation)
assert.NotEmpty(t, actual[0].Highlights[0].HighlightedAt)
}

func TestBuildPayload_SkipMalformedBookmarks(t *testing.T) {
Expand All @@ -131,7 +131,7 @@ func TestBuildPayload_SkipMalformedBookmarks(t *testing.T) {
{VolumeID: "mnt://kobo/blah/Good Book - An Author.epub", DateCreated: "2006-01-02T15:04:05.000"},
}
var actual, _ = BuildPayload(bookmarks, contentIndex)
assert.Equal(t, expected, actual)
assert.Equal(t, expected, actual[0])
}

func TestBuildPayload_LongHighlightChunks(t *testing.T) {
Expand Down Expand Up @@ -176,6 +176,19 @@ func TestBuildPayload_LongHighlightChunks(t *testing.T) {
bookmarks := []Bookmark{
{VolumeID: "mnt://kobo/blah/Good Book - An Author.epub", Text: text, DateCreated: "2006-01-02T15:04:05.000", Annotation: "Making a note here"},
}
var actual, _ = BuildPayload(bookmarks, contentIndex)
assert.Equal(t, expected, actual)
actual, err := BuildPayload(bookmarks, contentIndex)
assert.NoError(t, err)
assert.Equal(t, expected, actual[0])
}

func TestExcessiveHighlightAmounts(t *testing.T) {
expected := 3
contentIndex := map[string]Content{"mnt://kobo/blah/Good Book - An Author.epub": {ContentID: "mnt://kobo/blah/Good Book - An Author.epub", Title: "A Book", Attribution: "Computer"}}
bookmarks := []Bookmark{}
for i := 0; i < 4002; i++ {
bookmarks = append(bookmarks, Bookmark{VolumeID: "mnt://kobo/blah/Good Book - An Author.epub", Text: fmt.Sprintf("hi%d", i), DateCreated: "2006-01-02T15:04:05.000", Annotation: "Making a note here"})
}
actual, err := BuildPayload(bookmarks, contentIndex)
assert.NoError(t, err)
assert.Equal(t, expected, len(actual))
}
Loading

0 comments on commit ce28dbf

Please sign in to comment.