✨ Add Cassandra support and update README.md

app/config/cassandra.env.example

@@ -0,0 +1,9 @@
+CASSANDRA_USER=user
+CASSANDRA_PASSWORD=password
+CASSANDRA_KEYSPACE=ohmyurl
+CASSANDRA_DB=url
+CASSANDRA_CLUSTERS=ohmyurl-cassandra-1:9042,ohmyurl-cassandra-2:9042,ohmyurl-cassandra-3:9042
+LISTEN_PORT_CREATE=8080
+LISTEN_PORT_FORWARD=8081
+FORWARD_DOMAIN=http://localhost/
+CREATE_DOMAIN=http://127.0.0.1

app/config/settings.go

@@ -2,6 +2,7 @@ package config
 
 import (
 	"os"
+	"strconv"
 	"strings"
 )
 
@@ -28,6 +29,16 @@ type PostgresConfig struct {
 	DBNAME   string
 }
 
+type CassandraConfig struct {
+	USER         string
+	PASSWORD     string
+	KEYSPACE     string
+	CLUSTERS     []string
+	APP_LABEL    string
+	URL_START_ID int
+	URL_END_ID   int
+}
+
 func GetAppConfig() *AppConfig {
 	return &AppConfig{
 		LISTEN_PORT_CREATE:  GetEnvOrPanic("LISTEN_PORT_CREATE"),
@@ -55,6 +66,18 @@ func GetPostgresConfig() *PostgresConfig {
 	}
 }
 
+func GetCassandraConfig() *CassandraConfig {
+	return &CassandraConfig{
+		USER:         GetEnvOrPanic("CASSANDRA_USER"),
+		PASSWORD:     GetEnvOrPanic("CASSANDRA_PASSWORD"),
+		CLUSTERS:     strings.Split(GetEnvOrPanic("CASSANDRA_CLUSTERS"), ","),
+		KEYSPACE:     GetEnvOrPanic("CASSANDRA_KEYSPACE"),
+		APP_LABEL:    GetEnvOrPanic("CASSANDRA_APP_LABEL"),
+		URL_START_ID: Str2IntOrPanic(GetEnvOrPanic("CASSANDRA_URL_START_ID")),
+		URL_END_ID:   Str2IntOrPanic(GetEnvOrPanic("CASSANDRA_URL_END_ID")),
+	}
+}
+
 func GetDB() DBName {
 	dbName := strings.ToLower(GetEnvOrPanic("DB"))
 	switch dbName {
@@ -84,3 +107,11 @@ func GetEnvOrPanic(key string) string {
 	}
 	return value
 }
+
+func Str2IntOrPanic(value string) int {
+	i, err := strconv.Atoi(value)
+	if err != nil {
+		panic(err)
+	}
+	return i
+}

app/db/base.go

@@ -3,6 +3,7 @@ package db
 import (
 	"fmt"
 	"github.com/aykhans/oh-my-url/app/config"
+	"github.com/gocql/gocql"
 	gormPostgres "gorm.io/driver/postgres"
 	"gorm.io/gorm"
 	"time"
@@ -43,6 +44,32 @@ func GetDB() DB {
 			panic(err)
 		}
 		return &Postgres{gormDB: db}
+
+	case "cassandra":
+		cassandraConf := config.GetCassandraConfig()
+		cluster := gocql.NewCluster(cassandraConf.CLUSTERS...)
+		cluster.Keyspace = cassandraConf.KEYSPACE
+		cluster.Authenticator = gocql.PasswordAuthenticator{
+			Username: cassandraConf.USER,
+			Password: cassandraConf.PASSWORD,
+		}
+
+		var db *gocql.Session
+		var err error
+		for i := 0; i < 60; i++ {
+			db, err = cluster.CreateSession()
+			if err == nil {
+				break
+			}
+			time.Sleep(3 * time.Second)
+		}
+		if err != nil {
+			panic(err)
+		}
+		return &Cassandra{
+			db: db,
+			currentID: &CurrentID{},
+		}
 	}
-	return nil
+	panic("unknown db")
 }

app/db/cassandra.go

@@ -1 +1,99 @@
 package db
+
+import (
+	"errors"
+	"log"
+	"sync"
+
+	"github.com/aykhans/oh-my-url/app/utils"
+	"github.com/gocql/gocql"
+)
+
+type CurrentID struct {
+	ID int
+	mu sync.RWMutex
+}
+
+type Cassandra struct {
+	db        *gocql.Session
+	currentID *CurrentID
+}
+
+func (c *Cassandra) Init() {
+	err := c.db.Query(`
+		CREATE TABLE IF NOT EXISTS url (
+			id int,
+			key text,
+			url text,
+			count int,
+			PRIMARY KEY ((key), count)
+		) WITH CLUSTERING ORDER BY (count DESC)
+		AND compaction = {'class': 'LeveledCompactionStrategy'};
+	`).Consistency(gocql.All).Exec()
+	if err != nil {
+		panic(err)
+	}
+
+	var id int
+	err = c.db.
+		Query(`SELECT MAX(id) FROM url LIMIT 1`).
+		Consistency(gocql.One).
+		Scan(&id)
+
+	if err != nil {
+		panic(err)
+	}
+	c.currentID.ID = id
+}
+
+func (c *Cassandra) CreateURL(url string) (string, error) {
+	c.currentID.mu.Lock()
+	defer c.currentID.mu.Unlock()
+
+	id := c.currentID.ID + 1
+	key := utils.GenerateKey(id)
+	m := make(map[string]interface{})
+
+	query := `INSERT INTO url (id, key, url, count) VALUES (?, ?, ?, ?) IF NOT EXISTS`
+	applied, err := c.db.Query(query, id, key, url, 0).Consistency(gocql.All).MapScanCAS(m)
+	if err != nil {
+		log.Println(err)
+		return "", err
+	}
+	if !applied {
+		log.Println("Failed to insert unique key")
+		return "", errors.New("an error occurred, please try again later")
+	}
+	c.currentID.ID = id
+
+	return key, nil
+}
+
+func (c *Cassandra) GetURL(key string) (string, error) {
+	var url string
+	err := c.db.
+		Query(`SELECT url FROM url WHERE key = ? LIMIT 1`, key).
+		Consistency(gocql.One).
+		Scan(&url)
+	if err != nil {
+		return "", err
+	}
+	return url, nil
+}
+
+// just in case
+// const urlKeyCreateFunction = `
+// CREATE FUNCTION IF NOT EXISTS oh_my_url.generate_url_key(n int) RETURNS NULL ON NULL INPUT RETURNS text LANGUAGE java AS '
+//     String keyCharacters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
+//     StringBuilder result = new StringBuilder();
+//     int base = keyCharacters.length();
+
+//     while (n > 0) {
+//         n--;
+//         result.append(keyCharacters.charAt(n % base));
+//         n = (int) Math.floor(n / base);
+//     }
+
+//     return result.reverse().toString();
+// ';
+// `