Merge pull request #674 from splitgraph/bugfix/fix-socrata-data-source

Fix Socrata previews returning null

Author: mildbyte

splitgraph/hooks/data_source/fdw.py

@@ -235,9 +235,10 @@ def _preview_table(self, schema: str, table: str, limit: int = 10) -> List[Dict[
         result_json = cast(
             List[Dict[str, Any]],
             self.engine.run_sql(
-                SQL("SELECT row_to_json(t.*) FROM {}.{} t LIMIT %s").format(
-                    Identifier(schema), Identifier(table)
-                ),
+                SQL(
+                    "WITH p AS MATERIALIZED(SELECT * FROM {}.{} LIMIT %s) "
+                    "SELECT row_to_json(p.*) FROM p"
+                ).format(Identifier(schema), Identifier(table)),
                 (limit,),
                 return_shape=ResultShape.MANY_ONE,
             ),

test/splitgraph/ingestion/test_socrata.py

@@ -423,3 +423,7 @@ def test_socrata_data_source_introspection_smoke(local_engine_empty):
     schema, params = result["some_table"]
     assert len(schema) > 1
     assert params == {"socrata_id": "8wbx-tsch"}
+
+    preview_result = data_source.preview(tables=result)
+    assert len(preview_result["some_table"]) == 10
+    assert preview_result["some_table"][0][":id"] is not None