-
-
Notifications
You must be signed in to change notification settings - Fork 2.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
5 changed files
with
227 additions
and
6 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,107 @@ | ||
package explorer | ||
|
||
// A simple JSON database for storing and retrieving p2p network tokens and a name and description. | ||
|
||
import ( | ||
"encoding/json" | ||
"os" | ||
"sort" | ||
"sync" | ||
) | ||
|
||
// Database is a simple JSON database for storing and retrieving p2p network tokens and a name and description. | ||
type Database struct { | ||
sync.RWMutex | ||
path string | ||
data map[string]TokenData | ||
} | ||
|
||
// TokenData is a p2p network token with a name and description. | ||
type TokenData struct { | ||
Name string `json:"name"` | ||
Description string `json:"description"` | ||
} | ||
|
||
// NewDatabase creates a new Database with the given path. | ||
func NewDatabase(path string) (*Database, error) { | ||
db := &Database{ | ||
data: make(map[string]TokenData), | ||
path: path, | ||
} | ||
return db, db.load() | ||
} | ||
|
||
// Get retrieves a Token from the Database by its token. | ||
func (db *Database) Get(token string) (TokenData, bool) { | ||
db.RLock() | ||
defer db.RUnlock() | ||
t, ok := db.data[token] | ||
return t, ok | ||
} | ||
|
||
// Set stores a Token in the Database by its token. | ||
func (db *Database) Set(token string, t TokenData) error { | ||
db.Lock() | ||
db.data[token] = t | ||
db.Unlock() | ||
|
||
return db.save() | ||
} | ||
|
||
// Delete removes a Token from the Database by its token. | ||
func (db *Database) Delete(token string) error { | ||
db.Lock() | ||
delete(db.data, token) | ||
db.Unlock() | ||
return db.save() | ||
} | ||
|
||
func (db *Database) TokenList() []string { | ||
db.RLock() | ||
defer db.RUnlock() | ||
tokens := []string{} | ||
for k := range db.data { | ||
tokens = append(tokens, k) | ||
} | ||
|
||
sort.Slice(tokens, func(i, j int) bool { | ||
// sort by token | ||
return tokens[i] < tokens[j] | ||
}) | ||
|
||
return tokens | ||
} | ||
|
||
// load reads the Database from disk. | ||
func (db *Database) load() error { | ||
db.Lock() | ||
defer db.Unlock() | ||
|
||
if _, err := os.Stat(db.path); os.IsNotExist(err) { | ||
|
||
return nil | ||
} | ||
|
||
// Read the file from disk | ||
// Unmarshal the JSON into db.data | ||
f, err := os.ReadFile(db.path) | ||
if err != nil { | ||
return err | ||
} | ||
return json.Unmarshal(f, &db.data) | ||
} | ||
|
||
// save writes the Database to disk. | ||
func (db *Database) save() error { | ||
db.RLock() | ||
defer db.RUnlock() | ||
|
||
// Marshal db.data into JSON | ||
// Write the JSON to the file | ||
f, err := os.Create(db.path) | ||
if err != nil { | ||
return err | ||
} | ||
defer f.Close() | ||
return json.NewEncoder(f).Encode(db.data) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,92 @@ | ||
package explorer_test | ||
|
||
import ( | ||
"os" | ||
|
||
. "github.com/onsi/ginkgo/v2" | ||
. "github.com/onsi/gomega" | ||
|
||
"github.com/mudler/LocalAI/core/explorer" | ||
) | ||
|
||
var _ = Describe("Database", func() { | ||
var ( | ||
dbPath string | ||
db *explorer.Database | ||
err error | ||
) | ||
|
||
BeforeEach(func() { | ||
// Create a temporary file path for the database | ||
dbPath = "test_db.json" | ||
db, err = explorer.NewDatabase(dbPath) | ||
Expect(err).To(BeNil()) | ||
}) | ||
|
||
AfterEach(func() { | ||
// Clean up the temporary database file | ||
os.Remove(dbPath) | ||
}) | ||
|
||
Context("when managing tokens", func() { | ||
It("should add and retrieve a token", func() { | ||
token := "token123" | ||
t := explorer.TokenData{Name: "TokenName", Description: "A test token"} | ||
|
||
err = db.Set(token, t) | ||
Expect(err).To(BeNil()) | ||
|
||
retrievedToken, exists := db.Get(token) | ||
Expect(exists).To(BeTrue()) | ||
Expect(retrievedToken).To(Equal(t)) | ||
}) | ||
|
||
It("should delete a token", func() { | ||
token := "token123" | ||
t := explorer.TokenData{Name: "TokenName", Description: "A test token"} | ||
|
||
err = db.Set(token, t) | ||
Expect(err).To(BeNil()) | ||
|
||
err = db.Delete(token) | ||
Expect(err).To(BeNil()) | ||
|
||
_, exists := db.Get(token) | ||
Expect(exists).To(BeFalse()) | ||
}) | ||
|
||
It("should persist data to disk", func() { | ||
token := "token123" | ||
t := explorer.TokenData{Name: "TokenName", Description: "A test token"} | ||
|
||
err = db.Set(token, t) | ||
Expect(err).To(BeNil()) | ||
|
||
// Recreate the database object to simulate reloading from disk | ||
db, err = explorer.NewDatabase(dbPath) | ||
Expect(err).To(BeNil()) | ||
|
||
retrievedToken, exists := db.Get(token) | ||
Expect(exists).To(BeTrue()) | ||
Expect(retrievedToken).To(Equal(t)) | ||
|
||
// Check the token list | ||
tokenList := db.TokenList() | ||
Expect(tokenList).To(ContainElement(token)) | ||
}) | ||
}) | ||
|
||
Context("when loading an empty or non-existent file", func() { | ||
It("should start with an empty database", func() { | ||
dbPath = "empty_db.json" | ||
db, err = explorer.NewDatabase(dbPath) | ||
Expect(err).To(BeNil()) | ||
|
||
_, exists := db.Get("nonexistent") | ||
Expect(exists).To(BeFalse()) | ||
|
||
// Clean up | ||
os.Remove(dbPath) | ||
}) | ||
}) | ||
}) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,13 @@ | ||
package explorer_test | ||
|
||
import ( | ||
"testing" | ||
|
||
. "github.com/onsi/ginkgo/v2" | ||
. "github.com/onsi/gomega" | ||
) | ||
|
||
func TestExplorer(t *testing.T) { | ||
RegisterFailHandler(Fail) | ||
RunSpecs(t, "Explorer test suite") | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters