LCOV - code coverage report
Current view: top level - safekeeper/src - handler.rs (source / functions) Coverage Total Hit
Test: 6df3fc19ec669bcfbbf9aba41d1338898d24eaa0.info Lines: 16.2 % 370 60
Test Date: 2025-03-12 18:28:53 Functions: 7.3 % 41 3

            Line data    Source code
       1              : //! Part of Safekeeper pretending to be Postgres, i.e. handling Postgres
       2              : //! protocol commands.
       3              : 
       4              : use std::future::Future;
       5              : use std::str::{self, FromStr};
       6              : use std::sync::Arc;
       7              : 
       8              : use anyhow::Context;
       9              : use pageserver_api::models::ShardParameters;
      10              : use pageserver_api::shard::{ShardIdentity, ShardStripeSize};
      11              : use postgres_backend::{PostgresBackend, QueryError};
      12              : use postgres_ffi::PG_TLI;
      13              : use pq_proto::{BeMessage, FeStartupPacket, INT4_OID, RowDescriptor, TEXT_OID};
      14              : use regex::Regex;
      15              : use safekeeper_api::Term;
      16              : use safekeeper_api::models::ConnectionId;
      17              : use tokio::io::{AsyncRead, AsyncWrite};
      18              : use tracing::{Instrument, debug, info, info_span};
      19              : use utils::auth::{Claims, JwtAuth, Scope};
      20              : use utils::id::{TenantId, TenantTimelineId, TimelineId};
      21              : use utils::lsn::Lsn;
      22              : use utils::postgres_client::PostgresClientProtocol;
      23              : use utils::shard::{ShardCount, ShardNumber};
      24              : 
      25              : use crate::auth::check_permission;
      26              : use crate::metrics::{PG_QUERIES_GAUGE, TrafficMetrics};
      27              : use crate::timeline::TimelineError;
      28              : use crate::{GlobalTimelines, SafeKeeperConf};
      29              : 
      30              : /// Safekeeper handler of postgres commands
      31              : pub struct SafekeeperPostgresHandler {
      32              :     pub conf: Arc<SafeKeeperConf>,
      33              :     /// assigned application name
      34              :     pub appname: Option<String>,
      35              :     pub tenant_id: Option<TenantId>,
      36              :     pub timeline_id: Option<TimelineId>,
      37              :     pub ttid: TenantTimelineId,
      38              :     pub shard: Option<ShardIdentity>,
      39              :     pub protocol: Option<PostgresClientProtocol>,
      40              :     /// Unique connection id is logged in spans for observability.
      41              :     pub conn_id: ConnectionId,
      42              :     pub global_timelines: Arc<GlobalTimelines>,
      43              :     /// Auth scope allowed on the connections and public key used to check auth tokens. None if auth is not configured.
      44              :     auth: Option<(Scope, Arc<JwtAuth>)>,
      45              :     claims: Option<Claims>,
      46              :     io_metrics: Option<TrafficMetrics>,
      47              : }
      48              : 
      49              : /// Parsed Postgres command.
      50              : enum SafekeeperPostgresCommand {
      51              :     StartWalPush {
      52              :         proto_version: u32,
      53              :         // Eventually timelines will be always created explicitly by storcon.
      54              :         // This option allows legacy behaviour for compute to do that until we
      55              :         // fully migrate.
      56              :         allow_timeline_creation: bool,
      57              :     },
      58              :     StartReplication {
      59              :         start_lsn: Lsn,
      60              :         term: Option<Term>,
      61              :     },
      62              :     IdentifySystem,
      63              :     TimelineStatus,
      64              : }
      65              : 
      66            2 : fn parse_cmd(cmd: &str) -> anyhow::Result<SafekeeperPostgresCommand> {
      67            2 :     if cmd.starts_with("START_WAL_PUSH") {
      68              :         // Allow additional options in postgres START_REPLICATION style like
      69              :         //   START_WAL_PUSH (proto_version '3', allow_timeline_creation 'false').
      70              :         // Parsing here is very naive and breaks in case of commas or
      71              :         // whitespaces in values, but enough for our purposes.
      72            2 :         let re = Regex::new(r"START_WAL_PUSH(\s+?\((.*)\))?").unwrap();
      73            2 :         let caps = re
      74            2 :             .captures(cmd)
      75            2 :             .context(format!("failed to parse START_WAL_PUSH command {}", cmd))?;
      76              :         // capture () content
      77            2 :         let options = caps.get(2).map(|m| m.as_str()).unwrap_or("");
      78            2 :         // default values
      79            2 :         let mut proto_version = 2;
      80            2 :         let mut allow_timeline_creation = true;
      81            4 :         for kvstr in options.split(",") {
      82            4 :             if kvstr.is_empty() {
      83            1 :                 continue;
      84            3 :             }
      85            3 :             let mut kvit = kvstr.split_whitespace();
      86            3 :             let key = kvit.next().context(format!(
      87            3 :                 "failed to parse key in kv {} in command {}",
      88            3 :                 kvstr, cmd
      89            3 :             ))?;
      90            3 :             let value = kvit.next().context(format!(
      91            3 :                 "failed to parse value in kv {} in command {}",
      92            3 :                 kvstr, cmd
      93            3 :             ))?;
      94            3 :             let value_trimmed = value.trim_matches('\'');
      95            3 :             if key == "proto_version" {
      96            1 :                 proto_version = value_trimmed.parse::<u32>().context(format!(
      97            1 :                     "failed to parse proto_version value {} in command {}",
      98            1 :                     value, cmd
      99            1 :                 ))?;
     100            2 :             }
     101            3 :             if key == "allow_timeline_creation" {
     102            1 :                 allow_timeline_creation = value_trimmed.parse::<bool>().context(format!(
     103            1 :                     "failed to parse allow_timeline_creation value {} in command {}",
     104            1 :                     value, cmd
     105            1 :                 ))?;
     106            2 :             }
     107              :         }
     108            2 :         Ok(SafekeeperPostgresCommand::StartWalPush {
     109            2 :             proto_version,
     110            2 :             allow_timeline_creation,
     111            2 :         })
     112            0 :     } else if cmd.starts_with("START_REPLICATION") {
     113            0 :         let re = Regex::new(
     114            0 :             // We follow postgres START_REPLICATION LOGICAL options to pass term.
     115            0 :             r"START_REPLICATION(?: SLOT [^ ]+)?(?: PHYSICAL)? ([[:xdigit:]]+/[[:xdigit:]]+)(?: \(term='(\d+)'\))?",
     116            0 :         )
     117            0 :         .unwrap();
     118            0 :         let caps = re
     119            0 :             .captures(cmd)
     120            0 :             .context(format!("failed to parse START_REPLICATION command {}", cmd))?;
     121            0 :         let start_lsn =
     122            0 :             Lsn::from_str(&caps[1]).context("parse start LSN from START_REPLICATION command")?;
     123            0 :         let term = if let Some(m) = caps.get(2) {
     124            0 :             Some(m.as_str().parse::<u64>().context("invalid term")?)
     125              :         } else {
     126            0 :             None
     127              :         };
     128            0 :         Ok(SafekeeperPostgresCommand::StartReplication { start_lsn, term })
     129            0 :     } else if cmd.starts_with("IDENTIFY_SYSTEM") {
     130            0 :         Ok(SafekeeperPostgresCommand::IdentifySystem)
     131            0 :     } else if cmd.starts_with("TIMELINE_STATUS") {
     132            0 :         Ok(SafekeeperPostgresCommand::TimelineStatus)
     133              :     } else {
     134            0 :         anyhow::bail!("unsupported command {cmd}");
     135              :     }
     136            2 : }
     137              : 
     138            0 : fn cmd_to_string(cmd: &SafekeeperPostgresCommand) -> &str {
     139            0 :     match cmd {
     140            0 :         SafekeeperPostgresCommand::StartWalPush { .. } => "START_WAL_PUSH",
     141            0 :         SafekeeperPostgresCommand::StartReplication { .. } => "START_REPLICATION",
     142            0 :         SafekeeperPostgresCommand::TimelineStatus => "TIMELINE_STATUS",
     143            0 :         SafekeeperPostgresCommand::IdentifySystem => "IDENTIFY_SYSTEM",
     144              :     }
     145            0 : }
     146              : 
     147              : impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::Handler<IO>
     148              :     for SafekeeperPostgresHandler
     149              : {
     150              :     // tenant_id and timeline_id are passed in connection string params
     151            0 :     fn startup(
     152            0 :         &mut self,
     153            0 :         _pgb: &mut PostgresBackend<IO>,
     154            0 :         sm: &FeStartupPacket,
     155            0 :     ) -> Result<(), QueryError> {
     156            0 :         if let FeStartupPacket::StartupMessage { params, .. } = sm {
     157            0 :             if let Some(options) = params.options_raw() {
     158            0 :                 let mut shard_count: Option<u8> = None;
     159            0 :                 let mut shard_number: Option<u8> = None;
     160            0 :                 let mut shard_stripe_size: Option<u32> = None;
     161              : 
     162            0 :                 for opt in options {
     163              :                     // FIXME `ztenantid` and `ztimelineid` left for compatibility during deploy,
     164              :                     // remove these after the PR gets deployed:
     165              :                     // https://github.com/neondatabase/neon/pull/2433#discussion_r970005064
     166            0 :                     match opt.split_once('=') {
     167            0 :                         Some(("protocol", value)) => {
     168            0 :                             self.protocol =
     169            0 :                                 Some(serde_json::from_str(value).with_context(|| {
     170            0 :                                     format!("Failed to parse {value} as protocol")
     171            0 :                                 })?);
     172              :                         }
     173            0 :                         Some(("ztenantid", value)) | Some(("tenant_id", value)) => {
     174            0 :                             self.tenant_id = Some(value.parse().with_context(|| {
     175            0 :                                 format!("Failed to parse {value} as tenant id")
     176            0 :                             })?);
     177              :                         }
     178            0 :                         Some(("ztimelineid", value)) | Some(("timeline_id", value)) => {
     179            0 :                             self.timeline_id = Some(value.parse().with_context(|| {
     180            0 :                                 format!("Failed to parse {value} as timeline id")
     181            0 :                             })?);
     182              :                         }
     183            0 :                         Some(("availability_zone", client_az)) => {
     184            0 :                             if let Some(metrics) = self.io_metrics.as_ref() {
     185            0 :                                 metrics.set_client_az(client_az)
     186            0 :                             }
     187              :                         }
     188            0 :                         Some(("shard_count", value)) => {
     189            0 :                             shard_count = Some(value.parse::<u8>().with_context(|| {
     190            0 :                                 format!("Failed to parse {value} as shard count")
     191            0 :                             })?);
     192              :                         }
     193            0 :                         Some(("shard_number", value)) => {
     194            0 :                             shard_number = Some(value.parse::<u8>().with_context(|| {
     195            0 :                                 format!("Failed to parse {value} as shard number")
     196            0 :                             })?);
     197              :                         }
     198            0 :                         Some(("shard_stripe_size", value)) => {
     199            0 :                             shard_stripe_size = Some(value.parse::<u32>().with_context(|| {
     200            0 :                                 format!("Failed to parse {value} as shard stripe size")
     201            0 :                             })?);
     202              :                         }
     203            0 :                         _ => continue,
     204              :                     }
     205              :                 }
     206              : 
     207            0 :                 match self.protocol() {
     208              :                     PostgresClientProtocol::Vanilla => {
     209            0 :                         if shard_count.is_some()
     210            0 :                             || shard_number.is_some()
     211            0 :                             || shard_stripe_size.is_some()
     212              :                         {
     213            0 :                             return Err(QueryError::Other(anyhow::anyhow!(
     214            0 :                                 "Shard params specified for vanilla protocol"
     215            0 :                             )));
     216            0 :                         }
     217              :                     }
     218              :                     PostgresClientProtocol::Interpreted { .. } => {
     219            0 :                         match (shard_count, shard_number, shard_stripe_size) {
     220            0 :                             (Some(count), Some(number), Some(stripe_size)) => {
     221            0 :                                 let params = ShardParameters {
     222            0 :                                     count: ShardCount(count),
     223            0 :                                     stripe_size: ShardStripeSize(stripe_size),
     224            0 :                                 };
     225            0 :                                 self.shard =
     226            0 :                                     Some(ShardIdentity::from_params(ShardNumber(number), &params));
     227            0 :                             }
     228              :                             _ => {
     229            0 :                                 return Err(QueryError::Other(anyhow::anyhow!(
     230            0 :                                     "Shard params were not specified"
     231            0 :                                 )));
     232              :                             }
     233              :                         }
     234              :                     }
     235              :                 }
     236            0 :             }
     237              : 
     238            0 :             if let Some(app_name) = params.get("application_name") {
     239            0 :                 self.appname = Some(app_name.to_owned());
     240            0 :                 if let Some(metrics) = self.io_metrics.as_ref() {
     241            0 :                     metrics.set_app_name(app_name)
     242            0 :                 }
     243            0 :             }
     244              : 
     245            0 :             let ttid = TenantTimelineId::new(
     246            0 :                 self.tenant_id.unwrap_or(TenantId::from([0u8; 16])),
     247            0 :                 self.timeline_id.unwrap_or(TimelineId::from([0u8; 16])),
     248            0 :             );
     249            0 :             tracing::Span::current()
     250            0 :                 .record("ttid", tracing::field::display(ttid))
     251            0 :                 .record(
     252            0 :                     "application_name",
     253            0 :                     tracing::field::debug(self.appname.clone()),
     254            0 :                 );
     255              : 
     256            0 :             if let Some(shard) = self.shard.as_ref() {
     257            0 :                 if let Some(slug) = shard.shard_slug().strip_prefix("-") {
     258            0 :                     tracing::Span::current().record("shard", tracing::field::display(slug));
     259            0 :                 }
     260            0 :             }
     261              : 
     262            0 :             Ok(())
     263              :         } else {
     264            0 :             Err(QueryError::Other(anyhow::anyhow!(
     265            0 :                 "Safekeeper received unexpected initial message: {sm:?}"
     266            0 :             )))
     267              :         }
     268            0 :     }
     269              : 
     270            0 :     fn check_auth_jwt(
     271            0 :         &mut self,
     272            0 :         _pgb: &mut PostgresBackend<IO>,
     273            0 :         jwt_response: &[u8],
     274            0 :     ) -> Result<(), QueryError> {
     275            0 :         // this unwrap is never triggered, because check_auth_jwt only called when auth_type is NeonJWT
     276            0 :         // which requires auth to be present
     277            0 :         let (allowed_auth_scope, auth) = self
     278            0 :             .auth
     279            0 :             .as_ref()
     280            0 :             .expect("auth_type is configured but .auth of handler is missing");
     281            0 :         let data = auth
     282            0 :             .decode(str::from_utf8(jwt_response).context("jwt response is not UTF-8")?)
     283            0 :             .map_err(|e| QueryError::Unauthorized(e.0))?;
     284              : 
     285              :         // The handler might be configured to allow only tenant scope tokens.
     286            0 :         if matches!(allowed_auth_scope, Scope::Tenant)
     287            0 :             && !matches!(data.claims.scope, Scope::Tenant)
     288              :         {
     289            0 :             return Err(QueryError::Unauthorized(
     290            0 :                 "passed JWT token is for full access, but only tenant scope is allowed".into(),
     291            0 :             ));
     292            0 :         }
     293              : 
     294            0 :         if matches!(data.claims.scope, Scope::Tenant) && data.claims.tenant_id.is_none() {
     295            0 :             return Err(QueryError::Unauthorized(
     296            0 :                 "jwt token scope is Tenant, but tenant id is missing".into(),
     297            0 :             ));
     298            0 :         }
     299            0 : 
     300            0 :         debug!(
     301            0 :             "jwt scope check succeeded for scope: {:#?} by tenant id: {:?}",
     302              :             data.claims.scope, data.claims.tenant_id,
     303              :         );
     304              : 
     305            0 :         self.claims = Some(data.claims);
     306            0 :         Ok(())
     307            0 :     }
     308              : 
     309            0 :     fn process_query(
     310            0 :         &mut self,
     311            0 :         pgb: &mut PostgresBackend<IO>,
     312            0 :         query_string: &str,
     313            0 :     ) -> impl Future<Output = Result<(), QueryError>> {
     314            0 :         Box::pin(async move {
     315            0 :             if query_string
     316            0 :                 .to_ascii_lowercase()
     317            0 :                 .starts_with("set datestyle to ")
     318              :             {
     319              :                 // important for debug because psycopg2 executes "SET datestyle TO 'ISO'" on connect
     320            0 :                 pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
     321            0 :                 return Ok(());
     322            0 :             }
     323              : 
     324            0 :             let cmd = parse_cmd(query_string)?;
     325            0 :             let cmd_str = cmd_to_string(&cmd);
     326            0 : 
     327            0 :             let _guard = PG_QUERIES_GAUGE.with_label_values(&[cmd_str]).guard();
     328            0 : 
     329            0 :             info!("got query {:?}", query_string);
     330              : 
     331            0 :             let tenant_id = self.tenant_id.context("tenantid is required")?;
     332            0 :             let timeline_id = self.timeline_id.context("timelineid is required")?;
     333            0 :             self.check_permission(Some(tenant_id))?;
     334            0 :             self.ttid = TenantTimelineId::new(tenant_id, timeline_id);
     335            0 : 
     336            0 :             match cmd {
     337              :                 SafekeeperPostgresCommand::StartWalPush {
     338            0 :                     proto_version,
     339            0 :                     allow_timeline_creation,
     340            0 :                 } => {
     341            0 :                     self.handle_start_wal_push(pgb, proto_version, allow_timeline_creation)
     342            0 :                         .instrument(info_span!("WAL receiver"))
     343            0 :                         .await
     344              :                 }
     345            0 :                 SafekeeperPostgresCommand::StartReplication { start_lsn, term } => {
     346            0 :                     self.handle_start_replication(pgb, start_lsn, term)
     347            0 :                         .instrument(info_span!("WAL sender"))
     348            0 :                         .await
     349              :                 }
     350            0 :                 SafekeeperPostgresCommand::IdentifySystem => self.handle_identify_system(pgb).await,
     351            0 :                 SafekeeperPostgresCommand::TimelineStatus => self.handle_timeline_status(pgb).await,
     352              :             }
     353            0 :         })
     354            0 :     }
     355              : }
     356              : 
     357              : impl SafekeeperPostgresHandler {
     358            0 :     pub fn new(
     359            0 :         conf: Arc<SafeKeeperConf>,
     360            0 :         conn_id: u32,
     361            0 :         io_metrics: Option<TrafficMetrics>,
     362            0 :         auth: Option<(Scope, Arc<JwtAuth>)>,
     363            0 :         global_timelines: Arc<GlobalTimelines>,
     364            0 :     ) -> Self {
     365            0 :         SafekeeperPostgresHandler {
     366            0 :             conf,
     367            0 :             appname: None,
     368            0 :             tenant_id: None,
     369            0 :             timeline_id: None,
     370            0 :             ttid: TenantTimelineId::empty(),
     371            0 :             shard: None,
     372            0 :             protocol: None,
     373            0 :             conn_id,
     374            0 :             claims: None,
     375            0 :             auth,
     376            0 :             io_metrics,
     377            0 :             global_timelines,
     378            0 :         }
     379            0 :     }
     380              : 
     381            0 :     pub fn protocol(&self) -> PostgresClientProtocol {
     382            0 :         self.protocol.unwrap_or(PostgresClientProtocol::Vanilla)
     383            0 :     }
     384              : 
     385              :     // when accessing management api supply None as an argument
     386              :     // when using to authorize tenant pass corresponding tenant id
     387            0 :     fn check_permission(&self, tenant_id: Option<TenantId>) -> Result<(), QueryError> {
     388            0 :         if self.auth.is_none() {
     389              :             // auth is set to Trust, nothing to check so just return ok
     390            0 :             return Ok(());
     391            0 :         }
     392            0 :         // auth is some, just checked above, when auth is some
     393            0 :         // then claims are always present because of checks during connection init
     394            0 :         // so this expect won't trigger
     395            0 :         let claims = self
     396            0 :             .claims
     397            0 :             .as_ref()
     398            0 :             .expect("claims presence already checked");
     399            0 :         check_permission(claims, tenant_id).map_err(|e| QueryError::Unauthorized(e.0))
     400            0 :     }
     401              : 
     402            0 :     async fn handle_timeline_status<IO: AsyncRead + AsyncWrite + Unpin>(
     403            0 :         &mut self,
     404            0 :         pgb: &mut PostgresBackend<IO>,
     405            0 :     ) -> Result<(), QueryError> {
     406              :         // Get timeline, handling "not found" error
     407            0 :         let tli = match self.global_timelines.get(self.ttid) {
     408            0 :             Ok(tli) => Ok(Some(tli)),
     409            0 :             Err(TimelineError::NotFound(_)) => Ok(None),
     410            0 :             Err(e) => Err(QueryError::Other(e.into())),
     411            0 :         }?;
     412              : 
     413              :         // Write row description
     414            0 :         pgb.write_message_noflush(&BeMessage::RowDescription(&[
     415            0 :             RowDescriptor::text_col(b"flush_lsn"),
     416            0 :             RowDescriptor::text_col(b"commit_lsn"),
     417            0 :         ]))?;
     418              : 
     419              :         // Write row if timeline exists
     420            0 :         if let Some(tli) = tli {
     421            0 :             let (inmem, _state) = tli.get_state().await;
     422            0 :             let flush_lsn = tli.get_flush_lsn().await;
     423            0 :             let commit_lsn = inmem.commit_lsn;
     424            0 :             pgb.write_message_noflush(&BeMessage::DataRow(&[
     425            0 :                 Some(flush_lsn.to_string().as_bytes()),
     426            0 :                 Some(commit_lsn.to_string().as_bytes()),
     427            0 :             ]))?;
     428            0 :         }
     429              : 
     430            0 :         pgb.write_message_noflush(&BeMessage::CommandComplete(b"TIMELINE_STATUS"))?;
     431            0 :         Ok(())
     432            0 :     }
     433              : 
     434              :     ///
     435              :     /// Handle IDENTIFY_SYSTEM replication command
     436              :     ///
     437            0 :     async fn handle_identify_system<IO: AsyncRead + AsyncWrite + Unpin>(
     438            0 :         &mut self,
     439            0 :         pgb: &mut PostgresBackend<IO>,
     440            0 :     ) -> Result<(), QueryError> {
     441            0 :         let tli = self
     442            0 :             .global_timelines
     443            0 :             .get(self.ttid)
     444            0 :             .map_err(|e| QueryError::Other(e.into()))?;
     445              : 
     446            0 :         let lsn = if self.is_walproposer_recovery() {
     447              :             // walproposer should get all local WAL until flush_lsn
     448            0 :             tli.get_flush_lsn().await
     449              :         } else {
     450              :             // other clients shouldn't get any uncommitted WAL
     451            0 :             tli.get_state().await.0.commit_lsn
     452              :         }
     453            0 :         .to_string();
     454              : 
     455            0 :         let sysid = tli.get_state().await.1.server.system_id.to_string();
     456            0 :         let lsn_bytes = lsn.as_bytes();
     457            0 :         let tli = PG_TLI.to_string();
     458            0 :         let tli_bytes = tli.as_bytes();
     459            0 :         let sysid_bytes = sysid.as_bytes();
     460            0 : 
     461            0 :         pgb.write_message_noflush(&BeMessage::RowDescription(&[
     462            0 :             RowDescriptor {
     463            0 :                 name: b"systemid",
     464            0 :                 typoid: TEXT_OID,
     465            0 :                 typlen: -1,
     466            0 :                 ..Default::default()
     467            0 :             },
     468            0 :             RowDescriptor {
     469            0 :                 name: b"timeline",
     470            0 :                 typoid: INT4_OID,
     471            0 :                 typlen: 4,
     472            0 :                 ..Default::default()
     473            0 :             },
     474            0 :             RowDescriptor {
     475            0 :                 name: b"xlogpos",
     476            0 :                 typoid: TEXT_OID,
     477            0 :                 typlen: -1,
     478            0 :                 ..Default::default()
     479            0 :             },
     480            0 :             RowDescriptor {
     481            0 :                 name: b"dbname",
     482            0 :                 typoid: TEXT_OID,
     483            0 :                 typlen: -1,
     484            0 :                 ..Default::default()
     485            0 :             },
     486            0 :         ]))?
     487            0 :         .write_message_noflush(&BeMessage::DataRow(&[
     488            0 :             Some(sysid_bytes),
     489            0 :             Some(tli_bytes),
     490            0 :             Some(lsn_bytes),
     491            0 :             None,
     492            0 :         ]))?
     493            0 :         .write_message_noflush(&BeMessage::CommandComplete(b"IDENTIFY_SYSTEM"))?;
     494            0 :         Ok(())
     495            0 :     }
     496              : 
     497              :     /// Returns true if current connection is a replication connection, originating
     498              :     /// from a walproposer recovery function. This connection gets a special handling:
     499              :     /// safekeeper must stream all local WAL till the flush_lsn, whether committed or not.
     500            0 :     pub fn is_walproposer_recovery(&self) -> bool {
     501            0 :         match &self.appname {
     502            0 :             None => false,
     503            0 :             Some(appname) => {
     504            0 :                 appname == "wal_proposer_recovery" ||
     505              :                 // set by safekeeper peer recovery
     506            0 :                 appname.starts_with("safekeeper")
     507              :             }
     508              :         }
     509            0 :     }
     510              : }
     511              : 
     512              : #[cfg(test)]
     513              : mod tests {
     514              :     use super::SafekeeperPostgresCommand;
     515              : 
     516              :     /// Test parsing of START_WAL_PUSH command
     517              :     #[test]
     518            1 :     fn test_start_wal_push_parse() {
     519            1 :         let cmd = "START_WAL_PUSH";
     520            1 :         let parsed = super::parse_cmd(cmd).expect("failed to parse");
     521            1 :         match parsed {
     522              :             SafekeeperPostgresCommand::StartWalPush {
     523            1 :                 proto_version,
     524            1 :                 allow_timeline_creation,
     525            1 :             } => {
     526            1 :                 assert_eq!(proto_version, 2);
     527            1 :                 assert!(allow_timeline_creation);
     528              :             }
     529            0 :             _ => panic!("unexpected command"),
     530              :         }
     531              : 
     532            1 :         let cmd =
     533            1 :             "START_WAL_PUSH (proto_version '3', allow_timeline_creation 'false', unknown 'hoho')";
     534            1 :         let parsed = super::parse_cmd(cmd).expect("failed to parse");
     535            1 :         match parsed {
     536              :             SafekeeperPostgresCommand::StartWalPush {
     537            1 :                 proto_version,
     538            1 :                 allow_timeline_creation,
     539            1 :             } => {
     540            1 :                 assert_eq!(proto_version, 3);
     541            1 :                 assert!(!allow_timeline_creation);
     542              :             }
     543            0 :             _ => panic!("unexpected command"),
     544              :         }
     545            1 :     }
     546              : }
        

Generated by: LCOV version 2.1-beta