Skip to content

Commit f59d44b

Browse files
heiytorheiytor
committed
wip
1 parent af04767 commit f59d44b

27 files changed

+1267
-6
lines changed

‎api/server.go‎

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,17 @@ package main
22

33
import (
44
"context"
5+
"errors"
56
"os"
67

78
"github.com/getsentry/sentry-go"
89
"github.com/labstack/echo/v4"
910
"github.com/shellhub-io/shellhub/api/routes"
1011
"github.com/shellhub-io/shellhub/api/services"
12+
"github.com/shellhub-io/shellhub/api/store"
1113
"github.com/shellhub-io/shellhub/api/store/mongo"
1214
"github.com/shellhub-io/shellhub/api/store/mongo/options"
15+
"github.com/shellhub-io/shellhub/api/store/pg"
1316
"github.com/shellhub-io/shellhub/pkg/api/internalclient"
1417
"github.com/shellhub-io/shellhub/pkg/cache"
1518
"github.com/shellhub-io/shellhub/pkg/geoip/geolite2"
@@ -19,9 +22,22 @@ import (
1922
)
2023

2124
type env struct {
25+
Database string `env:"DATABASE,default=mongodb"`
26+
2227
// MongoURI specifies the connection string for MongoDB.
2328
MongoURI string `env:"MONGO_URI,default=mongodb://mongo:27017/main"`
2429

30+
// PostgresHost specifies the host for PostgreSQL.
31+
PostgresHost string `env:"POSTGRES_HOST,default=postgres"`
32+
// PostgresPort specifies the port for PostgreSQL.
33+
PostgresPort string `env:"POSTGRES_PORT,default=5432"`
34+
// PostgresUser specifies the username for authenticate PostgreSQL.
35+
PostgresUser string `env:"POSTGRES_USER,default=admin"`
36+
// PostgresUser specifies the password for authenticate PostgreSQL.
37+
PostgresPassword string `env:"POSTGRES_PASSWORD,default=admin"`
38+
// PostgresDB especifica o nome do banco de dados PostgreSQL a ser utilizado.
39+
PostgresDB string `env:"POSTGRES_DB,default=main"`
40+
2541
// RedisURI specifies the connection string for Redis.
2642
RedisURI string `env:"REDIS_URI,default=redis://redis:6379"`
2743
// RedisCachePoolSize defines the maximum number of concurrent connections to Redis cache.
@@ -72,14 +88,22 @@ func (s *Server) Setup(ctx context.Context) error {
7288

7389
log.Debug("Redis cache initialized successfully")
7490

75-
store, err := mongo.NewStore(ctx, s.env.MongoURI, cache, options.RunMigatrions)
76-
if err != nil {
77-
log.
78-
WithError(err).
79-
Fatal("failed to create the store")
91+
var store store.Store
92+
switch s.env.Database {
93+
case "mongodb":
94+
store, err = mongo.NewStore(ctx, s.env.MongoURI, cache, options.RunMigatrions)
95+
case "postgres":
96+
uri := pg.URI(s.env.PostgresHost, s.env.PostgresPort, s.env.PostgresUser, s.env.PostgresPassword, s.env.PostgresDB)
97+
store, err = pg.New(ctx, uri)
98+
default:
99+
log.WithField("database", s.env.Database).Error("invalid database")
100+
return errors.New("invalid database")
80101
}
81102

82-
log.Debug("MongoDB store connected successfully")
103+
if err != nil {
104+
log.WithError(err).Error("failed to create the store")
105+
return err
106+
}
83107

84108
apiClient, err := internalclient.NewClient(internalclient.WithAsynqWorker(s.env.RedisURI))
85109
if err != nil {

‎api/store/pg/api-key.go‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package pg
2+
3+
import (
4+
"context"
5+
6+
"github.com/shellhub-io/shellhub/api/store"
7+
"github.com/shellhub-io/shellhub/api/store/pg/entity"
8+
"github.com/shellhub-io/shellhub/pkg/clock"
9+
"github.com/shellhub-io/shellhub/pkg/models"
10+
)
11+
12+
func (pg *Pg) APIKeyCreate(ctx context.Context, apiKey *models.APIKey) (string, error) {
13+
apiKey.CreatedAt = clock.Now()
14+
apiKey.UpdatedAt = clock.Now()
15+
16+
if _, err := pg.driver.NewInsert().Model(entity.APIKeyFromModel(apiKey)).Exec(ctx); err != nil {
17+
return "", fromSqlError(err)
18+
}
19+
20+
return apiKey.ID, nil
21+
}
22+
23+
func (pg *Pg) APIKeyConflicts(ctx context.Context, tenantID string, target *models.APIKeyConflicts) ([]string, bool, error) {
24+
apiKeys := make([]map[string]any, 0)
25+
if err := pg.driver.NewSelect().Model((*entity.Namespace)(nil)).Column("name").Where("name = ?", target.Name).Scan(ctx, &apiKeys); err != nil {
26+
return nil, false, fromSqlError(err)
27+
}
28+
29+
conflicts := make([]string, 0)
30+
for _, apiKey := range apiKeys {
31+
if apiKey["name"] == target.Name {
32+
conflicts = append(conflicts, "name")
33+
}
34+
}
35+
36+
return conflicts, len(conflicts) > 0, nil
37+
}
38+
39+
func (pg *Pg) APIKeyList(ctx context.Context, opts ...store.QueryOption) ([]models.APIKey, int, error) {
40+
entities := make([]entity.APIKey, 0)
41+
42+
query := pg.driver.NewSelect().Model(&entities)
43+
if err := applyOptions(ctx, query, opts...); err != nil {
44+
return nil, 0, fromSqlError(err)
45+
}
46+
47+
count, err := query.ScanAndCount(ctx)
48+
if err != nil {
49+
return nil, 0, fromSqlError(err)
50+
}
51+
52+
apiKeys := make([]models.APIKey, len(entities))
53+
for i, e := range entities {
54+
apiKeys[i] = *entity.APIKeyToModel(&e)
55+
}
56+
57+
return apiKeys, count, nil
58+
}
59+
60+
func (pg *Pg) APIKeyGet(ctx context.Context, id string) (*models.APIKey, error) {
61+
a := new(entity.APIKey)
62+
if err := pg.driver.NewSelect().Model(a).Where("id = ?", id).Scan(ctx); err != nil {
63+
return nil, fromSqlError(err)
64+
}
65+
66+
return entity.APIKeyToModel(a), nil
67+
}
68+
69+
func (pg *Pg) APIKeyGetByName(ctx context.Context, tenantID string, name string) (*models.APIKey, error) {
70+
a := new(entity.APIKey)
71+
if err := pg.driver.NewSelect().Model(a).Where("namespace_id = ?", tenantID).Where("name = ?", name).Scan(ctx); err != nil {
72+
return nil, fromSqlError(err)
73+
}
74+
75+
return entity.APIKeyToModel(a), nil
76+
}
77+
78+
func (pg *Pg) APIKeySave(ctx context.Context, apiKey *models.APIKey) error {
79+
a := entity.APIKeyFromModel(apiKey)
80+
a.UpdatedAt = clock.Now()
81+
_, err := pg.driver.NewUpdate().Model(a).WherePK().Exec(ctx)
82+
83+
return fromSqlError(err)
84+
}
85+
86+
func (pg *Pg) APIKeyDelete(ctx context.Context, apiKey *models.APIKey) error {
87+
a := entity.APIKeyFromModel(apiKey)
88+
_, err := pg.driver.NewDelete().Model(a).WherePK().Exec(ctx)
89+
90+
return fromSqlError(err)
91+
}

‎api/store/pg/api-key_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
package pg_test

‎api/store/pg/dbtest/dbtest.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package dbtest
2+
3+
import (
4+
"context"
5+
"io"
6+
"log"
7+
"time"
8+
9+
"github.com/testcontainers/testcontainers-go"
10+
"github.com/testcontainers/testcontainers-go/modules/postgres"
11+
"github.com/testcontainers/testcontainers-go/wait"
12+
)
13+
14+
// Server represents a Postgres test server instance.
15+
type Server struct {
16+
container *postgres.PostgresContainer
17+
}
18+
19+
// Up starts a new Postgres container. Use [Server.ConnectionString] to access the connection string.
20+
func (srv *Server) Up(ctx context.Context, verbose bool) error {
21+
if !verbose {
22+
testcontainers.Logger = log.New(io.Discard, "", 0)
23+
}
24+
25+
opts := []testcontainers.ContainerCustomizer{
26+
postgres.WithDatabase("test"),
27+
postgres.WithUsername("admin"),
28+
postgres.WithPassword("admin"),
29+
testcontainers.WithWaitStrategy(wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5 * time.Second)),
30+
}
31+
32+
container, err := postgres.Run(ctx, "postgres:17", opts...)
33+
if err != nil {
34+
return err
35+
}
36+
37+
srv.container = container
38+
39+
return nil
40+
}
41+
42+
// Down gracefully terminates the Postgres container.
43+
func (srv *Server) Down(ctx context.Context) error {
44+
return srv.container.Terminate(ctx)
45+
}
46+
47+
func (srv *Server) ConnectionString(ctx context.Context) (string, error) {
48+
cIP, err := srv.container.ContainerIP(ctx)
49+
if err != nil {
50+
return "", err
51+
}
52+
53+
return "postgres://admin:admin@" + cIP + ":5432/test", nil
54+
}

‎api/store/pg/dbtest/fixtures.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
package dbtest
2+
3+
import (
4+
"path/filepath"
5+
"runtime"
6+
)
7+
8+
func FixturesPath() string {
9+
_, file, _, _ := runtime.Caller(0)
10+
11+
return filepath.Join(filepath.Dir(file), "fixtures")
12+
}

‎api/store/pg/dbtest/fixtures/users.yaml‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
- model: User
2+
rows:
3+
- id: 0195cefa-aa01-7efb-8098-c9c173056250
4+
created_at: 2025-01-15T10:30:00+00:00
5+
updated_at: 2025-01-15T10:30:00+00:00
6+
last_login: null
7+
status: confirmed
8+
origin: local
9+
external_id: ""
10+
name: Jonh Doe
11+
username: john_doe
12+
email: [email protected]
13+
security_email: [email protected]
14+
password_digest: "$2y$12$VVm2ETx7AvaGlfMYqNYK9uzU2M45YZ70YnT..O.s1o2zdE1pekhq6"
15+
auth_methods: [ local ]
16+
namespace_ownership_limit: -1
17+
email_marketing: true
18+
preferred_namespace_id: null

‎api/store/pg/entity/api-key.go‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
package entity
2+
3+
import (
4+
"time"
5+
6+
"github.com/shellhub-io/shellhub/pkg/api/authorizer"
7+
"github.com/shellhub-io/shellhub/pkg/models"
8+
"github.com/uptrace/bun"
9+
)
10+
11+
type APIKey struct {
12+
bun.BaseModel `bun:"table:api_keys"`
13+
14+
KeyDigest string `bun:"key_digest,pk"`
15+
NamespaceID string `bun:"namespace_id,pk"`
16+
Name string `bun:"name"`
17+
Role string `bun:"role"`
18+
UserID string `bun:"user_id"`
19+
CreatedAt time.Time `bun:"created_at"`
20+
UpdatedAt time.Time `bun:"updated_at"`
21+
ExpiresIn int64 `bun:"expires_in,nullzero"`
22+
}
23+
24+
func APIKeyFromModel(model *models.APIKey) *APIKey {
25+
return &APIKey{
26+
Name: model.Name,
27+
NamespaceID: model.TenantID,
28+
KeyDigest: model.ID,
29+
Role: model.Role.String(),
30+
UserID: model.CreatedBy,
31+
CreatedAt: model.CreatedAt,
32+
UpdatedAt: model.UpdatedAt,
33+
ExpiresIn: model.ExpiresIn,
34+
}
35+
}
36+
37+
func APIKeyToModel(entity *APIKey) *models.APIKey {
38+
return &models.APIKey{
39+
ID: entity.KeyDigest,
40+
Name: entity.Name,
41+
TenantID: entity.NamespaceID,
42+
Role: authorizer.Role(entity.Role),
43+
CreatedBy: entity.UserID,
44+
CreatedAt: entity.CreatedAt,
45+
UpdatedAt: entity.UpdatedAt,
46+
ExpiresIn: entity.ExpiresIn,
47+
}
48+
}

‎api/store/pg/entity/entity.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
package entity
2+
3+
func Entities() []any {
4+
return []any{
5+
(*APIKey)(nil),
6+
(*Membership)(nil),
7+
(*Namespace)(nil),
8+
(*User)(nil),
9+
}
10+
}

‎api/store/pg/entity/membership.go‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package entity
2+
3+
import (
4+
"time"
5+
6+
"github.com/shellhub-io/shellhub/pkg/api/authorizer"
7+
"github.com/shellhub-io/shellhub/pkg/models"
8+
"github.com/uptrace/bun"
9+
)
10+
11+
type Membership struct {
12+
bun.BaseModel `bun:"table:memberships"`
13+
14+
UserID string `bun:"user_id,pk,type:uuid"`
15+
NamespaceID string `bun:"namespace_id,pk,type:uuid"`
16+
CreatedAt time.Time `bun:"created_at"`
17+
UpdatedAt time.Time `bun:"updated_at"`
18+
Status string `bun:"status"`
19+
Role string `bun:"role"`
20+
21+
User *User `bun:"rel:belongs-to,join:user_id=id"`
22+
Namespace *Namespace `bun:"rel:belongs-to,join:namespace_id=id"`
23+
}
24+
25+
func MembershipFromModel(namespaceID string, member *models.Member) *Membership {
26+
return &Membership{
27+
UserID: member.ID,
28+
NamespaceID: namespaceID,
29+
CreatedAt: member.AddedAt,
30+
UpdatedAt: time.Time{},
31+
Status: string(member.Status),
32+
Role: string(member.Role),
33+
}
34+
}
35+
36+
func MembershipToModel(entity *Membership) *models.Member {
37+
return &models.Member{
38+
ID: entity.UserID,
39+
AddedAt: entity.CreatedAt,
40+
Role: authorizer.Role(entity.Role),
41+
Status: models.MemberStatus(entity.Status),
42+
Email: entity.User.Email,
43+
}
44+
}

0 commit comments

Comments
 (0)