github.com/omnigres/cli@v0.1.4/cmd/assemble.go (about)

     1  package cmd
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"database/sql/driver"
     7  	"encoding/json"
     8  	"fmt"
     9  	"maps"
    10  	"os"
    11  	"path"
    12  	"strings"
    13  
    14  	"github.com/charmbracelet/log"
    15  	"github.com/lib/pq"
    16  	"github.com/omnigres/cli/orb"
    17  	"github.com/spf13/cobra"
    18  )
    19  
    20  var assembleCmd = &cobra.Command{
    21  	Use:   "assemble",
    22  	Short: "Assemble or reassemble orb",
    23  	Long:  `By default, will assemble all listed orbs`,
    24  	Run: func(cmd *cobra.Command, args []string) {
    25  		var cluster orb.OrbCluster
    26  		var err error
    27  		cluster, err = getOrbCluster()
    28  		if err != nil {
    29  			log.Fatal(err)
    30  		}
    31  
    32  		log.Debug("Workspace", "workspace", workspace)
    33  		cwd, err := os.Getwd()
    34  		if err != nil {
    35  			log.Fatal(err)
    36  		}
    37  
    38  		orbs, err := currentOrbs(cluster, cwd)
    39  		if err != nil {
    40  			log.Fatal(err)
    41  		}
    42  
    43  		ctx := context.Background()
    44  		log.Debug("Capturing orbs", "orbs", orbs)
    45  		assembleOrbs(
    46  			ctx,
    47  			cluster,
    48  			dbReset,
    49  			orbs,
    50  			func(orbName string) string { return orbName },
    51  		)
    52  	},
    53  }
    54  
    55  func assembleOrbs(
    56  	ctx context.Context,
    57  	cluster orb.OrbCluster,
    58  	dbReset bool,
    59  	orbs []string,
    60  	databaseForOrb func(string) string,
    61  ) (err error) {
    62  	var db *sql.DB
    63  	db, err = cluster.Connect(ctx, "omnigres")
    64  	if err != nil {
    65      log.Error("Could not connect to orb. Ensure the docker container is running, perhaps 'omnigres start' will fix it.")
    66  		return
    67  	}
    68  	for _, orbName := range orbs {
    69  		log.Infof("Assembling orb %s", orbName)
    70  		dbName := databaseForOrb(orbName)
    71  		var dbExists bool
    72  		err := db.QueryRowContext(
    73  			ctx,
    74  			`select exists(select from pg_database where datname = $1)`,
    75  			dbName,
    76  		).Scan(&dbExists)
    77  		if err != nil {
    78  			log.Fatal(err)
    79  		}
    80  
    81  		if dbReset && dbExists {
    82  			if err != sql.ErrNoRows {
    83  				_, err = db.ExecContext(ctx, fmt.Sprintf(`drop database %q`, dbName))
    84  				if err != nil {
    85  					log.Fatal(err)
    86  				}
    87  			}
    88  		}
    89  
    90  		if dbReset || !dbExists {
    91  			_, err = db.ExecContext(ctx, fmt.Sprintf(`create database %q`, dbName))
    92  			if err != nil {
    93  				log.Fatal(err)
    94  			}
    95  		}
    96  
    97  		orbSource := path.Join(orbName, "src")
    98  		assembleSchema(ctx, db, orbSource, dbName)
    99  	}
   100  	return
   101  }
   102  
   103  func assembleSchema(ctx context.Context, db *sql.DB, orbSource string, dbName string) {
   104  	logger := log.New(os.Stdout)
   105  	logger.SetReportTimestamp(true)
   106  
   107  	logger.SetPrefix(fmt.Sprintf("[%s] ", dbName))
   108  
   109  	levels := make(map[string]log.Level)
   110  	levels["info"] = log.InfoLevel
   111  	levels["error"] = log.ErrorLevel
   112  	jsonMessageReporting := func(notice *pq.Error) {
   113  		var message map[string]interface{}
   114  		err := json.NewDecoder(strings.NewReader(notice.Message)).Decode(&message)
   115  		if err != nil {
   116  			logger.Errorf("Message is not valid JSON: %s", notice.Message)
   117  			return
   118  		}
   119  
   120  		mapToArray := func(m map[string]interface{}) []interface{} {
   121  			result := make([]interface{}, 0, len(m)*2) // Allocate slice with enough capacity
   122  			for key, value := range m {
   123  				result = append(result, key, value)
   124  			}
   125  			return result
   126  		}
   127  		strippedMessage := maps.Clone(message)
   128  		delete(strippedMessage, "type")
   129  		delete(strippedMessage, "message")
   130  		tags := mapToArray(strippedMessage)
   131  		logger.Log(levels[message["type"].(string)], message["message"], tags...)
   132  	}
   133  
   134  	conn, err := db.Conn(ctx)
   135  	if err != nil {
   136  		log.Fatal(err)
   137  	}
   138  	defer conn.Close()
   139  	conn.Raw(func(driverConn any) error {
   140  		pq.SetNoticeHandler(driverConn.(driver.Conn), jsonMessageReporting)
   141  		return nil
   142  	})
   143  
   144  	rows, err := conn.QueryContext(ctx,
   145  		`select migration_filename, migration_statement, execution_error from omni_schema.assemble_schema($1, omni_vfs.local_fs('/mnt/host'), $2) where execution_error is not null`,
   146  		fmt.Sprintf("dbname=%s user=omnigres", dbName), orbSource)
   147  
   148  	if err != nil {
   149  		log.Fatal(err)
   150  	}
   151  	defer rows.Close()
   152  
   153  	for rows.Next() {
   154  		var migration_filename, migration_statement, execution_error sql.NullString
   155  		err = rows.Scan(&migration_filename, &migration_statement, &execution_error)
   156  		if err != nil {
   157  			log.Error(err)
   158  			return
   159  		}
   160  	}
   161  }
   162  
   163  var dbReset bool
   164  
   165  func init() {
   166  	rootCmd.AddCommand(assembleCmd)
   167  	assembleCmd.Flags().BoolVarP(&dbReset, "dbReset", "r", false, "dbReset")
   168  }