@@ -21,9 +21,14 @@ const (
21
21
RDSSuperUserRole = "rds_superuser"
22
22
)
23
23
24
- var extensions = []string {"citext" , "uuid-ossp" ,
24
+ var defaultExtensions = []string {"citext" , "uuid-ossp" ,
25
25
"pgcrypto" , "hstore" , "pg_stat_statements" ,
26
- "plpgsql" , "pg_partman" , "hll" , "pg_cron" }
26
+ "plpgsql" , "hll" }
27
+
28
+ var specialExtensionsMap = map [string ]func (* client , string , string ) error {
29
+ "pg_partman" : (* client ).pg_partman ,
30
+ "pg_cron" : (* client ).pg_cron ,
31
+ }
27
32
28
33
type client struct {
29
34
dbType string
@@ -110,8 +115,9 @@ func (pc *client) CreateDataBase(name string) (bool, error) {
110
115
return pc .CreateDatabase (name )
111
116
}
112
117
118
+ // unit test override
113
119
var getDefaulExtensions = func () []string {
114
- return extensions
120
+ return defaultExtensions
115
121
}
116
122
117
123
func (pc * client ) CreateDatabase (dbName string ) (bool , error ) {
@@ -157,21 +163,77 @@ func (pc *client) CreateDefaultExtentions(dbName string) error {
157
163
}
158
164
pc .log .Info ("created extension " + s )
159
165
}
160
- var exists bool
161
- err = db .QueryRow ("SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = $1)" , "pg_cron" ).Scan (& exists )
166
+
167
+ return nil
168
+ }
169
+
170
+ func (pc * client ) CreateSpecialExtentions (dbName string , role string ) error {
171
+ db , err := pc .getDB (dbName )
162
172
if err != nil {
173
+ pc .log .Error (err , "could not connect to db" , "database" , dbName )
163
174
return err
164
175
}
165
- if exists {
166
- // pg_cron extension is enabled - grant permission
167
- _ , err = db . Exec ( "GRANT USAGE ON SCHEMA cron TO public" )
168
- if err != nil {
169
- pc .log .Error (err , "could not GRANT USAGE ON SCHEMA cron TO public" )
176
+ pc . log . Info ( "connected to " + dbName )
177
+ defer db . Close ()
178
+ for functionName := range specialExtensionsMap {
179
+ if err := specialExtensionsMap [ functionName ]( pc , dbName , role ); err != nil {
180
+ pc .log .Error (err , "error creating extention %s" , functionName )
170
181
return err
171
182
}
172
183
}
173
184
return nil
174
185
}
186
+ func (pc * client ) pg_cron (dbName string , role string ) error {
187
+ // create extension pg_cron and grant usage to public
188
+ db , err := pc .getDB (dbName )
189
+ if err != nil {
190
+ pc .log .Error (err , "Failed to connect to database" , "database" , dbName )
191
+ return err
192
+ }
193
+ defer db .Close ()
194
+
195
+ pc .log .Info ("Connected to database" , "database" , dbName )
196
+
197
+ _ , err = db .Exec (fmt .Sprintf (`
198
+ CREATE EXTENSION IF NOT EXISTS pg_cron;
199
+ GRANT USAGE ON SCHEMA cron TO %s;
200
+ ` , pq .QuoteIdentifier (role )))
201
+ if err != nil {
202
+ pc .log .Error (err , "Failed to create extension pg_cron and grant usage on schema cron to public" , "database" , dbName )
203
+ return fmt .Errorf ("failed to create extension pg_cron and grant usage on schema cron to public: %w" , err )
204
+ }
205
+
206
+ pc .log .Info ("Created pg_cron extension and granted usage on schema cron" , "role" , role )
207
+
208
+ return nil
209
+ }
210
+ func (pc * client ) pg_partman (dbName string , role string ) error {
211
+
212
+ db , err := pc .getDB (dbName )
213
+ if err != nil {
214
+ pc .log .Error (err , "could not connect to db" , "database" , dbName )
215
+ return err
216
+ }
217
+ defer db .Close ()
218
+ createPartman := fmt .Sprintf (`
219
+ CREATE SCHEMA IF NOT EXISTS partman;
220
+ CREATE EXTENSION IF NOT EXISTS pg_partman WITH SCHEMA partman;
221
+ GRANT ALL ON SCHEMA partman TO %s;
222
+ GRANT ALL ON ALL TABLES IN SCHEMA partman TO %s;
223
+ GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA partman TO %s;
224
+ GRANT EXECUTE ON ALL PROCEDURES IN SCHEMA partman TO %s;
225
+ ` , pq .QuoteIdentifier (role ), pq .QuoteIdentifier (role ),
226
+ pq .QuoteIdentifier (role ), pq .QuoteIdentifier (role ))
227
+
228
+ _ , err = db .Exec (createPartman )
229
+ if err != nil {
230
+ pc .log .Error (err , "could not create extension pg_partman" )
231
+ return err
232
+ }
233
+ pc .log .Info ("Created pg_partmann extension and granted usage on schema partman" , "role" , role )
234
+
235
+ return nil
236
+ }
175
237
176
238
func (pc * client ) ManageSystemFunctions (dbName string , functions map [string ]string ) error {
177
239
db , err := pc .getDB (dbName )
0 commit comments