client.go 3.65 KB
// This library is designed to replicate some of the basic functionality
// found in the Java version of the Twitter Tools for the TREC Mircoblog
// track.
//
// Example Usage
//
//	// Create a new client using the settings in conf.json
// 	client, err := twittertools.NewClientFromConf("conf.json")
// 	if err != nil {
// 		log.Fatal(err)
// 	}
// 	defer client.Close()
//
// 	// query the server for "glasgow" and retreive the top 50 results
// 	// -1 signifies that there is no maximum tweet ID
// 	results, err := client.Search("Glasgow", 50, -1)
// 	if err != nil {
// 		fmt.Println(err)
// 	}
//
// 	// Itterate over the results, printing the text from each tweet
// 	for k, result := range results {
// 		fmt.Println(k, result.Text)
// 	}
//

package twittertools

import (
	"encoding/json"
	"fmt"
	"math"
	"net"
	"os"
	"strconv"
)

import (
	"git.apache.org/thrift.git/lib/go/thrift"
	"mirgit.dcs.gla.ac.uk/JamesMcMinn/twitter-tools-go/ttthrift"
)

type Configuration struct {
	Host  string
	Port  int
	Group string
	Token string
}

type Client struct {
	Configuration
	transport  *thrift.TSocket
	TRECClient *ttthrift.TrecSearchClient
}

type Result struct {
	*ttthrift.TResult_
}

// Load a JSON configuration file which matches the format of the Configuration
// struct. Initiate a new Client using the details porvided.
func NewClientFromConf(filepath string) (client *Client, err error) {
	// Load configuration values from conf.js into the Configuration struct
	file, err := os.Open(filepath)
	if err != nil {
		fmt.Sprintln(os.Stderr, "Could not load conf.json\n\nError: %s", err)
		return nil, err
	}

	configuration := Configuration{}
	decoder := json.NewDecoder(file)
	if err = decoder.Decode(&configuration); err != nil {
		fmt.Sprintln(os.Stderr, "Could not parse conf.json\n\nError: %s", err)
		return nil, err
	}

	transport, trecClient, err := newClient(configuration.Host, configuration.Port,
		configuration.Group, configuration.Token)

	client = &Client{
		Configuration: configuration,
		transport:     transport,
		TRECClient:    trecClient,
	}

	return client, nil
}

// Initiate a new client using the details provided
func NewClientFromDetails(host string, port int, group, token string) (client *Client, err error) {
	transport, trecClient, err := newClient(host, port, group, token)

	client = &Client{
		Configuration: Configuration{
			Host:  host,
			Port:  port,
			Group: group,
			Token: token,
		},
		transport:  transport,
		TRECClient: trecClient,
	}

	return client, nil
}

// Create a connection to the Thrift server using the provided details.
func newClient(host string, port int, group, token string) (transport *thrift.TSocket, TRECClient *ttthrift.TrecSearchClient, err error) {
	portString := strconv.Itoa(port)
	transport, err = thrift.NewTSocket(net.JoinHostPort(host, portString))
	if err != nil {
		return nil, nil, err
	}

	protocolFactory := thrift.NewTBinaryProtocolFactoryDefault()
	TRECClient = ttthrift.NewTrecSearchClientFactory(transport, protocolFactory)
	if err := transport.Open(); err != nil {
		return nil, nil, err
	}

	return transport, TRECClient, nil
}

// Perform a search on the Thrift server.
func (t *Client) Search(q string, numResults int32, maxID int64) ([]Result, error) {
	if maxID < 0 {
		maxID = math.MaxInt64
	}
	query := ttthrift.NewTQuery()
	query.Group = t.Group
	query.Token = t.Token
	query.Text = q
	query.MaxId = maxID
	query.NumResults = numResults

	r, err := t.TRECClient.Search(query)
	if err != nil {
		return nil, err
	}

	results := []Result{}
	for k := range r {
		results = append(results, Result{r[k]})
	}

	return results, nil
}

// Close the current connection
func (t *Client) Close() {
	if t.transport != nil {
		t.transport.Close()
	}
}