Line data Source code
1 : use anyhow::{anyhow, bail, Result};
2 : use postgres::Client;
3 : use reqwest::StatusCode;
4 : use std::fs::File;
5 : use std::path::Path;
6 : use tracing::{error, info, instrument, warn};
7 :
8 : use crate::config;
9 : use crate::migration::MigrationRunner;
10 : use crate::params::PG_HBA_ALL_MD5;
11 : use crate::pg_helpers::*;
12 :
13 : use compute_api::responses::{ControlPlaneComputeStatus, ControlPlaneSpecResponse};
14 : use compute_api::spec::ComputeSpec;
15 :
16 : // Do control plane request and return response if any. In case of error it
17 : // returns a bool flag indicating whether it makes sense to retry the request
18 : // and a string with error message.
19 0 : fn do_control_plane_request(
20 0 : uri: &str,
21 0 : jwt: &str,
22 0 : ) -> Result<ControlPlaneSpecResponse, (bool, String)> {
23 0 : let resp = reqwest::blocking::Client::new()
24 0 : .get(uri)
25 0 : .header("Authorization", format!("Bearer {}", jwt))
26 0 : .send()
27 0 : .map_err(|e| {
28 0 : (
29 0 : true,
30 0 : format!("could not perform spec request to control plane: {}", e),
31 0 : )
32 0 : })?;
33 :
34 0 : match resp.status() {
35 0 : StatusCode::OK => match resp.json::<ControlPlaneSpecResponse>() {
36 0 : Ok(spec_resp) => Ok(spec_resp),
37 0 : Err(e) => Err((
38 0 : true,
39 0 : format!("could not deserialize control plane response: {}", e),
40 0 : )),
41 : },
42 : StatusCode::SERVICE_UNAVAILABLE => {
43 0 : Err((true, "control plane is temporarily unavailable".to_string()))
44 : }
45 : StatusCode::BAD_GATEWAY => {
46 : // We have a problem with intermittent 502 errors now
47 : // https://github.com/neondatabase/cloud/issues/2353
48 : // It's fine to retry GET request in this case.
49 0 : Err((true, "control plane request failed with 502".to_string()))
50 : }
51 : // Another code, likely 500 or 404, means that compute is unknown to the control plane
52 : // or some internal failure happened. Doesn't make much sense to retry in this case.
53 0 : _ => Err((
54 0 : false,
55 0 : format!(
56 0 : "unexpected control plane response status code: {}",
57 0 : resp.status()
58 0 : ),
59 0 : )),
60 : }
61 0 : }
62 :
63 : /// Request spec from the control-plane by compute_id. If `NEON_CONTROL_PLANE_TOKEN`
64 : /// env variable is set, it will be used for authorization.
65 0 : pub fn get_spec_from_control_plane(
66 0 : base_uri: &str,
67 0 : compute_id: &str,
68 0 : ) -> Result<Option<ComputeSpec>> {
69 0 : let cp_uri = format!("{base_uri}/compute/api/v2/computes/{compute_id}/spec");
70 0 : let jwt: String = match std::env::var("NEON_CONTROL_PLANE_TOKEN") {
71 0 : Ok(v) => v,
72 0 : Err(_) => "".to_string(),
73 : };
74 0 : let mut attempt = 1;
75 0 : let mut spec: Result<Option<ComputeSpec>> = Ok(None);
76 0 :
77 0 : info!("getting spec from control plane: {}", cp_uri);
78 :
79 : // Do 3 attempts to get spec from the control plane using the following logic:
80 : // - network error -> then retry
81 : // - compute id is unknown or any other error -> bail out
82 : // - no spec for compute yet (Empty state) -> return Ok(None)
83 : // - got spec -> return Ok(Some(spec))
84 0 : while attempt < 4 {
85 0 : spec = match do_control_plane_request(&cp_uri, &jwt) {
86 0 : Ok(spec_resp) => match spec_resp.status {
87 0 : ControlPlaneComputeStatus::Empty => Ok(None),
88 : ControlPlaneComputeStatus::Attached => {
89 0 : if let Some(spec) = spec_resp.spec {
90 0 : Ok(Some(spec))
91 : } else {
92 0 : bail!("compute is attached, but spec is empty")
93 : }
94 : }
95 : },
96 0 : Err((retry, msg)) => {
97 0 : if retry {
98 0 : Err(anyhow!(msg))
99 : } else {
100 0 : bail!(msg);
101 : }
102 : }
103 : };
104 :
105 0 : if let Err(e) = &spec {
106 0 : error!("attempt {} to get spec failed with: {}", attempt, e);
107 : } else {
108 0 : return spec;
109 : }
110 :
111 0 : attempt += 1;
112 0 : std::thread::sleep(std::time::Duration::from_millis(100));
113 : }
114 :
115 : // All attempts failed, return error.
116 0 : spec
117 0 : }
118 :
119 : /// Check `pg_hba.conf` and update if needed to allow external connections.
120 0 : pub fn update_pg_hba(pgdata_path: &Path) -> Result<()> {
121 0 : // XXX: consider making it a part of spec.json
122 0 : info!("checking pg_hba.conf");
123 0 : let pghba_path = pgdata_path.join("pg_hba.conf");
124 0 :
125 0 : if config::line_in_file(&pghba_path, PG_HBA_ALL_MD5)? {
126 0 : info!("updated pg_hba.conf to allow external connections");
127 : } else {
128 0 : info!("pg_hba.conf is up-to-date");
129 : }
130 :
131 0 : Ok(())
132 0 : }
133 :
134 : /// Create a standby.signal file
135 0 : pub fn add_standby_signal(pgdata_path: &Path) -> Result<()> {
136 0 : // XXX: consider making it a part of spec.json
137 0 : info!("adding standby.signal");
138 0 : let signalfile = pgdata_path.join("standby.signal");
139 0 :
140 0 : if !signalfile.exists() {
141 0 : info!("created standby.signal");
142 0 : File::create(signalfile)?;
143 : } else {
144 0 : info!("reused pre-existing standby.signal");
145 : }
146 0 : Ok(())
147 0 : }
148 :
149 0 : #[instrument(skip_all)]
150 : pub fn handle_neon_extension_upgrade(client: &mut Client) -> Result<()> {
151 : info!("handle neon extension upgrade");
152 : let query = "ALTER EXTENSION neon UPDATE";
153 : info!("update neon extension version with query: {}", query);
154 : client.simple_query(query)?;
155 :
156 : Ok(())
157 : }
158 :
159 0 : #[instrument(skip_all)]
160 : pub fn handle_migrations(client: &mut Client) -> Result<()> {
161 : info!("handle migrations");
162 :
163 : // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
164 : // !BE SURE TO ONLY ADD MIGRATIONS TO THE END OF THIS ARRAY. IF YOU DO NOT, VERY VERY BAD THINGS MAY HAPPEN!
165 : // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
166 :
167 : // Add new migrations in numerical order.
168 : let migrations = [
169 : include_str!("./migrations/0001-neon_superuser_bypass_rls.sql"),
170 : include_str!("./migrations/0002-alter_roles.sql"),
171 : include_str!("./migrations/0003-grant_pg_create_subscription_to_neon_superuser.sql"),
172 : include_str!("./migrations/0004-grant_pg_monitor_to_neon_superuser.sql"),
173 : include_str!("./migrations/0005-grant_all_on_tables_to_neon_superuser.sql"),
174 : include_str!("./migrations/0006-grant_all_on_sequences_to_neon_superuser.sql"),
175 : include_str!(
176 : "./migrations/0007-grant_all_on_tables_to_neon_superuser_with_grant_option.sql"
177 : ),
178 : include_str!(
179 : "./migrations/0008-grant_all_on_sequences_to_neon_superuser_with_grant_option.sql"
180 : ),
181 : include_str!("./migrations/0009-revoke_replication_for_previously_allowed_roles.sql"),
182 : include_str!(
183 : "./migrations/0010-grant_snapshot_synchronization_funcs_to_neon_superuser.sql"
184 : ),
185 : include_str!(
186 : "./migrations/0011-grant_pg_show_replication_origin_status_to_neon_superuser.sql"
187 : ),
188 : ];
189 :
190 : MigrationRunner::new(client, &migrations).run_migrations()?;
191 :
192 : Ok(())
193 : }
194 :
195 : /// Connect to the database as superuser and pre-create anon extension
196 : /// if it is present in shared_preload_libraries
197 0 : #[instrument(skip_all)]
198 : pub fn handle_extension_anon(
199 : spec: &ComputeSpec,
200 : db_owner: &str,
201 : db_client: &mut Client,
202 : grants_only: bool,
203 : ) -> Result<()> {
204 : info!("handle extension anon");
205 :
206 : if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
207 : if libs.contains("anon") {
208 : if !grants_only {
209 : // check if extension is already initialized using anon.is_initialized()
210 : let query = "SELECT anon.is_initialized()";
211 : match db_client.query(query, &[]) {
212 : Ok(rows) => {
213 : if !rows.is_empty() {
214 : let is_initialized: bool = rows[0].get(0);
215 : if is_initialized {
216 : info!("anon extension is already initialized");
217 : return Ok(());
218 : }
219 : }
220 : }
221 : Err(e) => {
222 : warn!(
223 : "anon extension is_installed check failed with expected error: {}",
224 : e
225 : );
226 : }
227 : };
228 :
229 : // Create anon extension if this compute needs it
230 : // Users cannot create it themselves, because superuser is required.
231 : let mut query = "CREATE EXTENSION IF NOT EXISTS anon CASCADE";
232 : info!("creating anon extension with query: {}", query);
233 : match db_client.query(query, &[]) {
234 : Ok(_) => {}
235 : Err(e) => {
236 : error!("anon extension creation failed with error: {}", e);
237 : return Ok(());
238 : }
239 : }
240 :
241 : // check that extension is installed
242 : query = "SELECT extname FROM pg_extension WHERE extname = 'anon'";
243 : let rows = db_client.query(query, &[])?;
244 : if rows.is_empty() {
245 : error!("anon extension is not installed");
246 : return Ok(());
247 : }
248 :
249 : // Initialize anon extension
250 : // This also requires superuser privileges, so users cannot do it themselves.
251 : query = "SELECT anon.init()";
252 : match db_client.query(query, &[]) {
253 : Ok(_) => {}
254 : Err(e) => {
255 : error!("anon.init() failed with error: {}", e);
256 : return Ok(());
257 : }
258 : }
259 : }
260 :
261 : // check that extension is installed, if not bail early
262 : let query = "SELECT extname FROM pg_extension WHERE extname = 'anon'";
263 : match db_client.query(query, &[]) {
264 : Ok(rows) => {
265 : if rows.is_empty() {
266 : error!("anon extension is not installed");
267 : return Ok(());
268 : }
269 : }
270 : Err(e) => {
271 : error!("anon extension check failed with error: {}", e);
272 : return Ok(());
273 : }
274 : };
275 :
276 : let query = format!("GRANT ALL ON SCHEMA anon TO {}", db_owner);
277 : info!("granting anon extension permissions with query: {}", query);
278 : db_client.simple_query(&query)?;
279 :
280 : // Grant permissions to db_owner to use anon extension functions
281 : let query = format!("GRANT ALL ON ALL FUNCTIONS IN SCHEMA anon TO {}", db_owner);
282 : info!("granting anon extension permissions with query: {}", query);
283 : db_client.simple_query(&query)?;
284 :
285 : // This is needed, because some functions are defined as SECURITY DEFINER.
286 : // In Postgres SECURITY DEFINER functions are executed with the privileges
287 : // of the owner.
288 : // In anon extension this it is needed to access some GUCs, which are only accessible to
289 : // superuser. But we've patched postgres to allow db_owner to access them as well.
290 : // So we need to change owner of these functions to db_owner.
291 : let query = format!("
292 : SELECT 'ALTER FUNCTION '||nsp.nspname||'.'||p.proname||'('||pg_get_function_identity_arguments(p.oid)||') OWNER TO {};'
293 : from pg_proc p
294 : join pg_namespace nsp ON p.pronamespace = nsp.oid
295 : where nsp.nspname = 'anon';", db_owner);
296 :
297 : info!("change anon extension functions owner to db owner");
298 : db_client.simple_query(&query)?;
299 :
300 : // affects views as well
301 : let query = format!("GRANT ALL ON ALL TABLES IN SCHEMA anon TO {}", db_owner);
302 : info!("granting anon extension permissions with query: {}", query);
303 : db_client.simple_query(&query)?;
304 :
305 : let query = format!("GRANT ALL ON ALL SEQUENCES IN SCHEMA anon TO {}", db_owner);
306 : info!("granting anon extension permissions with query: {}", query);
307 : db_client.simple_query(&query)?;
308 : }
309 : }
310 :
311 : Ok(())
312 : }
|
