@@ -70,6 +70,16 @@ func loadFile(responseArr []string, files []entity.File, db *sql.DB, tableCount
7070 if len (responseArr ) == 3 {
7171 path := responseArr [1 ]
7272 tableName := responseArr [2 ]
73+ mimeType , err := helpers .GetMimeType (path )
74+ if err != nil {
75+ fmt .Printf ("Error reading mimetype of file %v : %v\n " , path , err .Error ())
76+ return files
77+ }
78+
79+ if ! validFileType (mimeType ) {
80+ fmt .Printf ("Error reading file %v : file type not supported\n " , path )
81+ return files
82+ }
7383
7484 for i := range files {
7585 if files [i ].Table == tableName {
@@ -78,59 +88,72 @@ func loadFile(responseArr []string, files []entity.File, db *sql.DB, tableCount
7888 }
7989 }
8090
81- if _ , err := os .Stat (path ); err == nil || os .IsExist (err ) {
82- for {
83- fmt .Println ("File has a header row (y/n)?" )
84- response := strings .ToUpper (strings .TrimSpace (prompt .Input ("> " , helpers .Completer )))
85- if response == "Y" || response == "N" {
86- content , fileErr := helpers .ReadCSVFile (path )
87- if fileErr != nil {
88- fmt .Printf ("Error reading file %v : %v\n " , path , fileErr .Error ())
89- break
90- }
91+ for {
92+ fmt .Println ("File has a header row (y/n)?" )
93+ response := strings .ToUpper (strings .TrimSpace (prompt .Input ("> " , helpers .Completer )))
94+ if response == "Y" || response == "N" {
95+ var (
96+ content [][]string
97+ fileErr error
98+ headers = make ([]string , 0 )
99+ )
100+
101+ if mimeType == "text/csv" {
102+ content , fileErr = helpers .ReadCSVFile (path )
103+ }
91104
92- if len (content ) == 0 {
93- fmt .Println ("Empty file" )
94- break
105+ if mimeType == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" {
106+ sheets , err := helpers .GetXLSXSheetList (path )
107+ if err != nil {
108+ fmt .Printf ("Error reading file %v : %v\n " , path , err .Error ())
109+ return files
95110 }
111+ fmt .Printf ("What is the name of the sheet to be loaded? %v\n " , strings .Join (sheets , ", " ))
112+ sheet := strings .TrimSpace (prompt .Choose ("> " , sheets ))
113+ content , fileErr = helpers .ReadXLSXFile (path , sheet )
114+ }
96115
97- headers := make ([]string , 0 )
98-
99- if response == "Y" {
100- headers = content [0 ]
101- content = content [1 :]
102- } else {
103- for {
104- fmt .Println ("Enter " + strconv .Itoa (len (content [0 ])) + " headers seperated by commas" )
105- headers = strings .Split (strings .TrimSpace (prompt .Input ("> " , helpers .Completer )), "," )
106- if len (content [0 ]) == len (headers ) {
107- break
108- }
109- }
110- }
116+ if fileErr != nil {
117+ fmt .Printf ("Error reading file %v : %v\n " , path , fileErr .Error ())
118+ return files
119+ }
111120
112- for i := range headers {
113- headers [i ] = strings .TrimSpace (headers [i ])
114- }
121+ if len (content ) == 0 {
122+ fmt .Println ("Empty file" )
123+ return files
124+ }
115125
116- file := entity.File {
117- Path : path ,
118- Headers : headers ,
119- Table : strings .TrimSpace (tableName ),
126+ if response == "Y" {
127+ headers = content [0 ]
128+ content = content [1 :]
129+ } else {
130+ for {
131+ fmt .Println ("Enter " + strconv .Itoa (len (content [0 ])) + " headers seperated by commas" )
132+ headers = strings .Split (strings .TrimSpace (prompt .Input ("> " , helpers .Completer )), "," )
133+ if len (content [0 ]) == len (headers ) {
134+ break
135+ }
120136 }
137+ }
121138
122- files = append (files , file )
123- file .Content = content
124- helpers .PopulateTables (db , & file )
125- tableCount ++
126- break
139+ for i := range headers {
140+ headers [i ] = strings .TrimSpace (headers [i ])
127141 }
142+
143+ file := entity.File {
144+ Path : path ,
145+ Headers : headers ,
146+ Table : strings .TrimSpace (tableName ),
147+ }
148+
149+ files = append (files , file )
150+ file .Content = content
151+ helpers .PopulateTables (db , & file )
152+ tableCount ++
153+ break
128154 }
129- } else {
130- fmt .Printf ("File doesn't exists: %v\n " , path )
131155 }
132- } else {
133- fmt .Println ("Please use LOAD /path/to/file.csv table_name" )
156+
134157 }
135158 return files
136159}
@@ -166,6 +189,10 @@ func createDB(dbName string) *sql.DB {
166189 return db
167190}
168191
192+ func validFileType (mimeType string ) bool {
193+ return mimeType == "text/csv" || mimeType == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
194+ }
195+
169196func handleExit () {
170197 rawModeOff := exec .Command ("/bin/stty" , "-raw" , "echo" )
171198 rawModeOff .Stdin = os .Stdin
0 commit comments