github.com/infraboard/keyauth@v0.8.1/apps/ip2region/impl/ip2region.go (about) 1 package impl 2 3 import ( 4 "bytes" 5 "fmt" 6 "os" 7 8 "github.com/infraboard/mcube/exception" 9 10 "github.com/infraboard/keyauth/apps/ip2region" 11 "github.com/infraboard/keyauth/apps/ip2region/reader" 12 "github.com/infraboard/keyauth/apps/storage" 13 ) 14 15 func (s *service) UpdateDBFile(req *ip2region.UpdateDBFileRequest) error { 16 if err := req.Validate(); err != nil { 17 return exception.NewBadRequest("validate update db file requrest error, %s", err) 18 } 19 20 reader := req.ReadCloser() 21 defer reader.Close() 22 23 uploadReq := storage.NewUploadFileRequest(s.bucketName, s.dbFileName, reader) 24 uploadReq.WithToken(req.GetToken()) 25 return s.storage.UploadFile(uploadReq) 26 } 27 28 func (s *service) LookupIP(ip string) (*ip2region.IPInfo, error) { 29 dbReader, err := s.getDBReader() 30 if err != nil { 31 return nil, err 32 } 33 34 return dbReader.MemorySearch(ip) 35 } 36 37 func (s *service) getDBReader() (*reader.IPReader, error) { 38 s.Lock() 39 defer s.Unlock() 40 41 if s.dbReader != nil { 42 return s.dbReader, nil 43 } 44 45 // 优先从本地文件加载DB文件 46 if err := s.loadDBFileFromLocal(); err != nil { 47 s.log.Infof("load ip2region db file from local error, %s, retry other load method ", err) 48 } else { 49 return s.dbReader, nil 50 } 51 52 if err := s.loadDBFileFromBucket(); err != nil { 53 s.log.Infof("load ip2region db file from bucket error, %s", err) 54 } else { 55 return s.dbReader, nil 56 } 57 58 return nil, fmt.Errorf("load ip2region db file error") 59 } 60 61 func (s *service) loadDBFileFromLocal() error { 62 file, err := os.Open(s.dbFileName) 63 if err != nil { 64 return fmt.Errorf("open file error, %s", err) 65 } 66 67 reader, err := reader.New(file) 68 if err != nil { 69 return err 70 } 71 s.dbReader = reader 72 return nil 73 } 74 75 func (s *service) loadDBFileFromBucket() error { 76 buf := bytes.NewBuffer([]byte{}) 77 downloadReq := storage.NewDownloadFileRequest(s.bucketName, s.dbFileName, buf) 78 if err := s.storage.Download(downloadReq); err != nil { 79 return err 80 } 81 82 reader, err := reader.New(buf) 83 if err != nil { 84 return err 85 } 86 s.dbReader = reader 87 88 return nil 89 }