Merge pull request #678 from splitgraph/feature/postgres-fdw-extra-jsonschema

Add `postgres_fdw` options to its JSONSchema

Author: mildbyte

splitgraph/hooks/data_source/fdw.py

@@ -676,10 +676,34 @@ def get_description(cls) -> str:
                 "description": "Schema name to import data from",
                 "default": "public",
             },
+            "use_remote_estimate": {
+                "type": "boolean",
+                "default": False,
+                "title": "Use remote EXPLAIN",
+                "description": "Issue remote EXPLAIN commands to obtain cost estimates",
+            },
+            "fetch_size": {
+                "type": "integer",
+                "title": "Fetch size",
+                "description": "Number of rows from the remote server to get in each fetch operation",
+                "default": 10000,
+            },
         },
         "required": ["host", "port", "dbname", "remote_schema"],
     }
 
+    table_params_schema = merge_jsonschema(
+        ForeignDataWrapperDataSource.table_params_schema,
+        {
+            "type": "object",
+            "properties": {
+                "table_name": {"type": "string", "title": "Table name"},
+                "schema_name": {"type": "string", "title": "Schema name"},
+            },
+            "required": [],
+        },
+    )
+
     commandline_help: str = """Mount a Postgres database.
 
 Mounts a schema on a remote Postgres database as a set of foreign tables locally."""
@@ -700,6 +724,8 @@ def get_server_options(self):
             "host": self.params["host"],
             "port": str(self.params["port"]),
             "dbname": self.params["dbname"],
+            "use_remote_estimate": "true" if self.params.get("use_remote_estimate") else "false",
+            "fetch_size": int(self.params.get("fetch_size", 10000)),
             **self.params.get("extra_server_args", {}),
         }
 
@@ -708,7 +734,7 @@ def get_user_options(self):
 
     def get_table_options(self, table_name: str, tables: Optional[TableInfo] = None):
         options = super().get_table_options(table_name, tables)
-        options["schema_name"] = self.params["remote_schema"]
+        options["schema_name"] = options.get("schema_name", self.params["remote_schema"])
         return options
 
     def get_fdw_name(self):