Add Supabase pooler shard option

Author: Maxteabag

sqlit/cli.py

@@ -294,6 +294,10 @@ def main() -> int:
     )
     parser.add_argument("--supabase-region", help="Supabase region (temporary connection)")
     parser.add_argument("--supabase-project-id", help="Supabase project id (temporary connection)")
+    parser.add_argument(
+        "--supabase-aws-shard",
+        help="Supabase pooler shard prefix (temporary connection, e.g. aws-0, aws-1)",
+    )
     parser.add_argument(
         "--settings",
         metavar="PATH",

sqlit/domains/connections/providers/supabase/adapter.py

@@ -20,8 +20,16 @@ def supports_multiple_databases(self) -> bool:
     def connect(self, config: ConnectionConfig) -> Any:
         region = config.get_option("supabase_region", "")
         project_id = config.get_option("supabase_project_id", "")
+        shard = config.get_option("supabase_aws_shard", "aws-0")
+        shard = str(shard).strip() or "aws-0"
+        if shard.startswith("aws") and not shard.startswith("aws-"):
+            suffix = shard.removeprefix("aws")
+            if suffix.isdigit():
+                shard = f"aws-{suffix}"
+        elif shard.isdigit():
+            shard = f"aws-{shard}"
         transformed = config.with_endpoint(
-            host=f"aws-0-{region}.pooler.supabase.com",
+            host=f"{shard}-{region}.pooler.supabase.com",
             port="5432",
             username=f"postgres.{project_id}",
             database="postgres",

sqlit/domains/connections/providers/supabase/schema.py

@@ -26,6 +26,13 @@
             placeholder="abcdefghijklmnop",
             required=True,
         ),
+        SchemaField(
+            name="supabase_aws_shard",
+            label="AWS Shard",
+            placeholder="aws-0",
+            default="aws-0",
+            description="Pooler shard prefix (e.g. aws-0, aws-1)",
+        ),
         SchemaField(
             name="password",
             label="Password",