diff --git a/core/wren/src/wren/connector/redshift.py b/core/wren/src/wren/connector/redshift.py index 918c8a2244..4ac76b4940 100644 --- a/core/wren/src/wren/connector/redshift.py +++ b/core/wren/src/wren/connector/redshift.py @@ -52,7 +52,7 @@ def query(self, sql: str, limit: int | None = None) -> pa.Table: else: # Unlimited path also rejects trailing ``;`` for single statements # depending on driver/session settings — strip for consistency. - sql = _strip_trailing_semicolon(sql) + sql = strip_trailing_semicolon(sql) with closing(self.connection.cursor()) as cursor: cursor.execute(sql) cols = [desc[0] for desc in cursor.description] diff --git a/core/wren/src/wren/memory/seed_queries.py b/core/wren/src/wren/memory/seed_queries.py index cfc8a2e26d..d1f7ec4aad 100644 --- a/core/wren/src/wren/memory/seed_queries.py +++ b/core/wren/src/wren/memory/seed_queries.py @@ -31,13 +31,19 @@ def generate_seed_queries(manifest: dict) -> list[dict]: """Return a list of {"nl": ..., "sql": ...} seed pairs.""" pairs: list[dict] = [] + models = manifest.get("models", []) or [] + if not isinstance(models, list): + models = [] model_layers = { model["name"]: _prop_value(model, "dbtLayer", "dbt_layer") - for model in manifest.get("models", []) + for model in models + if isinstance(model, dict) and model.get("name") is not None } relationship_keys = _relationship_key_columns(manifest) - for model in manifest.get("models", []): + for model in models: + if not isinstance(model, dict) or model.get("name") is None: + continue if model_layers.get(model["name"]) == "raw": continue pairs.extend( @@ -46,10 +52,14 @@ def generate_seed_queries(manifest: dict) -> list[dict]: ) ) - for rel in manifest.get("relationships", []): - pair = _relationship_seed(rel, model_layers) - if pair: - pairs.append(pair) + rels = manifest.get("relationships", []) or [] + if isinstance(rels, list): + for rel in rels: + if not isinstance(rel, dict): + continue + pair = _relationship_seed(rel, model_layers) + if pair: + pairs.append(pair) return pairs @@ -58,8 +68,11 @@ def _model_seeds( model: dict, relationship_keys: frozenset[str] = frozenset() ) -> list[dict]: name = model["name"] + columns = model.get("columns", []) or [] + if not isinstance(columns, list): + columns = [] columns = [ - c for c in (model.get("columns") or []) if isinstance(c, dict) and c.get("name") + c for c in columns if isinstance(c, dict) and isinstance(c.get("name"), str) ] primary_keys = _primary_key_columns(model) pairs = [] @@ -166,7 +179,12 @@ def _relationship_key_columns(manifest: dict) -> dict[str, frozenset[str]]: of aggregation seeds. """ accum: dict[str, set[str]] = {} - for rel in manifest.get("relationships", []): + rels = manifest.get("relationships", []) or [] + if not isinstance(rels, list): + return {} + for rel in rels: + if not isinstance(rel, dict): + continue condition = rel.get("condition") or "" try: tree = sqlglot.parse_one(condition) diff --git a/core/wren/tests/unit/test_seed_queries_nonduct.py b/core/wren/tests/unit/test_seed_queries_nonduct.py new file mode 100644 index 0000000000..30248a5da4 --- /dev/null +++ b/core/wren/tests/unit/test_seed_queries_nonduct.py @@ -0,0 +1,22 @@ +from wren.memory.seed_queries import generate_seed_queries + + +def test_generate_seed_queries_skips_nonduct_models_and_columns(): + pairs = generate_seed_queries( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": "amount", "type": "double"}, + "bad", + {"type": "int"}, + ], + }, + "nope", + ], + "relationships": ["x"], + } + ) + assert any("orders" in p["nl"] for p in pairs) + assert all(isinstance(p.get("sql"), str) for p in pairs)