Skip custom extensions in fast import

.gitmodules

@@ -1,16 +1,16 @@
 [submodule "vendor/postgres-v14"]
 	path = vendor/postgres-v14
 	url = https://github.com/neondatabase/postgres.git
-	branch = REL_14_STABLE_neon
+	branch = 28934-pg-dump-schema-no-create-v14
 [submodule "vendor/postgres-v15"]
 	path = vendor/postgres-v15
 	url = https://github.com/neondatabase/postgres.git
-	branch = REL_15_STABLE_neon
+	branch = 28934-pg-dump-schema-no-create-v15
 [submodule "vendor/postgres-v16"]
 	path = vendor/postgres-v16
 	url = https://github.com/neondatabase/postgres.git
-	branch = REL_16_STABLE_neon
+	branch = 28934-pg-dump-schema-no-create-v16
 [submodule "vendor/postgres-v17"]
 	path = vendor/postgres-v17
 	url = https://github.com/neondatabase/postgres.git
-	branch = REL_17_STABLE_neon
+	branch = 28934-pg-dump-schema-no-create-v17

compute_tools/src/bin/fast_import.rs

@@ -70,6 +70,14 @@ enum Command {
         /// and maintenance_work_mem.
         #[clap(long, env = "NEON_IMPORTER_MEMORY_MB")]
         memory_mb: Option<usize>,
+
+        /// List of schemas to dump.
+        #[clap(long)]
+        schema: Vec<String>,
+
+        /// List of extensions to dump.
+        #[clap(long)]
+        extension: Vec<String>,
     },
 
     /// Runs pg_dump-pg_restore from source to destination without running local postgres.
@@ -82,6 +90,12 @@ enum Command {
         /// real scenario uses encrypted connection string in spec.json from s3.
         #[clap(long)]
         destination_connection_string: Option<String>,
+        /// List of schemas to dump.
+        #[clap(long)]
+        schema: Vec<String>,
+        /// List of extensions to dump.
+        #[clap(long)]
+        extension: Vec<String>,
     },
 }
 
@@ -117,6 +131,8 @@ struct Spec {
     source_connstring_ciphertext_base64: Vec<u8>,
     #[serde_as(as = "Option<serde_with::base64::Base64>")]
     destination_connstring_ciphertext_base64: Option<Vec<u8>>,
+    schemas: Option<Vec<String>>,
+    extensions: Option<Vec<String>>,
 }
 
 #[derive(serde::Deserialize)]
@@ -337,6 +353,8 @@ async fn run_dump_restore(
     pg_lib_dir: Utf8PathBuf,
     source_connstring: String,
     destination_connstring: String,
+    schemas: Vec<String>,
+    extensions: Vec<String>,
 ) -> Result<(), anyhow::Error> {
     let dumpdir = workdir.join("dumpdir");
     let num_jobs = num_cpus::get().to_string();
@@ -351,6 +369,7 @@ async fn run_dump_restore(
         "--no-subscriptions".to_string(),
         "--no-tablespaces".to_string(),
         "--no-event-triggers".to_string(),
+        "--enable-row-security".to_string(),
         // format
         "--format".to_string(),
         "directory".to_string(),
@@ -361,10 +380,36 @@ async fn run_dump_restore(
         "--verbose".to_string(),
     ];
 
+    let mut pg_dump_args = vec![
+        // this makes sure any unsupported extensions are not included in the dump
+        // even if we don't specify supported extensions explicitly
+        "--extension".to_string(),
+        "plpgsql".to_string(),
+    ];
+
+    // if no schemas are specified, try to import all schemas
+    if !schemas.is_empty() {
+        // always include public schema objects
+        // but never create the schema itself
+        // it already exists in any pg cluster by default
+        pg_dump_args.push("--schema-no-create".to_string());
+        pg_dump_args.push("public".to_string());
+        for schema in &schemas {
+            pg_dump_args.push("--schema".to_string());
+            pg_dump_args.push(schema.clone());
+        }
+    }
+
+    for extension in &extensions {
+        pg_dump_args.push("--extension".to_string());
+        pg_dump_args.push(extension.clone());
+    }
+
     info!("dump into the working directory");
     {
         let mut pg_dump = tokio::process::Command::new(pg_bin_dir.join("pg_dump"))
             .args(&common_args)
+            .args(&pg_dump_args)
             .arg("-f")
             .arg(&dumpdir)
             .arg("--no-sync")
@@ -455,6 +500,8 @@ async fn cmd_pgdata(
     maybe_s3_prefix: Option<s3_uri::S3Uri>,
     maybe_spec: Option<Spec>,
     source_connection_string: Option<String>,
+    schemas: Vec<String>,
+    extensions: Vec<String>,
     interactive: bool,
     pg_port: u16,
     workdir: Utf8PathBuf,
@@ -470,19 +517,25 @@ async fn cmd_pgdata(
         bail!("only one of spec or source_connection_string can be provided");
     }
 
-    let source_connection_string = if let Some(spec) = maybe_spec {
+    let (source_connection_string, schemas, extensions) = if let Some(spec) = maybe_spec {
         match spec.encryption_secret {
             EncryptionSecret::KMS { key_id } => {
-                decode_connstring(
+                let schemas = spec.schemas.unwrap_or(vec![]);
+                let extensions = spec.extensions.unwrap_or(vec![]);
+
+                let source = decode_connstring(
                     kms_client.as_ref().unwrap(),
                     &key_id,
                     spec.source_connstring_ciphertext_base64,
                 )
-                .await?
+                .await
+                .context("decrypt source connection string")?;
+
+                (source, schemas, extensions)
             }
         }
     } else {
-        source_connection_string.unwrap()
+        (source_connection_string.unwrap(), schemas, extensions)
     };
 
     let superuser = "cloud_admin";
@@ -504,6 +557,8 @@ async fn cmd_pgdata(
         pg_lib_dir,
         source_connection_string,
         destination_connstring,
+        schemas,
+        extensions,
     )
     .await?;
 
@@ -546,18 +601,26 @@ async fn cmd_pgdata(
     Ok(())
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn cmd_dumprestore(
     kms_client: Option<aws_sdk_kms::Client>,
     maybe_spec: Option<Spec>,
     source_connection_string: Option<String>,
     destination_connection_string: Option<String>,
+    schemas: Vec<String>,
+    extensions: Vec<String>,
     workdir: Utf8PathBuf,
     pg_bin_dir: Utf8PathBuf,
     pg_lib_dir: Utf8PathBuf,
 ) -> Result<(), anyhow::Error> {
-    let (source_connstring, destination_connstring) = if let Some(spec) = maybe_spec {
+    let (source_connstring, destination_connstring, schemas, extensions) = if let Some(spec) =
+        maybe_spec
+    {
         match spec.encryption_secret {
             EncryptionSecret::KMS { key_id } => {
+                let schemas = spec.schemas.unwrap_or(vec![]);
+                let extensions = spec.extensions.unwrap_or(vec![]);
+
                 let source = decode_connstring(
                     kms_client.as_ref().unwrap(),
                     &key_id,
@@ -578,18 +641,17 @@ async fn cmd_dumprestore(
                     );
                 };
 
-                (source, dest)
+                (source, dest, schemas, extensions)
             }
         }
     } else {
-        (
-            source_connection_string.unwrap(),
-            if let Some(val) = destination_connection_string {
-                val
-            } else {
-                bail!("destination connection string must be provided for dump_restore command");
-            },
-        )
+        let dest = if let Some(val) = destination_connection_string {
+            val
+        } else {
+            bail!("destination connection string must be provided for dump_restore command");
+        };
+
+        (source_connection_string.unwrap(), dest, schemas, extensions)
     };
 
     run_dump_restore(
@@ -598,6 +660,8 @@ async fn cmd_dumprestore(
         pg_lib_dir,
         source_connstring,
         destination_connstring,
+        schemas,
+        extensions,
     )
     .await
 }
@@ -679,6 +743,8 @@ pub(crate) async fn main() -> anyhow::Result<()> {
                 pg_port,
                 num_cpus,
                 memory_mb,
+                schema,
+                extension,
             } => {
                 cmd_pgdata(
                     s3_client.as_ref(),
@@ -686,6 +752,8 @@ pub(crate) async fn main() -> anyhow::Result<()> {
                     args.s3_prefix.clone(),
                     spec,
                     source_connection_string,
+                    schema,
+                    extension,
                     interactive,
                     pg_port,
                     args.working_directory.clone(),
@@ -699,12 +767,16 @@ pub(crate) async fn main() -> anyhow::Result<()> {
             Command::DumpRestore {
                 source_connection_string,
                 destination_connection_string,
+                schema,
+                extension,
             } => {
                 cmd_dumprestore(
                     kms_client,
                     spec,
                     source_connection_string,
                     destination_connection_string,
+                    schema,
+                    extension,
                     args.working_directory.clone(),
                     args.pg_bin_dir,
                     args.pg_lib_dir,

pgxn/neon/file_cache.c

@@ -21,7 +21,6 @@
 #include "access/xlog.h"
 #include "funcapi.h"
 #include "miscadmin.h"
-#include "common/file_utils.h"
 #include "common/hashfn.h"
 #include "pgstat.h"
 #include "port/pg_iovec.h"
@@ -65,7 +64,7 @@
  *
  * Cache is always reconstructed at node startup, so we do not need to save mapping somewhere and worry about
  * its consistency.
- *
+
  *
  * ## Holes
  *
@@ -77,15 +76,13 @@
  * fallocate(FALLOC_FL_PUNCH_HOLE) call. The nominal size of the file doesn't
  * shrink, but the disk space it uses does.
  *
- * Each hole is tracked in a freelist. The freelist consists of two parts: a
- * fixed-size array in shared memory, and a linked chain of on-disk
- * blocks. When the in-memory array fills up, it's flushed to a new on-disk
- * chunk. If the soft limit is raised again, we reuse the holes before
- * extending the nominal size of the file.
- *
- * The in-memory freelist array is protected by 'lfc_lock', while the on-disk
- * chain is protected by a separate 'lfc_freelist_lock'.  Locking rule to
- * avoid deadlocks: always acquire lfc_freelist_lock first, then lfc_lock.
+ * Each hole is tracked by a dummy FileCacheEntry, which are kept in the
+ * 'holes' linked list. They are entered into the chunk hash table, with a
+ * special key where the blockNumber is used to store the 'offset' of the
+ * hole, and all other fields are zero. Holes are never looked up in the hash
+ * table, we only enter them there to have a FileCacheEntry that we can keep
+ * in the linked list. If the soft limit is raised again, we reuse the holes
+ * before extending the nominal size of the file.
  */
 
 /* Local file storage allocation chunk.
@@ -103,8 +100,6 @@
 #define SIZE_MB_TO_CHUNKS(size) ((uint32)((size) * MB / BLCKSZ >> lfc_chunk_size_log))
 #define BLOCK_TO_CHUNK_OFF(blkno) ((blkno) & (lfc_blocks_per_chunk-1))
 
-#define INVALID_OFFSET (0xffffffff)
-
 /*
  * Blocks are read or written to LFC file outside LFC critical section.
  * To synchronize access to such block, writer set state of such block to PENDING.
@@ -128,11 +123,11 @@ typedef struct FileCacheEntry
 	uint32		hash;
 	uint32		offset;
 	uint32		access_count;
-	dlist_node	list_node;		/* LRU list node */
+	dlist_node	list_node;		/* LRU/holes list node */
 	uint32		state[FLEXIBLE_ARRAY_MEMBER]; /* two bits per block */
 } FileCacheEntry;
 
-#define FILE_CACHE_ENTRY_SIZE MAXALIGN(offsetof(FileCacheEntry, state) + (lfc_blocks_per_chunk*2+31)/32*4)
+#define FILE_CACHE_ENRTY_SIZE MAXALIGN(offsetof(FileCacheEntry, state) + (lfc_blocks_per_chunk*2+31)/32*4)
 #define GET_STATE(entry, i) (((entry)->state[(i) / 16] >> ((i) % 16 * 2)) & 3)
 #define SET_STATE(entry, i, new_state) (entry)->state[(i) / 16] = ((entry)->state[(i) / 16] & ~(3 << ((i) % 16 * 2))) | ((new_state) << ((i) % 16 * 2))
 
@@ -166,6 +161,7 @@ typedef struct FileCacheControl
 	uint64		evicted_pages;	/* number of evicted pages */
 	dlist_head	lru;			/* double linked list for LRU replacement
 								 * algorithm */
+	dlist_head  holes;          /* double linked list of punched holes */
 	HyperLogLogState wss_estimation; /* estimation of working set size */
 	ConditionVariable cv[N_COND_VARS]; /* turnstile of condition variables */
 	PrewarmWorkerState prewarm_workers[MAX_PREWARM_WORKERS];
@@ -176,35 +172,17 @@ typedef struct FileCacheControl
 	bool   prewarm_active;
 	bool   prewarm_canceled;
 	dsm_handle prewarm_lfc_state_handle;
-
-	/*
-	 * Free list. This is large enough to hold one chunks worth of entries.
-	 */
-	uint32		freelist_size;
-	uint32		freelist_head;
-	uint32		num_free_pages;
-	uint32		free_pages[FLEXIBLE_ARRAY_MEMBER];
 } FileCacheControl;
 
-typedef struct FreeListChunk
-{
-	uint32		next;
-	uint32		num_free_pages;
-	uint32		free_pages[FLEXIBLE_ARRAY_MEMBER];
-} FreeListChunk;
-
 #define FILE_CACHE_STATE_MAGIC 0xfcfcfcfc
 
 #define FILE_CACHE_STATE_BITMAP(fcs)	((uint8*)&(fcs)->chunks[(fcs)->n_chunks])
 #define FILE_CACHE_STATE_SIZE_FOR_CHUNKS(n_chunks)	(sizeof(FileCacheState) + (n_chunks)*sizeof(BufferTag) + (((n_chunks) * lfc_blocks_per_chunk)+7)/8)
 #define FILE_CACHE_STATE_SIZE(fcs)		(sizeof(FileCacheState) + (fcs->n_chunks)*sizeof(BufferTag) + (((fcs->n_chunks) << fcs->chunk_size_log)+7)/8)
 
-#define FREELIST_ENTRIES_PER_CHUNK(c) ((c) * BLCKSZ / sizeof(uint32) - 2)
-
 static HTAB *lfc_hash;
 static int	lfc_desc = -1;
 static LWLockId lfc_lock;
-static LWLockId lfc_freelist_lock;
 static int	lfc_max_size;
 static int	lfc_size_limit;
 static int	lfc_prewarm_limit;
@@ -227,11 +205,6 @@ bool AmPrewarmWorker;
 
 #define LFC_ENABLED() (lfc_ctl->limit != 0)
 
-static bool freelist_push(uint32 offset);
-static bool freelist_prepare_pop(void);
-static uint32 freelist_pop(void);
-static bool freelist_is_empty(void);
-
 /*
  * Close LFC file if opened.
  * All backends should close their LFC files once LFC is disabled.
@@ -275,9 +248,7 @@ lfc_switch_off(void)
 		lfc_ctl->used_pages = 0;
 		lfc_ctl->limit = 0;
 		dlist_init(&lfc_ctl->lru);
-
-		lfc_ctl->freelist_head = INVALID_OFFSET;
-		lfc_ctl->num_free_pages = 0;
+		dlist_init(&lfc_ctl->holes);
 
 		/*
 		 * We need to use unlink to to avoid races in LFC write, because it is not
@@ -346,7 +317,6 @@ lfc_ensure_opened(void)
 static void
 lfc_shmem_startup(void)
 {
-	size_t		size;
 	bool		found;
 	static HASHCTL info;
 
@@ -357,19 +327,15 @@ lfc_shmem_startup(void)
 
 	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
 
-	size = offsetof(FileCacheControl, free_pages);
-	size += FREELIST_ENTRIES_PER_CHUNK(lfc_blocks_per_chunk) * sizeof(uint32);
-
-	lfc_ctl = (FileCacheControl *) ShmemInitStruct("lfc", size, &found);
+	lfc_ctl = (FileCacheControl *) ShmemInitStruct("lfc", sizeof(FileCacheControl), &found);
 	if (!found)
 	{
 		int			fd;
 		uint32		n_chunks = SIZE_MB_TO_CHUNKS(lfc_max_size);
 
 		lfc_lock = (LWLockId) GetNamedLWLockTranche("lfc_lock");
-		lfc_freelist_lock = (LWLockId) GetNamedLWLockTranche("lfc_freelist_lock");
 		info.keysize = sizeof(BufferTag);
-		info.entrysize = FILE_CACHE_ENTRY_SIZE;
+		info.entrysize = FILE_CACHE_ENRTY_SIZE;
 
 		/*
 		 * n_chunks+1 because we add new element to hash table before eviction
@@ -379,12 +345,9 @@ lfc_shmem_startup(void)
 								 n_chunks + 1, n_chunks + 1,
 								 &info,
 								 HASH_ELEM | HASH_BLOBS);
-		memset(lfc_ctl, 0, offsetof(FileCacheControl, free_pages));
+		memset(lfc_ctl, 0, sizeof(FileCacheControl));
 		dlist_init(&lfc_ctl->lru);
-
-		lfc_ctl->freelist_size = FREELIST_ENTRIES_PER_CHUNK(lfc_blocks_per_chunk);
-		lfc_ctl->freelist_head = INVALID_OFFSET;
-		lfc_ctl->num_free_pages = 0;
+		dlist_init(&lfc_ctl->holes);
 
 		/* Initialize hyper-log-log structure for estimating working set size */
 		initSHLL(&lfc_ctl->wss_estimation);
@@ -413,20 +376,13 @@ lfc_shmem_startup(void)
 static void
 lfc_shmem_request(void)
 {
-	size_t		size;
-
 #if PG_VERSION_NUM>=150000
 	if (prev_shmem_request_hook)
 		prev_shmem_request_hook();
 #endif
 
-	size = offsetof(FileCacheControl, free_pages);
-	size += FREELIST_ENTRIES_PER_CHUNK(lfc_blocks_per_chunk) * sizeof(uint32);
-	size += hash_estimate_size(SIZE_MB_TO_CHUNKS(lfc_max_size) + 1, FILE_CACHE_ENTRY_SIZE);
-
-	RequestAddinShmemSpace(size);
+	RequestAddinShmemSpace(sizeof(FileCacheControl) + hash_estimate_size(SIZE_MB_TO_CHUNKS(lfc_max_size) + 1, FILE_CACHE_ENRTY_SIZE));
 	RequestNamedLWLockTranche("lfc_lock", 1);
-	RequestNamedLWLockTranche("lfc_freelist_lock", 2);
 }
 
 static bool
@@ -479,14 +435,12 @@ lfc_change_limit_hook(int newval, void *extra)
 	if (!lfc_ctl || !is_normal_backend())
 		return;
 
-	LWLockAcquire(lfc_freelist_lock, LW_EXCLUSIVE);
 	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
 
 	/* Open LFC file only if LFC was enabled or we are going to reenable it */
 	if (newval == 0 && !LFC_ENABLED())
 	{
 		LWLockRelease(lfc_lock);
-		LWLockRelease(lfc_freelist_lock);
 		/* File should be reopened if LFC is reenabled */
 		lfc_close_file();
 		return;
@@ -495,7 +449,6 @@ lfc_change_limit_hook(int newval, void *extra)
 	if (!lfc_ensure_opened())
 	{
 		LWLockRelease(lfc_lock);
-		LWLockRelease(lfc_freelist_lock);
 		return;
 	}
 
@@ -511,14 +464,18 @@ lfc_change_limit_hook(int newval, void *extra)
 		 * returning their space to file system
 		 */
 		FileCacheEntry *victim = dlist_container(FileCacheEntry, list_node, dlist_pop_head_node(&lfc_ctl->lru));
+		FileCacheEntry *hole;
 		uint32		offset = victim->offset;
+		uint32		hash;
+		bool		found;
+		BufferTag	holetag;
 
 		CriticalAssert(victim->access_count == 0);
 #ifdef FALLOC_FL_PUNCH_HOLE
 		if (fallocate(lfc_desc, FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE, (off_t) victim->offset * lfc_blocks_per_chunk * BLCKSZ, lfc_blocks_per_chunk * BLCKSZ) < 0)
 			neon_log(LOG, "Failed to punch hole in file: %m");
 #endif
-		/* We remove the entry, and enter a hole to the freelist */
+		/* We remove the old entry, and re-enter a hole to the hash table */
 		for (int i = 0; i < lfc_blocks_per_chunk; i++)
 		{
 			bool is_page_cached = GET_STATE(victim, i) == AVAILABLE;
@@ -527,14 +484,15 @@ lfc_change_limit_hook(int newval, void *extra)
 		}
 		hash_search_with_hash_value(lfc_hash, &victim->key, victim->hash, HASH_REMOVE, NULL);
 
-		if (!freelist_push(offset))
-		{
-			/* freelist_push already logged the error */
-			lfc_switch_off();
-			LWLockRelease(lfc_lock);
-			LWLockRelease(lfc_freelist_lock);
-			return;
-		}
+		memset(&holetag, 0, sizeof(holetag));
+		holetag.blockNum = offset;
+		hash = get_hash_value(lfc_hash, &holetag);
+		hole = hash_search_with_hash_value(lfc_hash, &holetag, hash, HASH_ENTER, &found);
+		hole->hash = hash;
+		hole->offset = offset;
+		hole->access_count = 0;
+		CriticalAssert(!found);
+		dlist_push_tail(&lfc_ctl->holes, &hole->list_node);
 
 		lfc_ctl->used -= 1;
 	}
@@ -546,7 +504,6 @@ lfc_change_limit_hook(int newval, void *extra)
 	neon_log(DEBUG1, "set local file cache limit to %d", new_size);
 
 	LWLockRelease(lfc_lock);
-	LWLockRelease(lfc_freelist_lock);
 }
 
 void
@@ -1423,7 +1380,7 @@ lfc_init_new_entry(FileCacheEntry* entry, uint32 hash)
 	 * options, in order of preference:
 	 *
 	 * Unless there is no space available, we can:
-	 *  1. Use an entry from the freelist, and
+	 *  1. Use an entry from the `holes` list, and
 	 *  2. Create a new entry.
 	 * We can always, regardless of space in the LFC:
 	 *  3. evict an entry from LRU, and
@@ -1431,10 +1388,17 @@ lfc_init_new_entry(FileCacheEntry* entry, uint32 hash)
 	 */
 	if (lfc_ctl->used < lfc_ctl->limit)
 	{
-		if (!freelist_is_empty())
+		if (!dlist_is_empty(&lfc_ctl->holes))
 		{
 			/* We can reuse a hole that was left behind when the LFC was shrunk previously */
-			uint32 offset = freelist_pop();
+			FileCacheEntry *hole = dlist_container(FileCacheEntry, list_node,
+												   dlist_pop_head_node(&lfc_ctl->holes));
+			uint32 offset = hole->offset;
+			bool hole_found;
+
+			hash_search_with_hash_value(lfc_hash, &hole->key,
+										hole->hash, HASH_REMOVE, &hole_found);
+			CriticalAssert(hole_found);
 
 			lfc_ctl->used += 1;
 			entry->offset = offset;			/* reuse the hole */
@@ -1548,7 +1512,6 @@ lfc_prefetch(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber blkno,
 	hash = get_hash_value(lfc_hash, &tag);
 	cv = &lfc_ctl->cv[hash % N_COND_VARS];
 
- retry:
 	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
 
 	if (!LFC_ENABLED() || !lfc_ensure_opened())
@@ -1557,9 +1520,6 @@ lfc_prefetch(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber blkno,
 		return false;
 	}
 
-	if (!freelist_prepare_pop())
-		goto retry;
-
 	lwlsn = neon_get_lwlsn(rinfo, forknum, blkno);
 
 	if (lwlsn > lsn)
@@ -1693,7 +1653,6 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 
 	CriticalAssert(BufTagGetRelNumber(&tag) != InvalidRelFileNumber);
 
- retry:
 	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
 
 	if (!LFC_ENABLED() || !lfc_ensure_opened())
@@ -1703,9 +1662,6 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	}
 	generation = lfc_ctl->generation;
 
-	if (!freelist_prepare_pop())
-		goto retry;
-
 	/*
 	 * For every chunk that has blocks we're interested in, we
 	 * 1. get the chunk header
@@ -1867,140 +1823,6 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 	LWLockRelease(lfc_lock);
 }
 
-/**** freelist management ****/
-
-
-/*
- * Prerequisites:
- * - The caller is holding 'lfc_lock'. XXX
- */
-static bool
-freelist_prepare_pop(void)
-{
-	/*
-	 * If the in-memory freelist is empty, but there are more blocks available, load them.
-	 *
-	 * TODO: if there
-	 */
-	if (lfc_ctl->num_free_pages == 0 && lfc_ctl->freelist_head != INVALID_OFFSET)
-	{
-		uint32		freelist_head;
-		FreeListChunk *freelist_chunk;
-		size_t		bytes_read;
-
-		LWLockRelease(lfc_lock);
-		LWLockAcquire(lfc_freelist_lock, LW_EXCLUSIVE);
-
-		if (!(lfc_ctl->num_free_pages == 0 && lfc_ctl->freelist_head != INVALID_OFFSET))
-		{
-			/* someone else did the work for us while we were not holding the lock */
-			LWLockRelease(lfc_freelist_lock);
-			return false;
-		}
-
-		freelist_head = lfc_ctl->freelist_head;
-		freelist_chunk = palloc(lfc_blocks_per_chunk * BLCKSZ);
-
-		bytes_read = 0;
-		while (bytes_read < lfc_blocks_per_chunk * BLCKSZ)
-		{
-			ssize_t		rc;
-
-			rc = pread(lfc_desc, freelist_chunk, lfc_blocks_per_chunk * BLCKSZ - bytes_read, (off_t) freelist_head * lfc_blocks_per_chunk * BLCKSZ + bytes_read);
-			if (rc < 0)
-			{
-				lfc_disable("read freelist page");
-				return false;
-			}
-			bytes_read += rc;
-		}
-
-		LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
-		if (lfc_generation != lfc_ctl->generation)
-		{
-			LWLockRelease(lfc_lock);
-			return false;
-		}
-
-		Assert(lfc_ctl->freelist_head == freelist_head);
-		Assert(lfc_ctl->num_free_pages == 0);
-		lfc_ctl->freelist_head = freelist_chunk->next;
-		lfc_ctl->num_free_pages = freelist_chunk->num_free_pages;
-		memcpy(lfc_ctl->free_pages, freelist_chunk->free_pages, lfc_ctl->num_free_pages * sizeof(uint32));
-		pfree(freelist_chunk);
-
-		LWLockRelease(lfc_lock);
-		LWLockRelease(lfc_freelist_lock);
-		return false;
-	}
-
-	return true;
-}
-
-/*
- * Prerequisites:
- * - The caller is holding 'lfc_lock' and 'lfc_freelist_lock'.
- *
- * Returns 'false' on error.
- */
-static bool
-freelist_push(uint32 offset)
-{
-	Assert(lfc_ctl->freelist_size == FREELIST_ENTRIES_PER_CHUNK(lfc_blocks_per_chunk));
-	if (lfc_ctl->num_free_pages == lfc_ctl->freelist_size)
-	{
-		FreeListChunk *freelist_chunk;
-		struct iovec iov;
-		ssize_t		rc;
-
-		freelist_chunk = palloc(lfc_blocks_per_chunk * BLCKSZ);
-
-		/* write the existing entries to the chunk on disk */
-		freelist_chunk->next = lfc_ctl->freelist_head;
-		freelist_chunk->num_free_pages = lfc_ctl->num_free_pages;
-		memcpy(freelist_chunk->free_pages, lfc_ctl->free_pages, lfc_ctl->num_free_pages * sizeof(uint32));
-
-		/* Use the passed-in offset to hold the freelist chunk itself */
-		iov.iov_base = freelist_chunk;
-		iov.iov_len = lfc_blocks_per_chunk * BLCKSZ;
-		rc = pg_pwritev_with_retry(lfc_desc, &iov, 1, (off_t) offset * lfc_blocks_per_chunk * BLCKSZ);
-
-		pfree(freelist_chunk);
-
-		if (rc < 0)
-			return false;
-
-		lfc_ctl->freelist_head = offset;
-		lfc_ctl->num_free_pages = 0;
-	}
-	else
-	{
-		lfc_ctl->free_pages[lfc_ctl->num_free_pages] = offset;
-		lfc_ctl->num_free_pages++;
-	}
-	return true;
-}
-
-static uint32
-freelist_pop(void)
-{
-	uint32		result;
-
-	/* The caller should've checked that the list is not empty */
-	Assert(lfc_ctl->num_free_pages > 0);
-
-	result = lfc_ctl->free_pages[lfc_ctl->num_free_pages - 1];
-	lfc_ctl->num_free_pages--;
-
-	return result;
-}
-
-static bool
-freelist_is_empty(void)
-{
-	return lfc_ctl->num_free_pages == 0;
-}
-
 typedef struct
 {
 	TupleDesc	tupdesc;
@@ -2227,8 +2049,12 @@ local_cache_pages(PG_FUNCTION_ARGS)
 				hash_seq_init(&status, lfc_hash);
 				while ((entry = hash_seq_search(&status)) != NULL)
 				{
-					for (int i = 0; i < lfc_blocks_per_chunk; i++)
-						n_pages += GET_STATE(entry, i) == AVAILABLE;
+					/* Skip hole tags */
+					if (NInfoGetRelNumber(BufTagGetNRelFileInfo(entry->key)) != 0)
+					{
+						for (int i = 0; i < lfc_blocks_per_chunk; i++)
+							n_pages += GET_STATE(entry, i) == AVAILABLE;
+					}
 				}
 			}
 		}
@@ -2256,16 +2082,19 @@ local_cache_pages(PG_FUNCTION_ARGS)
 			{
 				for (int i = 0; i < lfc_blocks_per_chunk; i++)
 				{
-					if (GET_STATE(entry, i) == AVAILABLE)
+					if (NInfoGetRelNumber(BufTagGetNRelFileInfo(entry->key)) != 0)
 					{
-						fctx->record[n].pageoffs = entry->offset * lfc_blocks_per_chunk + i;
-						fctx->record[n].relfilenode = NInfoGetRelNumber(BufTagGetNRelFileInfo(entry->key));
-						fctx->record[n].reltablespace = NInfoGetSpcOid(BufTagGetNRelFileInfo(entry->key));
-						fctx->record[n].reldatabase = NInfoGetDbOid(BufTagGetNRelFileInfo(entry->key));
-						fctx->record[n].forknum = entry->key.forkNum;
-						fctx->record[n].blocknum = entry->key.blockNum + i;
-						fctx->record[n].accesscount = entry->access_count;
-						n += 1;
+						if (GET_STATE(entry, i) == AVAILABLE)
+						{
+							fctx->record[n].pageoffs = entry->offset * lfc_blocks_per_chunk + i;
+							fctx->record[n].relfilenode = NInfoGetRelNumber(BufTagGetNRelFileInfo(entry->key));
+							fctx->record[n].reltablespace = NInfoGetSpcOid(BufTagGetNRelFileInfo(entry->key));
+							fctx->record[n].reldatabase = NInfoGetDbOid(BufTagGetNRelFileInfo(entry->key));
+							fctx->record[n].forknum = entry->key.forkNum;
+							fctx->record[n].blocknum = entry->key.blockNum + i;
+							fctx->record[n].accesscount = entry->access_count;
+							n += 1;
+						}
 					}
 				}
 			}

test_runner/regress/test_import_pgdata.py

@@ -9,7 +9,6 @@ from pathlib import Path
 from threading import Event
 
 import psycopg2
-import psycopg2.errors
 import pytest
 from fixtures.common_types import Lsn, TenantId, TenantShardId, TimelineId
 from fixtures.fast_import import (
@@ -1070,15 +1069,41 @@ def test_fast_import_restore_to_connstring_from_s3_spec(
         return mock_kms.encrypt(KeyId=key_id, Plaintext=x)
 
     # Start source postgres and ingest data
+    vanilla_pg.configure(["shared_preload_libraries='neon,neon_utils,neon_rmgr'"])
     vanilla_pg.start()
-    vanilla_pg.safe_psql("CREATE TABLE foo (a int); INSERT INTO foo SELECT generate_series(1, 10);")
+    res = vanilla_pg.safe_psql("SHOW shared_preload_libraries;")
+    log.info(f"shared_preload_libraries: {res}")
+    res = vanilla_pg.safe_psql("SELECT name FROM pg_available_extensions;")
+    log.info(f"pg_available_extensions: {res}")
+    res = vanilla_pg.safe_psql("SELECT extname FROM pg_extension;")
+    log.info(f"pg_extension: {res}")
+
+    # Create a number of extensions, we only will dump selected ones
+    vanilla_pg.safe_psql("CREATE EXTENSION neon;")
+    vanilla_pg.safe_psql("CREATE EXTENSION neon_utils;")
+    vanilla_pg.safe_psql("CREATE EXTENSION pg_visibility;")
+
+    # Default schema is always dumped
+    vanilla_pg.safe_psql(
+        "CREATE TABLE public.foo (a int); INSERT INTO public.foo SELECT generate_series(1, 7);"
+    )
+
+    # Create a number of schemas, we only will dump selected ones
+    vanilla_pg.safe_psql("CREATE SCHEMA custom;")
+    vanilla_pg.safe_psql(
+        "CREATE TABLE custom.foo (a int); INSERT INTO custom.foo SELECT generate_series(1, 13);"
+    )
+    vanilla_pg.safe_psql("CREATE SCHEMA other;")
+    vanilla_pg.safe_psql(
+        "CREATE TABLE other.foo (a int); INSERT INTO other.foo SELECT generate_series(1, 42);"
+    )
 
     # Start target postgres
     pgdatadir = test_output_dir / "destination-pgdata"
     pg_bin = PgBin(test_output_dir, pg_distrib_dir, pg_version)
     port = port_distributor.get_port()
     with VanillaPostgres(pgdatadir, pg_bin, port) as destination_vanilla_pg:
-        destination_vanilla_pg.configure(["shared_preload_libraries='neon_rmgr'"])
+        destination_vanilla_pg.configure(["shared_preload_libraries='neon,neon_utils,neon_rmgr'"])
         destination_vanilla_pg.start()
 
         # Encrypt connstrings and put spec into S3
@@ -1092,6 +1117,8 @@ def test_fast_import_restore_to_connstring_from_s3_spec(
             "destination_connstring_ciphertext_base64": base64.b64encode(
                 destination_connstring_encrypted["CiphertextBlob"]
             ).decode("utf-8"),
+            "schemas": ["custom"],
+            "extensions": ["plpgsql", "neon"],
         }
 
         bucket = "test-bucket"
@@ -1117,9 +1144,31 @@ def test_fast_import_restore_to_connstring_from_s3_spec(
         }, f"got status: {job_status}"
         vanilla_pg.stop()
 
-        res = destination_vanilla_pg.safe_psql("SELECT count(*) FROM foo;")
+        res = destination_vanilla_pg.safe_psql("SELECT count(*) FROM public.foo;")
         log.info(f"Result: {res}")
-        assert res[0][0] == 10
+        assert res[0][0] == 7
+
+        res = destination_vanilla_pg.safe_psql("SELECT count(*) FROM custom.foo;")
+        log.info(f"Result: {res}")
+        assert res[0][0] == 13
+
+        # Check that other schema is not restored
+        with pytest.raises(psycopg2.errors.UndefinedTable):
+            destination_vanilla_pg.safe_psql("SELECT count(*) FROM other.foo;")
+
+        # Check that all schemas are listed correctly
+        res = destination_vanilla_pg.safe_psql("SELECT nspname FROM pg_namespace;")
+        log.info(f"Result: {res}")
+        schemas = [row[0] for row in res]
+        assert "other" not in schemas
+
+        # Check that only selected extensions are restored
+        res = destination_vanilla_pg.safe_psql("SELECT extname FROM pg_extension;")
+        log.info(f"Result: {res}")
+        assert len(res) == 2
+        extensions = set([str(row[0]) for row in res])
+        assert "plpgsql" in extensions
+        assert "neon" in extensions
 
 
 def test_fast_import_restore_to_connstring_error_to_s3_bad_destination(

vendor/revisions.json

@@ -1,18 +1,18 @@
 {
   "v17": [
     "17.5",
-    "8be779fd3ab9e87206da96a7e4842ef1abf04f44"
+    "10c002910447b3138e13213befca662df7cbe1d0"
   ],
   "v16": [
     "16.9",
-    "0bf96bd6d70301a0b43b0b3457bb3cf8fb43c198"
+    "94ad7e11cd43cce32f5af5674af29b3f334551a7"
   ],
   "v15": [
     "15.13",
-    "de7640f55da07512834d5cc40c4b3fb376b5f04f"
+    "cd0b534a761c18d8ef4654d4f749c63c5663215f"
   ],
   "v14": [
     "14.18",
-    "55c0d45abe6467c02084c2192bca117eda6ce1e7"
+    "b1e9959858f0529ea33d3cc5e833c0acc43f583a"
   ]
 }