restructure project

5 files changed, 158 insertions, 118 deletions

diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml
index 2f15d75..e70b4c4 100644
--- a/.github/workflows/release.yaml
+++ b/.github/workflows/release.yaml
@@ -29,7 +29,7 @@ jobs:
       - name: Build binary for ${{ matrix.arch }}
         run: |
           echo "Building for GOARCH=${{ matrix.arch }}"
-          go build -v -o v2stat-${{ matrix.arch }} .
+          go build -v -o v2stat go.rikki.moe/v2stat/cmd/v2stat
         env:
           GOARCH: ${{ matrix.arch }}
 
diff --git a/cmd/v2stat/daemon.go b/cmd/v2stat/daemon.go
new file mode 100644
index 0000000..8c5e91f
--- /dev/null
+++ b/cmd/v2stat/daemon.go
@@ -0,0 +1,45 @@
+package main
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"go.rikki.moe/v2stat"
+)
+
+func runDaemon(v2s *v2stat.V2Stat) {
+	if err := v2s.InitDB(); err != nil {
+		logger.Fatalf("Failed to initialize database: %v", err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+	defer signal.Stop(sigCh)
+
+	ticker := time.NewTicker(time.Duration(*flagInterval) * time.Second)
+	defer ticker.Stop()
+
+	for {
+		logger.Info("Recording stats...")
+		if err := v2s.RecordNow(ctx); err != nil {
+			logger.Errorf("Failed to record stats: %v", err)
+		}
+
+		select {
+		case <-ticker.C:
+			continue
+		case <-sigCh:
+			logger.Info("Received shutdown signal, exiting.")
+			return
+		case <-ctx.Done():
+			logger.Info("Context canceled, exiting.")
+			return
+		}
+	}
+}
diff --git a/main.go b/cmd/v2stat/main.go
index 7451840..ce0c3ed 100644
--- a/main.go
+++ b/cmd/v2stat/main.go
@@ -1,14 +1,10 @@
 package main
 
 import (
-	"context"
 	"database/sql"
 	"flag"
 	"fmt"
 	"os"
-	"os/signal"
-	"syscall"
-	"time"
 
 	"github.com/jedib0t/go-pretty/table"
 	"github.com/jedib0t/go-pretty/text"
@@ -16,7 +12,7 @@ import (
 	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 
-	"go.rikki.moe/v2stat/command"
+	"go.rikki.moe/v2stat"
 )
 
 var (
@@ -34,11 +30,7 @@ var DefaultDBPaths = []string{
 	"/opt/apps/v2stat/traffic.db",
 }
 
-type V2Stat struct {
-	logger *logrus.Logger
-	db     *sql.DB
-	stat   command.StatsServiceClient
-}
+var logger *logrus.Logger
 
 func main() {
 	flag.Parse()
@@ -50,14 +42,11 @@ func main() {
 		os.Exit(1)
 	}
 
-	logger := setupLogger(*flagLogLevel)
+	logger = setupLogger(*flagLogLevel)
 	db := setupDatabase(logger, *flagDatabase)
-	defer db.Close()
 
-	v2stat := &V2Stat{
-		logger: logger,
-		db:     db,
-	}
+	v2s := v2stat.NewV2Stat(logger, db, nil)
+	defer v2s.Close()
 
 	switch args[0] {
 	case "daemon":
@@ -65,12 +54,11 @@ func main() {
 		if err != nil {
 			logger.Fatalf("Failed to dial gRPC server: %v", err)
 		}
-		defer conn.Close()
-		v2stat.stat = command.NewStatsServiceClient(conn)
-		runServer(v2stat)
+		v2s.SetConn(conn)
+		runDaemon(v2s)
 
 	case "query":
-		handleQuery(v2stat, args[1:])
+		handleQuery(v2s, args[1:])
 
 	default:
 		fmt.Println("Unknown command:", args[0])
@@ -110,49 +98,13 @@ func setupDatabase(logger *logrus.Logger, dbpath string) *sql.DB {
 	return db
 }
 
-func runServer(v2stat *V2Stat) {
-	logger := v2stat.logger
-
-	if err := v2stat.InitDB(); err != nil {
-		logger.Fatalf("Failed to initialize database: %v", err)
-	}
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
-	defer signal.Stop(sigCh)
-
-	ticker := time.NewTicker(time.Duration(*flagInterval) * time.Second)
-	defer ticker.Stop()
-
-	for {
-		logger.Info("Recording stats...")
-		if err := v2stat.RecordNow(ctx); err != nil {
-			logger.Errorf("Failed to record stats: %v", err)
-		}
-
-		select {
-		case <-ticker.C:
-			continue
-		case <-sigCh:
-			logger.Info("Received shutdown signal, exiting.")
-			return
-		case <-ctx.Done():
-			logger.Info("Context canceled, exiting.")
-			return
-		}
-	}
-}
-
-func handleQuery(v2stat *V2Stat, args []string) {
+func handleQuery(v2s *v2stat.V2Stat, args []string) {
 	if len(args) == 0 {
 		fmt.Println("Usage: v2stat query <connection_name>")
 		fmt.Println("Available connections:")
-		conns, err := v2stat.QueryConn()
+		conns, err := v2s.QueryConn()
 		if err != nil {
-			v2stat.logger.Fatalf("Failed to query connection: %v", err)
+			logger.Fatalf("Failed to query connection: %v", err)
 		}
 		for _, c := range conns {
 			fmt.Printf("\t%s\n", c.String())
@@ -161,20 +113,20 @@ func handleQuery(v2stat *V2Stat, args []string) {
 	}
 
 	connStr := args[0]
-	conn, ok := ParseConnInfo(connStr)
+	conn, ok := v2stat.ParseConnInfo(connStr)
 	if !ok {
-		v2stat.logger.Fatalf("Invalid connection format: %s", connStr)
+		logger.Fatalf("Invalid connection format: %s", connStr)
 	}
 
-	stats, err := v2stat.QueryStatsHourly(&conn)
+	stats, err := v2s.QueryStatsHourly(&conn)
 	if err != nil {
-		v2stat.logger.Fatalf("Failed to query stats: %v", err)
+		logger.Fatalf("Failed to query stats: %v", err)
 	}
 
 	printStatsTable(stats)
 }
 
-func printStatsTable(stats []TrafficStat) {
+func printStatsTable(stats []v2stat.TrafficStat) {
 	tb := table.NewWriter()
 	tb.SetOutputMirror(os.Stdout)
 	tb.AppendHeader(table.Row{"Time", "Downlink", "Uplink"})
@@ -204,4 +156,4 @@ func sizeToHuman(size int64) string {
 		value /= 1024
 	}
 	return fmt.Sprintf("%7.2f %s", value, "PiB")
-}
\ No newline at end of file
+}
diff --git a/def.go b/def.go
new file mode 100644
index 0000000..e9951ad
--- /dev/null
+++ b/def.go
@@ -0,0 +1,66 @@
+package v2stat
+
+import (
+	"strings"
+)
+
+type TrafficDirection int
+
+const (
+	DirectionDownlink TrafficDirection = iota
+	DirectionUplink
+)
+
+type ConnectionType int
+
+const (
+	ConnTypeUser ConnectionType = iota
+	ConnTypeInbound
+	ConnTypeOutbound
+)
+
+type ConnInfo struct {
+	Type ConnectionType `json:"type"`
+	Name string         `json:"name"`
+}
+
+type TrafficStat struct {
+	Time     string `json:"time"`
+	Downlink int64  `json:"downlink"`
+	Uplink   int64  `json:"uplink"`
+}
+
+func (ci *ConnInfo) String() string {
+	switch ci.Type {
+	case ConnTypeUser:
+		return "user:" + ci.Name
+	case ConnTypeInbound:
+		return "inbound:" + ci.Name
+	case ConnTypeOutbound:
+		return "outbound:" + ci.Name
+	default:
+		return "unknown:" + ci.Name
+	}
+}
+
+func ParseConnInfo(s string) (ConnInfo, bool) {
+	parts := strings.Split(s, ":")
+	if len(parts) != 2 {
+		return ConnInfo{}, false
+	}
+	var connType ConnectionType
+	switch parts[0] {
+	case "user":
+		connType = ConnTypeUser
+	case "inbound":
+		connType = ConnTypeInbound
+	case "outbound":
+		connType = ConnTypeOutbound
+	default:
+		return ConnInfo{}, false
+	}
+	return ConnInfo{
+		Type: connType,
+		Name: parts[1],
+	}, true
+}
\ No newline at end of file
diff --git a/db.go b/v2stat.go
index 8f7a5c7..49b37a8 100644
--- a/db.go
+++ b/v2stat.go
@@ -1,72 +1,47 @@
-package main
+package v2stat
 
 import (
 	"context"
+	"database/sql"
 	"strings"
 	"time"
 
+	"github.com/sirupsen/logrus"
 	"go.rikki.moe/v2stat/command"
+	"google.golang.org/grpc"
 )
 
-type TrafficDirection int
-
-const (
-	DirectionDownlink TrafficDirection = iota
-	DirectionUplink
-)
-
-type ConnectionType int
-
-const (
-	ConnTypeUser ConnectionType = iota
-	ConnTypeInbound
-	ConnTypeOutbound
-)
-
-type ConnInfo struct {
-	Type ConnectionType `json:"type"`
-	Name string         `json:"name"`
+type V2Stat struct {
+	logger *logrus.Logger
+	db     *sql.DB
+	conn *grpc.ClientConn
+	stat   command.StatsServiceClient
 }
 
-type TrafficStat struct {
-	Time     string `json:"time"`
-	Downlink int64  `json:"downlink"`
-	Uplink   int64  `json:"uplink"`
+func NewV2Stat(logger *logrus.Logger, db *sql.DB, conn *grpc.ClientConn) *V2Stat {
+	return &V2Stat{
+		logger: logger,
+		db:     db,
+		conn:   conn,
+		stat:   command.NewStatsServiceClient(conn),
+	}
 }
 
-func (ci *ConnInfo) String() string {
-	switch ci.Type {
-	case ConnTypeUser:
-		return "user:" + ci.Name
-	case ConnTypeInbound:
-		return "inbound:" + ci.Name
-	case ConnTypeOutbound:
-		return "outbound:" + ci.Name
-	default:
-		return "unknown:" + ci.Name
+func (v *V2Stat) Close() {
+	if v.db != nil {
+		v.db.Close()
+	}
+	if v.conn != nil {
+		v.conn.Close()
 	}
 }
 
-func ParseConnInfo(s string) (ConnInfo, bool) {
-	parts := strings.Split(s, ":")
-	if len(parts) != 2 {
-		return ConnInfo{}, false
-	}
-	var connType ConnectionType
-	switch parts[0] {
-	case "user":
-		connType = ConnTypeUser
-	case "inbound":
-		connType = ConnTypeInbound
-	case "outbound":
-		connType = ConnTypeOutbound
-	default:
-		return ConnInfo{}, false
+func (v *V2Stat) SetConn(conn *grpc.ClientConn) {
+	if v.conn != nil {
+		v.conn.Close()
 	}
-	return ConnInfo{
-		Type: connType,
-		Name: parts[1],
-	}, true
+	v.conn = conn
+	v.stat = command.NewStatsServiceClient(conn)
 }
 
 
@@ -255,3 +230,5 @@ func (v *V2Stat) QueryStatsHourly(conn *ConnInfo) ([]TrafficStat, error) {
 	}
 	return stats, nil		
 }
+
+