github.com/Ryan-Johnson-1315/socketlogger@v0.0.2/csv.go (about)

     1  package socketlogger
     2  
     3  import (
     4  	"encoding/csv"
     5  	"fmt"
     6  	"log"
     7  	"os"
     8  	"path/filepath"
     9  	"strings"
    10  )
    11  
    12  type CsvServer interface {
    13  	SetOutputCsvDirectory(string)
    14  	Server
    15  }
    16  
    17  type csvserver struct {
    18  	writers   map[string]*csv.Writer
    19  	outputDir string
    20  	flush     chan bool
    21  }
    22  
    23  func (c *csvserver) SetOutputCsvDirectory(dir string) {
    24  	c.outputDir = dir
    25  	if !fileDirExists(dir, "") {
    26  		err := os.MkdirAll(dir, os.ModePerm)
    27  		if err != nil {
    28  			log.Print(newLogMessage(MessageLevelWrn, "Could not make directory \"%s\", %v", dir, err))
    29  		}
    30  	}
    31  }
    32  
    33  func (c *csvserver) buildCsvFile(msg *CsvMessage) *csv.Writer {
    34  	if !fileDirExists(c.outputDir, msg.Filename) ||
    35  		(fileDirExists(c.outputDir, msg.Filename) && c.writers[msg.Filename] == nil) {
    36  		duplicate := false
    37  		fname := filepath.Join(c.outputDir, msg.Filename)
    38  		for i := 1; ; i++ {
    39  			if fileExists(fname) {
    40  				duplicate = true
    41  				fname = filepath.Join(c.outputDir, fmt.Sprintf("%s_%d.csv", strings.Split(msg.Filename, ".csv")[0], i))
    42  			} else {
    43  				break
    44  			}
    45  		}
    46  
    47  		if duplicate {
    48  			log.Print(newLogMessage(MessageLevelWrn, "Found previous %s, creating %s", msg.Filename, fname))
    49  		}
    50  
    51  		fptr, err := os.OpenFile(fname, os.O_CREATE|os.O_APPEND|os.O_RDWR, 0o666)
    52  
    53  		if err != nil {
    54  			log.Print(newLogMessage(MessageLevelErr, "Could not open file: %s -> %v", fname, err))
    55  			return nil
    56  		} else {
    57  			log.Print(newLogMessage(MessageLevelSuccess, "File created %s", fname))
    58  		}
    59  
    60  		csvWriter := csv.NewWriter(fptr)
    61  		if err != nil {
    62  			log.Print(newLogMessage(MessageLevelErr, "Could not open build csv.Writer -> %v", err))
    63  			return nil
    64  		}
    65  		c.writers[msg.Filename] = csvWriter
    66  	}
    67  
    68  	return c.writers[msg.Filename]
    69  }
    70  
    71  func (c *csvserver) getMessageType() SocketMessage {
    72  	return &CsvMessage{}
    73  }
    74  
    75  func (c *csvserver) initCsvServer() {
    76  	c.writers = make(map[string]*csv.Writer)
    77  }
    78  
    79  func (c *csvserver) write(msgs chan SocketMessage) {
    80  	for msg := range msgs {
    81  		if msg.Type() == Csv {
    82  			inst := msg.(*CsvMessage)
    83  			if inst.Filename != "" {
    84  				writer := c.buildCsvFile(inst)
    85  				if writer == nil {
    86  					log.Print(newLogMessage(MessageLevelErr, "csv writer returned as nil!"))
    87  					continue
    88  				}
    89  				// Only need to write the row if it is there
    90  				if len(inst.Row) > 0 {
    91  					writer.Write(transform(inst.Row))
    92  					writer.Flush() // flushes headers & data
    93  				}
    94  			}
    95  		}
    96  	}
    97  	c.flush <- true
    98  }
    99  
   100  func (c *csvserver) setFlushChannel(flush chan bool) {
   101  	c.flush = flush
   102  }
   103  
   104  type TcpCsvServer struct {
   105  	tcpserver
   106  	csvserver
   107  }
   108  
   109  func NewTcpCsvServer() CsvServer {
   110  	t := &TcpCsvServer{}
   111  	t.init(t)
   112  	t.initCsvServer()
   113  	return t
   114  }
   115  
   116  type UdpCsvServer struct {
   117  	udpserver
   118  	csvserver
   119  }
   120  
   121  func NewUdpCsvServer() CsvServer {
   122  	u := &UdpCsvServer{}
   123  	u.init(u)
   124  	u.initCsvServer()
   125  	return u
   126  }
   127  
   128  type CsvClient interface {
   129  	NewCsvFile(fname string, headers []interface{})
   130  	AppendRow(fname string, row []interface{})
   131  	Client
   132  }
   133  
   134  type csvclient struct {
   135  	msgsToSend chan SocketMessage
   136  }
   137  
   138  func (c *csvclient) setMsgChannel(msgsToSend chan SocketMessage) {
   139  	c.msgsToSend = msgsToSend
   140  }
   141  
   142  func (c *csvclient) NewCsvFile(fname string, headers []interface{}) {
   143  	c.msgsToSend <- newCsvMessage(fname, headers)
   144  }
   145  
   146  func (c *csvclient) AppendRow(fname string, row []interface{}) {
   147  	c.msgsToSend <- newCsvMessage(fname, row)
   148  }
   149  
   150  type UdpCsvClient struct {
   151  	csvclient
   152  	udpClient
   153  }
   154  
   155  func NewUdpCsvClient() CsvClient {
   156  	u := &UdpCsvClient{}
   157  	u.init(u)
   158  	return u
   159  }
   160  
   161  type TcpCsvClient struct {
   162  	csvclient
   163  	tcpClient
   164  }
   165  
   166  func NewTcpCsvClient() CsvClient {
   167  	t := &TcpCsvClient{}
   168  	t.init(t)
   169  	return t
   170  }