kinmedai.go - a quick Google Natural Language API test in go

kinmedai.go

package main

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"io/ioutil"
	"log"
	"net/http"
)

type DBH struct {
	*sql.DB
}

type EntityRequest struct {
	EncodingType string                `json:"encodingType"`
	Document     EntityRequestDocument `json:"document"`
}

type EntityRequestDocument struct {
	TypeName string `json:"type"`
	Content  string `json:"content"`
	Language string `json:"language"`
}

type EntityResponse struct {
	Entities []DetectedEntity `json:"entities"`
	Language string           `json:"language"`
}

type DetectedEntity struct {
	Name       string          `json:"name"`
	EntityType string          `json:"type"`
	Salience   float64         `json:"salience"`
	Mentions   []EntityMention `json:"mentions"`
	Metadata   struct {
		WikipediaUrl string `json:wikipedia_url"`
	} `json:"metadata"`
}

type EntityMention struct {
	Text struct {
		Content     string `json:"content"`
		BeginOffset string `json:"beginOffset"`
	} `json:"text"`
}

const (
	ENTITIES_URL = "https://language.googleapis.com/v1beta1/documents:analyzeEntities"
)

var (
	dbh    *DBH
	dbPath = flag.String("d", "databasename.db", "location of sqlite3 db file")
)

func main() {
	flag.Parse()

	var accessToken string
	fmt.Scan(&accessToken)

	entityRequests := createEntityRequests()
	detectedEntities := make(map[string]int)

	for i := 0; i < len(entityRequests); i++ {
		var entityResponse EntityResponse
		body := postEntity(accessToken, entityRequests[i])
		json.Unmarshal(body, &entityResponse)

		for j := 0; j < len(entityResponse.Entities); j++ {
			entity := entityResponse.Entities[j]
			detectedEntities[entity.Name] = detectedEntities[entity.Name] + 1
		}
	}

	for k, v := range detectedEntities {
		fmt.Printf("%s: %d\n", k, v)
	}
}

func initDB() {
	handle, err := sql.Open("sqlite3", *dbPath)
	if err != nil {
		log.Fatal(err)
	}

	dbh = &DBH{handle}
}

func getDBH() *DBH {
	if dbh == nil {
		initDB()
	}
	return dbh
}

func createEntityRequests() []*EntityRequest {
	dbh := getDBH()
	rows, err := dbh.Query(`SELECT content FROM review WHERE rating < 3`)
	if err != nil {
		log.Fatal(err)
	}

	var entities []*EntityRequest

	for rows.Next() {
		var comment string
		err = rows.Scan(&comment)
		if err != nil {
			log.Fatal(err)
		}

		// Google Play lets users submit ratings with no comments (stars only ratings) so skip those
		if len(comment) == 0 {
			continue
		}

		entityRequest := &EntityRequest{
			EncodingType: "UTF8",
			Document: EntityRequestDocument{
				TypeName: "PLAIN_TEXT",
				Content:  comment,
				Language: "JA",
			},
		}
		entities = append(entities, entityRequest)
	}

	return entities
}

func postEntity(accessToken string, entityRequest *EntityRequest) []byte {
	jsonEntity, _ := json.Marshal(entityRequest)
	req, err := http.NewRequest("POST", ENTITIES_URL, bytes.NewBuffer(jsonEntity))
	if err != nil {
		log.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+accessToken)

	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		log.Fatal(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)

	if res.StatusCode != http.StatusOK || err != nil {
		log.Fatal(res.Status)
		log.Fatal(err)
	}

	return body
}