Merge pull request #62 from Hu6li/client

Implementation of small client

Author: meeehow

client/client.go

@@ -0,0 +1,81 @@
+package main
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"flag"
+	"fmt"
+
+	"cloud.google.com/go/spanner"
+	"github.com/google/hashr/client/cloudspanner"
+	"github.com/google/hashr/client/postgres"
+	_ "github.com/lib/pq"
+
+	"github.com/golang/glog"
+)
+
+var (
+	hashStorage   = flag.String("hashStorage", "", "Storage used for computed hashes, can have one of the two values: postgres, cloudspanner")
+	spannerDBPath = flag.String("spanner_db_path", "", "Path to spanner DB.")
+
+	// Postgres DB flags
+	postgresHost     = flag.String("postgres_host", "localhost", "PostgreSQL instance address.")
+	postgresPort     = flag.Int("postgres_port", 5432, "PostgresSQL instance port.")
+	postgresUser     = flag.String("postgres_user", "hashr", "PostgresSQL user.")
+	postgresPassword = flag.String("postgres_password", "hashr", "PostgresSQL password.")
+	postgresDBName   = flag.String("postgres_db", "hashr", "PostgresSQL database.")
+)
+
+// Storage represents  storage that is used to store data about processed sources.
+type Storage interface {
+	GetSamples(ctx context.Context) (map[string]map[string]string, error)
+}
+
+func main() {
+	ctx := context.Background()
+	flag.Parse()
+
+	var storage Storage
+	switch *hashStorage {
+	case "postgres":
+		psqlInfo := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
+			*postgresHost, *postgresPort, *postgresUser, *postgresPassword, *postgresDBName)
+
+		db, err := sql.Open("postgres", psqlInfo)
+		if err != nil {
+			glog.Exitf("Error initializing Postgres client: %v", err)
+		}
+		defer db.Close()
+
+		storage, err = postgres.NewStorage(db)
+		if err != nil {
+			glog.Exitf("Error initializing Postgres storage: %v", err)
+		}
+	case "cloudspanner":
+		spannerClient, err := spanner.NewClient(ctx, *spannerDBPath)
+		if err != nil {
+			glog.Exitf("Error initializing Spanner client: %v", err)
+		}
+
+		storage, err = cloudspanner.NewStorage(ctx, spannerClient)
+		if err != nil {
+			glog.Exitf("Error initializing Postgres storage: %v", err)
+		}
+	default:
+		glog.Exit("hashStorage flag needs to have one of the two values: postgres, cloudspanner")
+
+	}
+	samples, err := storage.GetSamples(ctx)
+	if err != nil {
+		glog.Exitf("Error retriving samples: %v", err)
+	}
+
+	jsonData, err := json.Marshal(samples)
+	if err != nil {
+		fmt.Println("Error:", err)
+		return
+	}
+
+	fmt.Println(string(jsonData))
+}

client/cloudspanner/cloudspanner.go

@@ -0,0 +1,64 @@
+package cloudspanner
+
+import (
+	"context"
+	"strconv"
+
+	"cloud.google.com/go/spanner"
+
+	"google.golang.org/api/iterator"
+)
+
+// Storage allows to interact with cloud spanner.
+type Storage struct {
+	spannerClient *spanner.Client
+}
+
+// NewStorage creates new Storage struct that allows to interact with cloud spanner.
+func NewStorage(ctx context.Context, spannerClient *spanner.Client) (*Storage, error) {
+	return &Storage{spannerClient: spannerClient}, nil
+}
+
+// GetSamples fetches processing samples from cloud spanner.
+func (s *Storage) GetSamples(ctx context.Context) (map[string]map[string]string, error) {
+	samples := make(map[string]map[string]string)
+	iter := s.spannerClient.Single().Read(ctx, "samples",
+		spanner.AllKeys(), []string{"sha256", "mimetype", "file_output", "size"})
+	defer iter.Stop()
+	for {
+		row, err := iter.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			return nil, err
+		}
+		var sha256, mimetype, fileOutput string
+		var size int64
+		err = row.ColumnByName("sha256", &sha256)
+		if err != nil {
+			return nil, err
+		}
+		err = row.ColumnByName("mimetype", &mimetype)
+		if err != nil {
+			return nil, err
+		}
+		err = row.ColumnByName("file_output", &fileOutput)
+		if err != nil {
+			return nil, err
+		}
+		err = row.ColumnByName("size", &size)
+		if err != nil {
+			return nil, err
+		}
+		samples[sha256] = make(map[string]string)
+
+		// Assign values to the nested map
+		samples[sha256]["sha256"] = sha256
+		samples[sha256]["mimetype"] = mimetype
+		samples[sha256]["file_output"] = fileOutput
+		samples[sha256]["size"] = strconv.FormatInt(size, 10)
+
+	}
+	return samples, nil
+}

client/postgres/postgres.go

@@ -0,0 +1,83 @@
+// Package postgres implements PostgreSQL as a hashR storage.
+package postgres
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+
+	// Blank import below is needed for the SQL driver.
+	_ "github.com/lib/pq"
+)
+
+// Storage allows to interact with PostgreSQL instance.
+type Storage struct {
+	sqlDB *sql.DB
+}
+
+// NewStorage creates new Storage struct that allows to interact with PostgreSQL instance and all the necessary tables, if they don't exist.
+func NewStorage(sqlDB *sql.DB) (*Storage, error) {
+	return &Storage{sqlDB: sqlDB}, nil
+}
+
+// GetSamples fetches processed samples from postgres.
+func (s *Storage) GetSamples(ctx context.Context) (map[string]map[string]string, error) {
+	exists, err := tableExists(s.sqlDB, "samples")
+	if err != nil {
+		return nil, err
+	}
+
+	samples := make(map[string]map[string]string)
+
+	if exists {
+		var sql = `SELECT * FROM samples;`
+
+		rows, err := s.sqlDB.Query(sql)
+
+		if err != nil {
+			return nil, err
+		}
+
+		defer rows.Close()
+
+		for rows.Next() {
+			var sha256, mimetype, fileOutput, size string
+			err := rows.Scan(&sha256, &mimetype, &fileOutput, &size)
+			if err != nil {
+				return nil, err
+			}
+
+			samples[sha256] = make(map[string]string)
+
+			// Assign values to the nested map
+			samples[sha256]["sha256"] = sha256
+			samples[sha256]["mimetype"] = mimetype
+			samples[sha256]["file_output"] = fileOutput
+			samples[sha256]["size"] = size
+		}
+
+	} else {
+		return nil, fmt.Errorf("table samples does not exist")
+	}
+
+	return samples, nil
+}
+
+func tableExists(db *sql.DB, tableName string) (bool, error) {
+	// Query to check if the table exists in PostgreSQL
+	query := `
+        SELECT EXISTS (
+            SELECT 1
+            FROM   information_schema.tables
+            WHERE  table_name = $1
+        )
+    `
+
+	var exists bool
+	err := db.QueryRow(query, tableName).Scan(&exists)
+	if err != nil {
+		return false, err
+	}
+
+	return exists, nil
+}