|
1 |
| -#include <stdlib.h> |
2 |
| - |
3 | 1 | #include "postgres.h"
|
4 | 2 | #include "fmgr.h"
|
5 | 3 |
|
6 | 4 | #include "libpq/auth.h"
|
7 |
| -#include <unistd.h> |
| 5 | +#include "executor/spi.h" |
| 6 | +#include "miscadmin.h" |
8 | 7 |
|
9 | 8 | PG_MODULE_MAGIC;
|
10 | 9 |
|
11 | 10 | static ClientAuthentication_hook_type original_client_auth_hook = NULL;
|
12 | 11 |
|
| 12 | +static void execute_command(const char *cmd) { |
| 13 | + int result = system(cmd); |
| 14 | + |
| 15 | + if (result != 0) { |
| 16 | + ereport(ERROR, ( |
| 17 | + errmsg("command failed"), |
| 18 | + errdetail("Command: %s\nExit code: %d", cmd, WEXITSTATUS(result)) |
| 19 | + ) |
| 20 | + ); |
| 21 | + } |
| 22 | +} |
| 23 | + |
13 | 24 | static void ensure_role_and_database_exists(Port *port, int status) {
|
14 |
| - char *cmd; |
15 |
| - char *postgres_user; |
16 |
| - char *role_attributes; |
| 25 | + char *cmd = NULL; |
| 26 | + const char *postgres_user = getenv("POSTGRES_USER"); |
| 27 | + const char *role_attributes = getenv("POSTGRES_ROLE_ATTRIBUTES"); |
17 | 28 |
|
18 | 29 | if (original_client_auth_hook) {
|
19 | 30 | original_client_auth_hook(port, status);
|
20 | 31 | }
|
21 | 32 |
|
22 |
| - postgres_user = getenv("POSTGRES_USER"); |
| 33 | + if (!postgres_user) { |
| 34 | + ereport(ERROR, errmsg("POSTGRES_USER environment variable is not set.")); |
| 35 | + } |
23 | 36 |
|
24 |
| - fprintf(stderr, "handling connection for username '%s' to database '%s'\n", port->user_name, port->database_name); |
| 37 | + if (!role_attributes) { |
| 38 | + ereport(ERROR, errmsg("POSTGRES_ROLE_ATTRIBUTES environment variable is not set.")); |
| 39 | + } |
25 | 40 |
|
26 | 41 | // don't infinitely recurse when connecting as superuser
|
27 |
| - if (strcmp(port->database_name, postgres_user) == 0 && strcmp(port->user_name, postgres_user) == 0) { |
| 42 | + if (strcmp(port->user_name, postgres_user) == 0 && strcmp(port->database_name, postgres_user) == 0) { |
28 | 43 | return;
|
29 | 44 | }
|
30 | 45 |
|
31 |
| - role_attributes = getenv("POSTGRES_ROLE_ATTRIBUTES"); |
| 46 | + elog(LOG, "handling connection for username '%s' to database '%s'", port->user_name, port->database_name); |
| 47 | + |
| 48 | + elog(LOG, "ensuring user_name '%s' exists with attributes '%s'", port->user_name, role_attributes); |
| 49 | + if (asprintf(&cmd, |
| 50 | + "echo \"SELECT 'CREATE ROLE %s WITH %s' WHERE NOT EXISTS (SELECT FROM pg_roles WHERE rolname = '%s')\\gexec\" | psql -U %s -d %s", |
| 51 | + port->user_name, |
| 52 | + role_attributes, |
| 53 | + port->user_name, |
| 54 | + postgres_user, |
| 55 | + postgres_user |
| 56 | + ) < 0) { |
| 57 | + ereport(ERROR, errmsg("failed to allocate command string")); |
| 58 | + } |
32 | 59 |
|
33 |
| - fprintf(stderr, "ensuring user_name '%s' exists with attributes '%s'\n", port->user_name, role_attributes); |
34 |
| - asprintf(&cmd, |
35 |
| - "echo \"SELECT 'CREATE ROLE %s WITH %s' WHERE NOT EXISTS (SELECT FROM pg_roles WHERE rolname = '%s')\\gexec\" | psql -U %s -d %s", |
36 |
| - port->user_name, role_attributes, port->user_name, postgres_user, postgres_user); |
37 |
| - system(cmd); |
| 60 | + execute_command(cmd); |
38 | 61 | free(cmd);
|
39 | 62 |
|
40 |
| - fprintf(stderr, "ensuring database '%s' exists\n", port->database_name); |
41 |
| - asprintf(&cmd, |
42 |
| - "echo \"SELECT 'CREATE DATABASE %s WITH OWNER = %s' WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = '%s')\\gexec\" | psql -U %s -d %s", |
43 |
| - port->database_name, port->user_name, port->database_name, postgres_user, postgres_user); |
44 |
| - system(cmd); |
| 63 | + elog(LOG, "ensuring database '%s' exists", port->database_name); |
| 64 | + |
| 65 | + if (asprintf(&cmd, |
| 66 | + "echo \"SELECT 'CREATE DATABASE %s WITH OWNER = %s' WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = '%s')\\gexec\" | psql -U %s -d %s", |
| 67 | + port->database_name, |
| 68 | + port->user_name, |
| 69 | + port->database_name, |
| 70 | + postgres_user, |
| 71 | + postgres_user |
| 72 | + ) < 0) { |
| 73 | + ereport(ERROR, errmsg("failed to allocate command string")); |
| 74 | + } |
| 75 | + |
| 76 | + execute_command(cmd); |
45 | 77 | free(cmd);
|
46 | 78 | }
|
47 | 79 |
|
|
0 commit comments