Mirrors/juanfont.headscale
1
0
Fork 0
mirror of https://github.com/juanfont/headscale.git synced 2025-08-24 13:46:53 +02:00
Code Issues Projects Releases Wiki Activity

combine ListNodes and ListNodesSubset into one function

Browse Source
This commit is contained in:
Nils Enkelmann 2025-04-01 08:34:05 +02:00
parent e3d44d3b44
commit e0df114d5a
3 changed files with 24 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
hscontrol/db/node.go
View File

@ -58,41 +58,26 @@ func ListPeers(tx *gorm.DB, nodeID types.NodeID) (types.Nodes, error) {
return nodes, nil return nodes, nil
} }
func (hsdb *HSDatabase) ListNodes() (types.Nodes, error) { func (hsdb *HSDatabase) ListNodes(nodeIDs ...types.NodeIDs) (types.Nodes, error) {
return Read(hsdb.DB, func(rx *gorm.DB) (types.Nodes, error) { return Read(hsdb.DB, func(rx *gorm.DB) (types.Nodes, error) {
return ListNodes(rx) return ListNodes(rx, nodeIDs...)
}) })
} }
func ListNodes(tx *gorm.DB) (types.Nodes, error) { func ListNodes(tx *gorm.DB, nodeIDs ...types.NodeIDs) (types.Nodes, error) {
nodes := types.Nodes{} if len(nodeIDs) > 0 && len(nodeIDs[0]) == 0 {
if err := tx.
Preload("AuthKey").
Preload("AuthKey.User").
Preload("User").
Find(&nodes).Error; err != nil {
return nil, err
}
return nodes, nil
}
func (hsdb *HSDatabase) ListNodesSubset(nodeIDs types.NodeIDs) (types.Nodes, error) {
return Read(hsdb.DB, func(rx *gorm.DB) (types.Nodes, error) {
return ListNodesSubset(rx, nodeIDs)
})
}
func ListNodesSubset(tx *gorm.DB, nodeIDs types.NodeIDs) (types.Nodes, error) {
if len(nodeIDs) < 1 {
return types.Nodes{}, nil return types.Nodes{}, nil
} }
var nodeFilter types.NodeIDs = nil
if len(nodeIDs) > 0 {
nodeFilter = nodeIDs[0]
}
nodes := types.Nodes{} nodes := types.Nodes{}
if err := tx. if err := tx.
Preload("AuthKey"). Preload("AuthKey").
Preload("AuthKey.User"). Preload("AuthKey.User").
Preload("User"). Preload("User").
Where(nodeIDs).Find(&nodes).Error; err != nil { Where(nodeFilter).Find(&nodes).Error; err != nil {
return nil, err return nil, err
} }

17
hscontrol/db/node_test.go
View File

@ -753,7 +753,7 @@ func TestRenameNode(t *testing.T) {
assert.ErrorContains(t, err, "name is not unique") assert.ErrorContains(t, err, "name is not unique")
} }
func TestListNodesSubset(t *testing.T) { func TestListNodes(t *testing.T) {
// Setup test database // Setup test database
db, err := newSQLiteTestDB() db, err := newSQLiteTestDB()
if err != nil { if err != nil {
@ -807,24 +807,31 @@ func TestListNodesSubset(t *testing.T) {
assert.Len(t, nodes, 2) assert.Len(t, nodes, 2)
// No parameter means no filter, should return all nodes
nodes, err = db.ListNodes()
require.NoError(t, err)
assert.Equal(t, len(nodes), 2)
assert.Equal(t, "test1", nodes[0].Hostname)
assert.Equal(t, "test2", nodes[1].Hostname)
// Empty node list should return empty list // Empty node list should return empty list
nodes, err = db.ListNodesSubset(types.NodeIDs{}) nodes, err = db.ListNodes(types.NodeIDs{})
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, len(nodes), 0) assert.Equal(t, len(nodes), 0)
// No match in IDs should return empty list and no error // No match in IDs should return empty list and no error
nodes, err = db.ListNodesSubset(types.NodeIDs{3, 4, 5}) nodes, err = db.ListNodes(types.NodeIDs{3, 4, 5})
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, len(nodes), 0) assert.Equal(t, len(nodes), 0)
// Partial match in IDs // Partial match in IDs
nodes, err = db.ListNodesSubset(types.NodeIDs{2, 3}) nodes, err = db.ListNodes(types.NodeIDs{2, 3})
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, len(nodes), 1) assert.Equal(t, len(nodes), 1)
assert.Equal(t, "test2", nodes[0].Hostname) assert.Equal(t, "test2", nodes[0].Hostname)
// Several matched IDs // Several matched IDs
nodes, err = db.ListNodesSubset(types.NodeIDs{1, 2, 3}) nodes, err = db.ListNodes(types.NodeIDs{1, 2, 3})
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, len(nodes), 2) assert.Equal(t, len(nodes), 2)
assert.Equal(t, "test1", nodes[0].Hostname) assert.Equal(t, "test1", nodes[0].Hostname)

6
hscontrol/mapper/mapper.go
View File

@ -269,7 +269,7 @@ func (m *Mapper) PeerChangedResponse(
} }
} }
changedNodes, err := m.ListNodesSubset(changedIDs) changedNodes, err := m.ListNodes(changedIDs)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -491,8 +491,8 @@ func (m *Mapper) ListPeers(nodeID types.NodeID) (types.Nodes, error) {
return peers, nil return peers, nil
} }
func (m *Mapper) ListNodesSubset(nodeIDs []types.NodeID) (types.Nodes, error) { func (m *Mapper) ListNodes(nodeIDs ...types.NodeIDs) (types.Nodes, error) {
nodes, err := m.db.ListNodesSubset(nodeIDs) nodes, err := m.db.ListNodes(nodeIDs...)
if err != nil { if err != nil {
return nil, err return nil, err
} }