feat: add 53-bit unique ID generator & database-generated sharded sequences (#858)

### Smaller unique ID

_This feature is still experimental and subject to change_.

Add support for 53-bit unique ID. This makes them smaller and
"javascript-safe" in case identifiers are exposed as integers to the app
API.

This is configurable thourough `pgdog.toml`:

```toml
[general]
unique_id_function = "compact" # or "standard", which is default
```

Trade-offs: it only allows 64 pgdog nodes and 64,000 IDs generated _per
second_ per node. This is considerably lower than the standard 64-bit
unique ID. It should be used for omnisharded tables only with a low
write frequency.

### Database-generated sharded sequences

Add the ability to generate cross-shard unique IDs using sequences. This
produces smaller starting integers, but can only work with
direct-to-shard `INSERT`s. For omnisharded tables, the unique ID
function should continue to be used.

This is configurable through `pgdog.toml`:

```toml
[rewrite]
primary_key = "rewrite_omni" # Will only inject Unique ID to omnisharded table inserts, leaving others to be
                                                # generated by the database instead.
```

The sharded sequence is installed automatically on all tables with a
`BIGINT` primary key when running `pgdog setup` or `SETUP SCHEMA` via
the admin DB.

### Identity columns

`GENERATED ... WITH IDENTITY` is now removed entirely during schema
sync. This constraint blocks our implementation of replication and
doesn't work with sharded databases, incl. with our unique ID generator
because it prevents the column value from being inserted by the query.

### Smaller features

- Add `RESET PREPARED` admin command which evicts all unused prepared
statements from the global cache
- log any notice/warning received from Postgres when executing queries
via `execute_checked`



### Bug fixes

- `unique_id_min` setting was being ignored by the unique ID generator.
This didn't affect any deployments because unique IDs start very large
already.
- Admin command `SETUP SCHEMA` was applied to all database/user pairs;
it's now applied to the `schema_owner` only

Author: levkk

.schema/pgdog.schema.json

@@ -102,6 +102,7 @@
         "tls_verify": "prefer",
         "two_phase_commit": false,
         "two_phase_commit_auto": null,
+        "unique_id_function": "standard",
         "unique_id_min": 0,
         "workers": 2
       }
@@ -1025,6 +1026,11 @@
           ],
           "default": null
         },
+        "unique_id_function": {
+          "description": "Unique ID generation function.",
+          "$ref": "#/$defs/UniqueIdFunction",
+          "default": "standard"
+        },
         "unique_id_min": {
           "description": "Minimum ID for unique ID generator.",
           "type": "integer",
@@ -1495,6 +1501,11 @@
           "description": "Automatically rewrite the query and execute it.",
           "type": "string",
           "const": "rewrite"
+        },
+        {
+          "description": "Rewrite only for omnisharded tables.",
+          "type": "string",
+          "const": "rewrite_omni"
         }
       ]
     },
@@ -1817,6 +1828,20 @@
         }
       ]
     },
+    "UniqueIdFunction": {
+      "oneOf": [
+        {
+          "description": "Standard 64-bit function using the entire 64-bit range.",
+          "type": "string",
+          "const": "standard"
+        },
+        {
+          "description": "Compact function using the leftest 53-bit range, making it\nJavaScript-safe, so you can pass it as an integer directly\nto the frontend apps.\n\nThe year is 2026 and JavaScript continues to be a pain in the ass.",
+          "type": "string",
+          "const": "compact"
+        }
+      ]
+    },
     "Vector": {
       "type": "object",
       "properties": {

integration/copy_data/pgdog.toml

@@ -9,6 +9,11 @@ name = "source"
 host = "127.0.0.1"
 database_name = "pgdog"
 
+[[sharded_tables]]
+database = "source"
+column = "tenant_id"
+data_type = "bigint"
+
 #
 # Reshard 0 -> 2
 #

integration/copy_data/setup.sql

@@ -45,6 +45,7 @@ CREATE TABLE copy_data.with_identity(
     tenant_id BIGINT NOT NULL
 );
 
+
 -- Omni (non-sharded) tables: no tenant_id column.
 CREATE TABLE IF NOT EXISTS copy_data.countries (
     id BIGSERIAL PRIMARY KEY,
@@ -65,6 +66,12 @@ CREATE TABLE IF NOT EXISTS copy_data.categories (
     parent_id INT
 );
 
+CREATE TABLE copy_data.settings (
+    id BIGSERIAL PRIMARY KEY,
+    setting_name TEXT NOT NULL UNIQUE,
+    setting_value TEXT NOT NULL
+);
+
 DROP PUBLICATION IF EXISTS pgdog;
 CREATE PUBLICATION pgdog FOR TABLES IN SCHEMA copy_data;
 

integration/pgdog.toml

@@ -24,7 +24,7 @@ query_parser_engine = "pg_query_raw"
 system_catalogs = "omnisharded_sticky"
 reload_schema_on_ddl = false
 #idle_healthcheck_delay = 50000000
-
+unique_id_function = "standard"
 
 [memory]
 net_buffer = 8096

integration/rust/src/setup.rs

@@ -45,9 +45,19 @@ pub async fn connections_sqlx() -> Vec<Pool<Postgres>> {
 }
 
 pub async fn connection_sqlx_direct() -> Pool<Postgres> {
+    connection_sqlx_direct_db("pgdog").await
+}
+
+pub async fn connection_sqlx_direct_db(name: &str) -> Pool<Postgres> {
     PgPoolOptions::new()
         .max_connections(1)
-        .connect("postgres://pgdog:pgdog@127.0.0.1:5432/pgdog?application_name=sqlx_direct")
+        .connect(
+            format!(
+                "postgres://pgdog:pgdog@127.0.0.1:5432/{}?application_name=sqlx_direct",
+                name
+            )
+            .as_str(),
+        )
         .await
         .unwrap()
 }

integration/rust/tests/integration/mod.rs

@@ -22,6 +22,7 @@ pub mod prepared;
 pub mod reload;
 pub mod reset;
 pub mod rewrite;
+pub mod rewrite_omni;
 pub mod savepoint;
 pub mod set_in_transaction;
 pub mod set_sharding_key;

integration/rust/tests/integration/rewrite_omni.rs

@@ -0,0 +1,131 @@
+use rust::setup::*;
+use sqlx::Executor;
+
+#[tokio::test]
+async fn test_omni_only_pk_rewrite() {
+    let shard_0 = connection_sqlx_direct_db("shard_0").await;
+    let shard_1 = connection_sqlx_direct_db("shard_1").await;
+
+    for (shard, pool) in [&shard_0, &shard_1].iter().enumerate() {
+        pool.execute(
+            "DROP TABLE IF EXISTS public.test_omni_rewrite_pk_omni, test_omni_rewrite_pk_sharded CASCADE",
+        )
+        .await
+        .unwrap();
+
+        pool.execute("CREATE TABLE IF NOT EXISTS public.test_omni_rewrite_pk_omni(id BIGSERIAL PRIMARY KEY, value TEXT NOT NULL)").await.unwrap();
+        pool.execute("CREATE TABLE IF NOT EXISTS public.test_omni_rewrite_pk_sharded(id BIGSERIAL PRIMARY KEY, customer_id BIGINT NOT NULL, value TEXT NOT NULL)").await.unwrap();
+
+        pool.execute(
+            "SELECT pgdog.install_sharded_sequence('public', 'test_omni_rewrite_pk_sharded', 'id')",
+        )
+        .await
+        .unwrap();
+
+        // Configure sharding.
+        let mut t = pool.begin().await.unwrap();
+        t.execute("DELETE FROM pgdog.config").await.unwrap();
+        t.execute(
+            format!(
+                "INSERT INTO pgdog.config (shard, shards) VALUES ({}, 2)",
+                shard
+            )
+            .as_str(),
+        )
+        .await
+        .unwrap();
+        t.commit().await.unwrap();
+    }
+
+    let sharded = connections_sqlx().await.pop().unwrap();
+    let admin = admin_sqlx().await;
+
+    // This will reload the schema as well.
+    admin
+        .execute("SET rewrite_primary_key TO 'rewrite_omni'")
+        .await
+        .unwrap();
+
+    let starting_id: (i64,) = sqlx::query_as("SELECT pgdog.unique_id()")
+        .fetch_one(&sharded)
+        .await
+        .unwrap();
+
+    for run in 0..25 {
+        let omni_id: (i64,) = sqlx::query_as(
+            "INSERT INTO public.test_omni_rewrite_pk_omni (value) VALUES ($1) RETURNING id",
+        )
+        .bind(format!("test_{}", run))
+        .fetch_one(&sharded)
+        .await
+        .unwrap();
+
+        assert!(
+            omni_id.0 > starting_id.0,
+            "omni ID should be unique_id, but got {}",
+            omni_id.0
+        );
+
+        let sharded_id: (i64,) = sqlx::query_as(
+            "INSERT INTO public.test_omni_rewrite_pk_sharded (customer_id, value) VALUES ($1, $2) RETURNING id",
+    )
+        .bind(run as i64)
+        .bind(format!("test_{}", run))
+    .fetch_one(&sharded)
+    .await
+    .unwrap();
+
+        assert!(
+            sharded_id.0 < omni_id.0,
+            "sharded ID should not be unique_id, but got {}",
+            sharded_id.0
+        );
+    }
+
+    sharded.close().await;
+    let sharded = connections_sqlx().await.pop().unwrap();
+
+    // Re-enable sharded unique ID.
+    admin
+        .execute("SET rewrite_primary_key TO 'rewrite'")
+        .await
+        .unwrap();
+
+    // The rewrite is cached in prepared statements
+    // and the query cache, so we need to be careful to rest it
+    // _after_ the client disconnected. In production, changing this setting
+    // definitely requires a restart.
+    admin.execute("RESET QUERY_CACHE").await.unwrap();
+    admin.execute("RESET PREPARED").await.unwrap();
+
+    for run in 25..50 {
+        let omni_id: (i64,) = sqlx::query_as(
+            "INSERT INTO public.test_omni_rewrite_pk_omni (value) VALUES ($1) RETURNING id",
+        )
+        .bind(format!("test_{}", run))
+        .fetch_one(&sharded)
+        .await
+        .unwrap();
+
+        assert!(
+            omni_id.0 > starting_id.0,
+            "omni ID should be unique_id, but got {}",
+            omni_id.0
+        );
+
+        let sharded_id: (i64,) = sqlx::query_as(
+            "INSERT INTO public.test_omni_rewrite_pk_sharded (customer_id, value) VALUES ($1, $2) RETURNING id",
+    )
+        .bind(run as i64)
+        .bind(format!("test_{}", run))
+    .fetch_one(&sharded)
+    .await
+    .unwrap();
+
+        assert!(
+            sharded_id.0 > omni_id.0,
+            "sharded ID should be unique_id, but got {}",
+            sharded_id.0
+        );
+    }
+}

integration/schema_sync/dev.sh

@@ -81,12 +81,25 @@ EOF
 
 diff source.sql destination.sql > diff.txt || true
 
-# Extract column name and type from diff lines, ignoring everything else
-# This normalizes across different PG versions and constraint syntaxes
-ACTUAL_CONVERSIONS=$(grep '^[<>]' diff.txt | \
-    grep -E '\b(integer|bigint)\b' | \
-    sed -E 's/.*[[:space:]]([a-z_]+)[[:space:]]+(integer|bigint).*/\1 \2/' | \
-    sort -u)
+# Extract integer -> bigint conversions from the diff
+# 1. Get removed (< integer) and added (> bigint) column lines
+# 2. Only keep columns that appear as integer in source AND bigint in destination
+REMOVED_INT=$(grep '^<' diff.txt | \
+    sed -E 's/.*[[:space:]]([a-z_]+)[[:space:]]+integer\b.*/\1/' | \
+    grep -E '^[a-z_]+$' | sort -u)
+
+ADDED_BIGINT=$(grep '^>' diff.txt | \
+    sed -E 's/.*[[:space:]]([a-z_]+)[[:space:]]+bigint\b.*/\1/' | \
+    grep -E '^[a-z_]+$' | sort -u)
+
+# Columns that changed from integer to bigint
+CONVERTED=$(comm -12 <(echo "$REMOVED_INT") <(echo "$ADDED_BIGINT"))
+
+# Build the expected format: column_name integer \n column_name bigint
+ACTUAL_CONVERSIONS=$(echo "$CONVERTED" | while read col; do
+    echo "$col integer"
+    echo "$col bigint"
+done | sort -u)
 
 EXPECTED_SORTED=$(echo "$EXPECTED_CONVERSIONS" | sort -u)
 

integration/schema_sync/pgdog.toml

@@ -14,3 +14,7 @@ database_name = "pgdog1"
 name = "destination"
 host = "127.0.0.1"
 database_name = "pgdog2"
+
+[[sharded_tables]]
+database = "destination"
+column = "user_id"

pgdog-config/src/general.rs

@@ -8,6 +8,7 @@ use std::str::FromStr;
 use std::time::Duration;
 
 use crate::pooling::ConnectionRecovery;
+use crate::UniqueIdFunction;
 use crate::{
     CopyFormat, CutoverTimeoutAction, LoadSchema, QueryParserEngine, QueryParserLevel,
     SystemCatalogsBehavior,
@@ -545,6 +546,10 @@ pub struct General {
     #[serde(default)]
     pub unique_id_min: u64,
 
+    /// Unique ID generation function.
+    #[serde(default)]
+    pub unique_id_function: UniqueIdFunction,
+
     /// Changes how system catalog tables (like `pg_database`, `pg_class`, etc.) are treated by the query router.
     ///
     /// _Default:_ `omnisharded_sticky`
@@ -718,6 +723,7 @@ impl Default for General {
             cutover_timeout: Self::cutover_timeout(),
             cutover_timeout_action: Self::cutover_timeout_action(),
             cutover_save_config: bool::default(),
+            unique_id_function: Self::unique_id_function(),
         }
     }
 }
@@ -817,6 +823,10 @@ impl General {
         Self::env_or_default("PGDOG_BAN_REPLICA_LAG_BYTES", i64::MAX as u64)
     }
 
+    fn unique_id_function() -> UniqueIdFunction {
+        Self::env_enum_or_default("PGDOG_UNIQUE_ID_FUNCTION")
+    }
+
     fn cutover_replication_lag_threshold() -> u64 {
         Self::env_or_default("PGDOG_CUTOVER_REPLICATION_LAG_THRESHOLD", 0)
         // 0 bytes