join

Author: brunovcosta

abstra_json_sql/apply.py

@@ -511,7 +511,11 @@ def apply_select_fields(fields: List[SelectField], data: List[dict], ctx: dict):
         result_row = {}
         for field in fields:
             if isinstance(field.expression, Wildcard):
-                result_row.update(row)
+                for key, value in row.items():
+                    assert isinstance(key, str), "Key should be a string"
+                    parts = key.split(".")
+                    last_part = parts[-1]
+                    result_row[last_part] = value
             elif isinstance(field.expression, Expression):
                 result_row[field_name(field)] = apply_expression(
                     field.expression, {**ctx, **row}
@@ -522,6 +526,11 @@ def apply_select_fields(fields: List[SelectField], data: List[dict], ctx: dict):
     return result
 
 
+def add_scope_to_keys(prefix: str, data: dict) -> dict:
+    change = {f"{prefix}.{key}": value for key, value in data.items()}
+    return {**data, **change}
+
+
 def apply_from(
     from_part: Optional[From], tables: ITablesSnapshot, ctx: dict
 ) -> List[dict]:
@@ -537,11 +546,23 @@ def apply_from(
             if not join_table:
                 raise ValueError(f"Table {join.table} not found")
             data = [
-                {**row, **join_row}
+                {
+                    **add_scope_to_keys(table.name, row),
+                    **add_scope_to_keys(join_table.name, join_row),
+                }
                 for row in data
                 for join_row in join_table.data
-                if apply_expression(join.on, {**ctx, **row, **join_row})
+                if apply_expression(
+                    join.on,
+                    {
+                        **ctx,
+                        **add_scope_to_keys(table.name, row),
+                        **add_scope_to_keys(join_table.name, join_row),
+                    },
+                )
             ]
+    else:
+        data = [{**add_scope_to_keys(table.name, row)} for row in data]
     return data
 
 

abstra_json_sql/eval_test.py

@@ -467,6 +467,72 @@ def test_group_by(self):
             ],
         )
 
+    def test_select_wildcard(self):
+        code = "select * from bar"
+        tables = InMemoryTables(
+            tables=[
+                Table(
+                    name="bar",
+                    columns=[Column(name="foo", type="text")],
+                    data=[
+                        {"foo": "a"},
+                        {"foo": "b"},
+                        {"foo": None},
+                        {"foo": "c"},
+                    ],
+                )
+            ],
+        )
+        ctx = {}
+        result = eval_sql(code=code, tables=tables, ctx=ctx)
+        self.assertEqual(
+            result,
+            [
+                {"foo": "a"},
+                {"foo": "b"},
+                {"foo": None},
+                {"foo": "c"},
+            ],
+        )
+
+    def test_join(self):
+        code = "select a.foo, b.bar from a join b on a.id = b.a_id"
+        tables = InMemoryTables(
+            tables=[
+                Table(
+                    name="a",
+                    columns=[
+                        Column(name="id", type="int"),
+                        Column(name="foo", type="text"),
+                    ],
+                    data=[
+                        {"id": 1, "foo": "a1"},
+                        {"id": 2, "foo": "a2"},
+                    ],
+                ),
+                Table(
+                    name="b",
+                    columns=[
+                        Column(name="a_id", type="int"),
+                        Column(name="bar", type="text"),
+                    ],
+                    data=[
+                        {"a_id": 1, "bar": "b1"},
+                        {"a_id": 2, "bar": "b2"},
+                    ],
+                ),
+            ],
+        )
+        ctx = {}
+        result = eval_sql(code=code, tables=tables, ctx=ctx)
+        self.assertEqual(
+            result,
+            [
+                {"foo": "a1", "bar": "b1"},
+                {"foo": "a2", "bar": "b2"},
+            ],
+        )
+
     def test_complete(self):
         code = "\n".join(
             [
@@ -499,31 +565,3 @@ def test_complete(self):
         ctx = {}
         result = eval_sql(code=code, tables=tables, ctx=ctx)
         self.assertEqual(result, [{"foo": 3, "count": 2}])
-
-    def test_select_wildcard(self):
-        code = "select * from bar"
-        tables = InMemoryTables(
-            tables=[
-                Table(
-                    name="bar",
-                    columns=[Column(name="foo", type="text")],
-                    data=[
-                        {"foo": "a"},
-                        {"foo": "b"},
-                        {"foo": None},
-                        {"foo": "c"},
-                    ],
-                )
-            ],
-        )
-        ctx = {}
-        result = eval_sql(code=code, tables=tables, ctx=ctx)
-        self.assertEqual(
-            result,
-            [
-                {"foo": "a"},
-                {"foo": "b"},
-                {"foo": None},
-                {"foo": "c"},
-            ],
-        )

abstra_json_sql/field_name.py

@@ -16,7 +16,9 @@ def expression_name(expression: Expression):
     Get the field name from an Expression object.
     """
     if isinstance(expression, NameExpression):
-        return expression.name
+        name_parts = expression.name.split(".")
+        last_part = name_parts[-1]
+        return last_part
     elif isinstance(expression, FunctionCallExpression):
         return expression.name
     else:

abstra_json_sql/lexer.py

@@ -32,7 +32,7 @@ def extract_space(code: str) -> str:
 
 
 def start_with_name(code: str):
-    return code[0].isalnum() or code[0] == "_"
+    return code[0].isalnum() or code[0] == "_" or code[0] == "."
 
 
 def extract_name(code: str):