github.com/terramate-io/tf@v0.0.0-20230830114523-fce866b4dfcd/backend/remote-state/pg/backend.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package pg
     5  
     6  import (
     7  	"context"
     8  	"database/sql"
     9  	"fmt"
    10  	"os"
    11  	"strconv"
    12  
    13  	"github.com/lib/pq"
    14  	"github.com/terramate-io/tf/backend"
    15  	"github.com/terramate-io/tf/legacy/helper/schema"
    16  )
    17  
    18  const (
    19  	statesTableName = "states"
    20  	statesIndexName = "states_by_name"
    21  )
    22  
    23  func defaultBoolFunc(k string, dv bool) schema.SchemaDefaultFunc {
    24  	return func() (interface{}, error) {
    25  		if v := os.Getenv(k); v != "" {
    26  			return strconv.ParseBool(v)
    27  		}
    28  
    29  		return dv, nil
    30  	}
    31  }
    32  
    33  // New creates a new backend for Postgres remote state.
    34  func New() backend.Backend {
    35  	s := &schema.Backend{
    36  		Schema: map[string]*schema.Schema{
    37  			"conn_str": {
    38  				Type:        schema.TypeString,
    39  				Optional:    true,
    40  				Description: "Postgres connection string; a `postgres://` URL",
    41  				DefaultFunc: schema.EnvDefaultFunc("PG_CONN_STR", nil),
    42  			},
    43  
    44  			"schema_name": {
    45  				Type:        schema.TypeString,
    46  				Optional:    true,
    47  				Description: "Name of the automatically managed Postgres schema to store state",
    48  				DefaultFunc: schema.EnvDefaultFunc("PG_SCHEMA_NAME", "terraform_remote_state"),
    49  			},
    50  
    51  			"skip_schema_creation": {
    52  				Type:        schema.TypeBool,
    53  				Optional:    true,
    54  				Description: "If set to `true`, Terraform won't try to create the Postgres schema",
    55  				DefaultFunc: defaultBoolFunc("PG_SKIP_SCHEMA_CREATION", false),
    56  			},
    57  
    58  			"skip_table_creation": {
    59  				Type:        schema.TypeBool,
    60  				Optional:    true,
    61  				Description: "If set to `true`, Terraform won't try to create the Postgres table",
    62  				DefaultFunc: defaultBoolFunc("PG_SKIP_TABLE_CREATION", false),
    63  			},
    64  
    65  			"skip_index_creation": {
    66  				Type:        schema.TypeBool,
    67  				Optional:    true,
    68  				Description: "If set to `true`, Terraform won't try to create the Postgres index",
    69  				DefaultFunc: defaultBoolFunc("PG_SKIP_INDEX_CREATION", false),
    70  			},
    71  		},
    72  	}
    73  
    74  	result := &Backend{Backend: s}
    75  	result.Backend.ConfigureFunc = result.configure
    76  	return result
    77  }
    78  
    79  type Backend struct {
    80  	*schema.Backend
    81  
    82  	// The fields below are set from configure
    83  	db         *sql.DB
    84  	configData *schema.ResourceData
    85  	connStr    string
    86  	schemaName string
    87  }
    88  
    89  func (b *Backend) configure(ctx context.Context) error {
    90  	// Grab the resource data
    91  	b.configData = schema.FromContextBackendConfig(ctx)
    92  	data := b.configData
    93  
    94  	b.connStr = data.Get("conn_str").(string)
    95  	b.schemaName = pq.QuoteIdentifier(data.Get("schema_name").(string))
    96  
    97  	db, err := sql.Open("postgres", b.connStr)
    98  	if err != nil {
    99  		return err
   100  	}
   101  
   102  	// Prepare database schema, tables, & indexes.
   103  	var query string
   104  
   105  	if !data.Get("skip_schema_creation").(bool) {
   106  		// list all schemas to see if it exists
   107  		var count int
   108  		query = `select count(1) from information_schema.schemata where schema_name = $1`
   109  		if err := db.QueryRow(query, data.Get("schema_name").(string)).Scan(&count); err != nil {
   110  			return err
   111  		}
   112  
   113  		// skip schema creation if schema already exists
   114  		// `CREATE SCHEMA IF NOT EXISTS` is to be avoided if ever
   115  		// a user hasn't been granted the `CREATE SCHEMA` privilege
   116  		if count < 1 {
   117  			// tries to create the schema
   118  			query = `CREATE SCHEMA IF NOT EXISTS %s`
   119  			if _, err := db.Exec(fmt.Sprintf(query, b.schemaName)); err != nil {
   120  				return err
   121  			}
   122  		}
   123  	}
   124  
   125  	if !data.Get("skip_table_creation").(bool) {
   126  		if _, err := db.Exec("CREATE SEQUENCE IF NOT EXISTS public.global_states_id_seq AS bigint"); err != nil {
   127  			return err
   128  		}
   129  
   130  		query = `CREATE TABLE IF NOT EXISTS %s.%s (
   131  			id bigint NOT NULL DEFAULT nextval('public.global_states_id_seq') PRIMARY KEY,
   132  			name text UNIQUE,
   133  			data text
   134  			)`
   135  		if _, err := db.Exec(fmt.Sprintf(query, b.schemaName, statesTableName)); err != nil {
   136  			return err
   137  		}
   138  	}
   139  
   140  	if !data.Get("skip_index_creation").(bool) {
   141  		query = `CREATE UNIQUE INDEX IF NOT EXISTS %s ON %s.%s (name)`
   142  		if _, err := db.Exec(fmt.Sprintf(query, statesIndexName, b.schemaName, statesTableName)); err != nil {
   143  			return err
   144  		}
   145  	}
   146  
   147  	// Assign db after its schema is prepared.
   148  	b.db = db
   149  
   150  	return nil
   151  }