Dragonfly2
77 строк · 1.8 Кб
1/*
2* Copyright 2022 The Dragonfly Authors
3*
4* Licensed under the Apache License, Version 2.0 (the "License");
5* you may not use this file except in compliance with the License.
6* You may obtain a copy of the License at
7*
8* http://www.apache.org/licenses/LICENSE-2.0
9*
10* Unless required by applicable law or agreed to in writing, software
11* distributed under the License is distributed on an "AS IS" BASIS,
12* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13* See the License for the specific language governing permissions and
14* limitations under the License.
15*/
16
17package database18
19import (20"fmt"21
22"gorm.io/driver/postgres"23"gorm.io/gorm"24"gorm.io/gorm/schema"25"moul.io/zapgorm2"26
27logger "d7y.io/dragonfly/v2/internal/dflog"28"d7y.io/dragonfly/v2/manager/config"29)
30
31func newPostgres(cfg *config.Config) (*gorm.DB, error) {32postgresCfg := &cfg.Database.Postgres33
34// Format dsn string.35dsn := formatPostgresDSN(postgresCfg)36
37// Connect to postgres.38db, err := gorm.Open(postgres.New(postgres.Config{39DSN: dsn,40PreferSimpleProtocol: cfg.Database.Postgres.PreferSimpleProtocol,41}), &gorm.Config{42NamingStrategy: schema.NamingStrategy{43SingularTable: true,44},45DisableForeignKeyConstraintWhenMigrating: true,46Logger: zapgorm2.New(logger.CoreLogger.Desugar()),47})48if err != nil {49return nil, err50}51
52// Run migration.53if postgresCfg.Migrate {54if err := migrate(db); err != nil {55return nil, err56}57}58
59// Run seed.60if err := seed(db); err != nil {61return nil, err62}63
64return db, nil65}
66
67func formatPostgresDSN(cfg *config.PostgresConfig) string {68return fmt.Sprintf("host=%v user=%v password=%v dbname=%v port=%v sslmode=%v TimeZone=%v",69cfg.Host,70cfg.User,71cfg.Password,72cfg.DBName,73cfg.Port,74cfg.SSLMode,75cfg.Timezone,76)77}
78