refactor: add support for new sharded mapping configuration

scripts/generate-schema.py

@@ -190,6 +190,14 @@ def load_overrides():
             },
         },
     },
+    "flexibleType": {
+        "description": (
+            "pgdog FlexibleType: a bare integer, or a string (UUID or "
+            "arbitrary key). Numbers render as bare TOML integers; strings "
+            "render quoted."
+        ),
+        "type": ["integer", "string"],
+    },
 }
 
 # Properties that should use a $ref instead of being auto-generated

scripts/schema-overrides.yaml

@@ -341,8 +341,27 @@ shardedTables:
         type: string
       dataType:
         type: string
-
-# --- Sharded mappings ---
+      mapping:
+        type: array
+        description: List/range/default sharding rules for this table
+        items:
+          type: object
+          additionalProperties: false
+          required: ["shard"]
+          properties:
+            shard:
+              type: integer
+              minimum: 0
+            values:
+              type: array
+              items:
+                $ref: "#/$defs/flexibleType"
+            start:
+              $ref: "#/$defs/flexibleType"
+            end:
+              $ref: "#/$defs/flexibleType"
+
+# --- Sharded mappings (Deprecated) ---
 shardedMappings:
   type: array
   description: Sharded mapping entries
@@ -363,10 +382,12 @@ shardedMappings:
         minimum: 0
       values:
         type: array
+        items:
+          $ref: "#/$defs/flexibleType"
       start:
-        type: integer
+        $ref: "#/$defs/flexibleType"
       end:
-        type: integer
+        $ref: "#/$defs/flexibleType"
 
 # --- Omnisharded tables ---
 omnishardedTables:

templates/_helpers.tpl

@@ -97,6 +97,28 @@ as scientific notation (e.g. 8.64e+07). TOML rejects this.
 {{- if kindIs "string" . -}}{{ . }}{{- else -}}{{ int64 . }}{{- end -}}
 {{- end -}}
 
+{{/*
+Render a single FlexibleType value (integer | uuid | string) as a TOML literal.
+pgdog's FlexibleType accepts integers, UUIDs and strings. Numbers must render as
+bare TOML integers (see pgdog.intval); everything else is a string and must be
+quoted to be valid TOML.
+- Numeric inputs (float64 from YAML): render bare via int64
+- String inputs (UUIDs, arbitrary keys, quoted numbers): render quoted
+*/}}
+{{- define "pgdog.flexval" -}}
+{{- if kindIs "string" . -}}{{ . | quote }}{{- else -}}{{ int64 . }}{{- end -}}
+{{- end -}}
+
+{{/*
+Render a list of FlexibleType values as a TOML array literal.
+Each element is rendered via pgdog.flexval (numbers bare, strings quoted),
+comma-separated and wrapped in brackets, e.g. [1, 2, 3] or ["a", "b"].
+Call as: include "pgdog.flexlist" .values
+*/}}
+{{- define "pgdog.flexlist" -}}
+[{{- range $i, $v := . }}{{ if $i }}, {{ end }}{{ include "pgdog.flexval" $v }}{{- end }}]
+{{- end -}}
+
 {{/*
 Render a resources block, omitting CPU limits when noCpuLimits is true.
 Call as: include "pgdog.resources" (dict "resources" .Values.resources "noCpuLimits" .Values.noCpuLimits)

templates/config.yaml

@@ -370,6 +370,19 @@ data:
         {{- end }}
         column = {{ .column | quote }}
         data_type = {{ .dataType | quote }}
+        {{- range .mapping }}
+        [[sharded_tables.mapping]]
+        {{- if hasKey . "values" }}
+        values = {{ include "pgdog.flexlist" .values }}
+        {{- end }}
+        {{- if hasKey . "start" }}
+        start = {{ include "pgdog.flexval" .start }}
+        {{- end }}
+        {{- if hasKey . "end" }}
+        end = {{ include "pgdog.flexval" .end }}
+        {{- end }}
+        shard = {{ include "pgdog.intval" .shard }}
+        {{- end }}
         {{- end }}
 
         {{- range .Values.shardedMappings }}
@@ -379,13 +392,13 @@ data:
         kind = {{ .kind | quote }}
         shard = {{ include "pgdog.intval" .shard }}
         {{- if hasKey . "values" }}
-        values = {{ .values| toToml }}
+        values = {{ include "pgdog.flexlist" .values }}
         {{- end }}
         {{- if hasKey . "start" }}
-        start = {{ include "pgdog.intval" .start }}
+        start = {{ include "pgdog.flexval" .start }}
         {{- end }}
         {{- if hasKey . "end" }}
-        end = {{ include "pgdog.intval" .end }}
+        end = {{ include "pgdog.flexval" .end }}
         {{- end}}
         {{- end }}
 

test/values-deprecated-mappings.yaml

@@ -0,0 +1,26 @@
+# Test: the DEPRECATED top-level `shardedMappings` form still validates and
+# renders. Superseded by the `mapping` field on shardedTables entries, but kept
+# for backward compatibility. Covers list and range rules.
+databases:
+  - name: primary
+    host: postgres-primary.example.com
+    port: 5432
+
+shardedTables:
+  - database: primary
+    name: users
+    column: id
+    dataType: bigint
+
+shardedMappings:
+  - database: primary
+    column: id
+    kind: list
+    values: ["1", "2", "3"]
+    shard: 0
+  - database: primary
+    column: id
+    kind: range
+    start: 4
+    end: 100
+    shard: 1

test/values-full.yaml

@@ -102,25 +102,18 @@ shardedTables:
     name: users
     column: id
     dataType: bigint
+    mapping:
+      - values: ["1", "2", "3"]
+        shard: 0
+      - start: 4
+        end: 100
+        shard: 1
 
 shardedSchemas:
   - database: primary
     name: tenant_a
     shard: 0
 
-shardedMappings:
-  - database: primary
-    column: id
-    kind: list
-    values: ["1", "2", "3"]
-    shard: 0
-  - database: primary
-    column: id
-    kind: range
-    start: 4
-    end: 100
-    shard: 1
-
 mirrors:
   - sourceDb: primary
     destinationDb: replica

test/values-sharded-mapping.yaml

@@ -0,0 +1,39 @@
+# Test: inline `mapping` on shardedTables entries renders valid pgdog.toml.
+# Covers every FlexibleType value shape (integer / uuid / string) across all
+# three rule shapes (list / range / default) and several dataType values, to
+# confirm integers render bare and strings/uuids render quoted.
+databases:
+  - name: primary
+    host: postgres-primary.example.com
+    port: 5432
+
+shardedTables:
+  # Integer keys (bigint): list, range and default rules. Ints render bare.
+  - database: primary
+    column: tenant_id
+    dataType: bigint
+    mapping:
+      - values: [1, 2, 3]
+        shard: 0
+      - start: 100
+        end: 200
+        shard: 1
+      - shard: 2
+  # UUID keys: list and range rules. UUIDs are strings and must render quoted.
+  - database: primary
+    column: tenant_uuid
+    dataType: uuid
+    mapping:
+      - values: ["11111111-1111-1111-1111-111111111111"]
+        shard: 0
+      - start: "aaaaaaaa-0000-0000-0000-000000000000"
+        end: "bbbbbbbb-0000-0000-0000-000000000000"
+        shard: 1
+  # String keys (varchar): list of arbitrary strings, quoted.
+  - database: primary
+    column: region
+    dataType: varchar
+    mapping:
+      - values: ["us-east", "eu-west"]
+        shard: 0
+      - shard: 1

values.yaml

@@ -238,15 +238,32 @@ shardedSchemas: []
 
 # shardedTables contains the list of sharded table entries in pgdog.toml
 # Each entry requires: database, column and dataType; name is optional
+# An entry may also carry a list/range/default mapping via the optional `mapping` field.
+# Each mapping rule has a target `shard`; the rule type is inferred from the fields set:
+#   - `values` -> list rule
+#   - `start` and/or `end` -> range rule
+#   - `shard` only -> default (catch-all) rule
 # Example:
 # shardedTables:
 #   - database: "prod"
 #     name: "users"
 #     column: "id"
 #     dataType: "bigint"
+#   - database: "prod"
+#     column: "tenant_id"
+#     dataType: "bigint"
+#     mapping:
+#       - values: [1, 5, 1_000]
+#         shard: 0
+#       - start: 100
+#         end: 200
+#         shard: 1
+#       - shard: 2
 shardedTables: []
 
 # shardedMappings contains the list of sharded mapping entries in pgdog.toml
+# DEPRECATED: prefer the `mapping` field on shardedTables entries above. This
+# legacy top-level form is still supported for backward compatibility.
 # Each entry requires: database, column, kind, shard; values, start and end are optional
 # Example:
 # shardedMappings: