🚧 This instance is under construction; expect occasional downtime. Runners available in /repos. Questions? Ask in #wikimedia-gitlab on libera.chat, or under GitLab on Phabricator.

Commit 45926e1e authored by Gmodena's avatar Gmodena
Browse files

Add cassandra client example

parent 5b02fa88
build:
go build .
run:
go run .
test:
echo "Not implemented"
.PHONY: build run test
package main
import (
"github.com/gocql/gocql"
"log"
"os"
"strconv"
)
const (
CASSANDRA_KEYSPACE = "keyspace"
CASSANDRA_HOST = "localhost"
CASSANDRA_USERNAME = "cassandra"
CASSANDRA_PASSWORD = "cassandra" // TODO: storing a password in env var is a bad idea. Maybe move secrets to a config file?
CASSANDRA_CQL_PORT = "9042"
PAGE_SIZE = 10
)
type Cassandra struct {
// get version, host, date, etc. from Makefile
}
type clusterConfig struct {
hosts string // TODO(gmodena) should be a list of hosts
keyspace string
cql_port int
authenticator gocql.Authenticator
}
func lookupEnvOrElse(key string, fallback string) string {
val, err := os.LookupEnv(key)
if err == false {
val = fallback
}
return val
}
func getConfig() clusterConfig {
host := lookupEnvOrElse("CASSANDRA_HOST", CASSANDRA_HOST)
username := lookupEnvOrElse("CASSANDRA_USERNAME", CASSANDRA_USERNAME)
password := lookupEnvOrElse("CASSANDRA_PASSWORD", CASSANDRA_PASSWORD)
keyspace := lookupEnvOrElse("CASSANDRA_KEYSPACE", CASSANDRA_KEYSPACE)
cql_port, err := strconv.Atoi(lookupEnvOrElse("CASSANDRA_CQL_PORT", CASSANDRA_CQL_PORT))
if err != nil {
log.Fatal(err)
}
authenticator := gocql.PasswordAuthenticator{Username: username, Password: password}
return clusterConfig{hosts: host, keyspace: keyspace, cql_port: cql_port, authenticator: authenticator}
}
func initCluster(config clusterConfig) *gocql.ClusterConfig {
cluster := gocql.NewCluster(config.hosts)
cluster.Authenticator = config.authenticator
cluster.Keyspace = config.keyspace
cluster.Consistency = gocql.One
cluster.Port = config.cql_port
cluster.PageSize = PAGE_SIZE
return cluster
}
func IterRows(query gocql.Query) chan map[string]interface{} {
err := query.Exec()
if err != nil {
log.Fatal(err)
}
iter := query.Iter()
c := make(chan map[string]interface{})
go func() {
rowValues := make(map[string]interface{})
for iter.MapScan(rowValues) {
c <- rowValues
rowValues = make(map[string]interface{})
}
close(c)
}()
return c
}
func getSession() *gocql.Session {
config := getConfig()
cluster := initCluster(config)
session, err := cluster.CreateSession()
if err != nil {
log.Fatal(err)
}
return session
}
func GetCassandraSession() *gocql.Session {
return getSession()
}
\ No newline at end of file
package main
import (
"encoding/json"
"fmt"
"log"
"net/http"
)
type ImageMatching struct {
Prediction [] map[string]interface{} `json:"prediction"`
ModelVersion string `json:"model_version"`
}
func (* ImageMatching) ServeHTTP(w http.ResponseWriter, r *http.Request) {
session := GetCassandraSession() // we should not re-init at every request.
wiki, ok := r.URL.Query()["wiki"]
if ok {}
page_id, ok := r.URL.Query()["page_id"]
if ok {}
maxrows := 3
stmt := fmt.Sprintf("select image_id from imagerec.matches where wiki = '%s' and page_id = '%s' limit %d;", wiki[0], page_id[0], maxrows)
query := session.Query(stmt)
var recs []map[string]interface{};
for row := range IterRows(*query) {
recs = append(recs, row)
}
ima := ImageMatching{Prediction: recs, ModelVersion: "1a"}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
response, err := json.Marshal(ima)
if err != nil {
log.Print(err)
}
w.Write(response)
}
package main
import "net/http"
type server struct{}
var version = "unknown"
var buildDate = "unknown"
func main() {
ima := ImageMatching{}
http.Handle("/predict", &ima)
http.ListenAndServe(":8000", nil)
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment