finish implementation

hogql_parser/CONTRIBUTING.md

@@ -68,3 +68,9 @@ The three pages below are must-reads though. They're key to writing production-r
     ```bash
     pytest posthog/hogql/
     ```
+
+## How to install dependencies on Ubuntu
+
+Antlr runtime provided in Ubuntu packages might be of older version, which results in compilation error.
+
+In that case run commands from [this step](https://github.com/PostHog/posthog/blob/4fba6a63e351131fdb27b85e7ba436446fdb3093/.github/actions/run-backend-tests/action.yml#L100).

hogql_parser/parser.cpp

@@ -857,9 +857,6 @@ class HogQLParseTreeConverter : public HogQLParserBaseVisitor {
   VISIT_UNSUPPORTED(JoinOpCross)
 
   VISIT(JoinConstraintClause) {
-    if (ctx->USING()) {
-      throw NotImplementedError("Unsupported: JOIN ... USING");
-    }
     PyObject* column_expr_list = visitAsPyObject(ctx->columnExprList());
     Py_ssize_t column_expr_list_size = PyList_Size(column_expr_list);
     if (column_expr_list_size == -1) {
@@ -872,7 +869,7 @@ class HogQLParseTreeConverter : public HogQLParserBaseVisitor {
     }
     PyObject* expr = Py_NewRef(PyList_GET_ITEM(column_expr_list, 0));
     Py_DECREF(column_expr_list);
-    RETURN_NEW_AST_NODE("JoinConstraint", "{s:N}", "expr", expr);
+    RETURN_NEW_AST_NODE("JoinConstraint", "{s:N,s:s}", "expr", expr, "constraint_type", ctx->USING() ? "USING" : "ON");
   }
 
   VISIT(SampleClause) {

posthog/hogql/ast.py

@@ -570,8 +570,8 @@ class Call(Expr):
 
 @dataclass(kw_only=True)
 class JoinConstraint(Expr):
-    join_type: Literal["ON", "USING"]
     expr: Expr
+    constraint_type: Literal["ON", "USING"]
 
 
 @dataclass(kw_only=True)

posthog/hogql/parser.py

@@ -369,7 +369,7 @@ def visitJoinConstraintClause(self, ctx: HogQLParser.JoinConstraintClauseContext
         column_expr_list = self.visit(ctx.columnExprList())
         if len(column_expr_list) != 1:
             raise NotImplementedError(f"Unsupported: JOIN ... ON with multiple expressions")
-        return ast.JoinConstraint(expr=column_expr_list[0], join_type="USING" if ctx.USING() else "ON")
+        return ast.JoinConstraint(expr=column_expr_list[0], constraint_type="USING" if ctx.USING() else "ON")
 
     def visitSampleClause(self, ctx: HogQLParser.SampleClauseContext):
         ratio_expressions = ctx.ratioExpr()

posthog/hogql/printer.py

@@ -452,7 +452,7 @@ def visit_join_expr(self, node: ast.JoinExpr) -> JoinExprResponse:
                 join_strings.append(sample_clause)
 
         if node.constraint is not None:
-            join_strings.append(f"{node.constraint.join_type} {self.visit(node.constraint)}")
+            join_strings.append(f"{node.constraint.constraint_type} {self.visit(node.constraint)}")
 
         return JoinExprResponse(printed_sql=" ".join(join_strings), where=extra_where)
 

posthog/hogql/resolver.py

@@ -312,7 +312,7 @@ def visit_join_expr(self, node: ast.JoinExpr):
                 node_type = node_table_type
 
             node = cast(ast.JoinExpr, clone_expr(node))
-            if node.constraint and node.constraint.join_type == "USING":
+            if node.constraint and node.constraint.constraint_type == "USING":
                 # visit constraint before adding the table to not have ambiguous names
                 # TODO: add a test for it
                 node.constraint = self.visit(node.constraint)
@@ -336,7 +336,7 @@ def visit_join_expr(self, node: ast.JoinExpr):
             if is_global:
                 node.next_join.join_type = "GLOBAL JOIN"
 
-            if node.constraint and node.constraint.join_type == "ON":
+            if node.constraint and node.constraint.constraint_type == "ON":
                 node.constraint = self.visit(node.constraint)
             node.sample = self.visit(node.sample)
 
@@ -348,7 +348,7 @@ def visit_join_expr(self, node: ast.JoinExpr):
 
         elif isinstance(node.table, ast.SelectQuery) or isinstance(node.table, ast.SelectUnionQuery):
             node = cast(ast.JoinExpr, clone_expr(node))
-            if node.constraint and node.constraint.join_type == "USING":
+            if node.constraint and node.constraint.constraint_type == "USING":
                 # visit constraint before adding the table to not have ambiguous names
                 node.constraint = self.visit(node.constraint)
 
@@ -375,7 +375,7 @@ def visit_join_expr(self, node: ast.JoinExpr):
 
             # :TRICKY: Make sure to clone and visit _all_ JoinExpr fields/nodes.
             node.next_join = self.visit(node.next_join)
-            if node.constraint and node.constraint.join_type == "ON":
+            if node.constraint and node.constraint.constraint_type == "ON":
                 node.constraint = self.visit(node.constraint)
             node.sample = self.visit(node.sample)
 

posthog/hogql/test/_test_parser.py

@@ -790,7 +790,7 @@ def test_select_from_join(self):
                         next_join=ast.JoinExpr(
                             join_type="JOIN",
                             table=ast.Field(chain=["events2"]),
-                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1)),
+                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1), constraint_type="ON"),
                         ),
                     ),
                 ),
@@ -804,7 +804,7 @@ def test_select_from_join(self):
                         next_join=ast.JoinExpr(
                             join_type="LEFT OUTER JOIN",
                             table=ast.Field(chain=["events2"]),
-                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1)),
+                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1), constraint_type="ON"),
                         ),
                     ),
                 ),
@@ -818,16 +818,30 @@ def test_select_from_join(self):
                         next_join=ast.JoinExpr(
                             join_type="LEFT OUTER JOIN",
                             table=ast.Field(chain=["events2"]),
-                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1)),
+                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1), constraint_type="ON"),
                             next_join=ast.JoinExpr(
                                 join_type="RIGHT ANY JOIN",
                                 table=ast.Field(chain=["events3"]),
-                                constraint=ast.JoinConstraint(expr=ast.Constant(value=2)),
+                                constraint=ast.JoinConstraint(expr=ast.Constant(value=2), constraint_type="ON"),
                             ),
                         ),
                     ),
                 ),
             )
+            self.assertEqual(
+                self._select("select 1 from events JOIN events2 USING 1"),
+                ast.SelectQuery(
+                    select=[ast.Constant(value=1)],
+                    select_from=ast.JoinExpr(
+                        table=ast.Field(chain=["events"]),
+                        next_join=ast.JoinExpr(
+                            join_type="JOIN",
+                            table=ast.Field(chain=["events2"]),
+                            constraint=ast.JoinConstraint(expr=ast.Constant(value=1), constraint_type="USING"),
+                        ),
+                    ),
+                ),
+            )
 
         def test_select_from_join_multiple(self):
             node = self._select(
@@ -863,7 +877,8 @@ def test_select_from_join_multiple(self):
                                     op=ast.CompareOperationOp.Eq,
                                     left=ast.Field(chain=["pdi", "distinct_id"]),
                                     right=ast.Field(chain=["e", "distinct_id"]),
-                                )
+                                ),
+                                constraint_type="ON",
                             ),
                             next_join=ast.JoinExpr(
                                 join_type="LEFT JOIN",
@@ -874,7 +889,8 @@ def test_select_from_join_multiple(self):
                                         op=ast.CompareOperationOp.Eq,
                                         left=ast.Field(chain=["p", "id"]),
                                         right=ast.Field(chain=["pdi", "person_id"]),
-                                    )
+                                    ),
+                                    constraint_type="ON",
                                 ),
                             ),
                         ),

posthog/hogql/test/test_resolver.py

@@ -31,7 +31,7 @@ class TestResolver(BaseTest):
     def _select(self, query: str, placeholders: Optional[dict[str, ast.Expr]] = None) -> ast.SelectQuery:
         return cast(
             ast.SelectQuery,
-            clone_expr(parse_select(query, placeholders=placeholders, backend="python"), clear_locations=True),
+            clone_expr(parse_select(query, placeholders=placeholders), clear_locations=True),
         )
 
     def _print_hogql(self, select: str):
@@ -268,17 +268,18 @@ def test_ctes_with_union_all(self):
                     UNION ALL
                     SELECT * FROM (SELECT 1 AS a) AS cte1
                         """),
-            )
+        )
 
     def test_join_using(self):
-        self.assertEqual(
-            self._print_hogql(
-                "with my_table as (select 1 as a) select q1.a from my_table as q1 inner join my_table as q2 USING a"
-            ),
-            self._print_hogql(
-                "with my_table as (select 1 as a) select q1.a from my_table as q1 inner join my_table as q2 USING a"
-            ),
+        node = self._select(
+            "with my_table as (select 1 as a) select q1.a from my_table as q1 inner join my_table as q2 USING a"
         )
+        node = cast(ast.SelectQuery, resolve_types(node, self.context, dialect="clickhouse"))
+        assert cast(ast.SelectQuery, node).select_from.next_join.constraint.constraint_type == "USING"
+
+        node = self._select("select q1.event from events as q1 inner join events as q2 USING event")
+        node = cast(ast.SelectQuery, resolve_types(node, self.context, dialect="clickhouse"))
+        assert cast(ast.SelectQuery, node).select_from.next_join.constraint.constraint_type == "USING"
 
     @override_settings(PERSON_ON_EVENTS_OVERRIDE=False, PERSON_ON_EVENTS_V2_OVERRIDE=False)
     @pytest.mark.usefixtures("unittest_snapshot")

posthog/hogql/visitor.py

@@ -529,7 +529,7 @@ def visit_window_frame_expr(self, node: ast.WindowFrameExpr):
         )
 
     def visit_join_constraint(self, node: ast.JoinConstraint):
-        return ast.JoinConstraint(expr=self.visit(node.expr), join_type=node.join_type)
+        return ast.JoinConstraint(expr=self.visit(node.expr), constraint_type=node.constraint_type)
 
     def visit_hogqlx_tag(self, node: ast.HogQLXTag):
         return ast.HogQLXTag(kind=node.kind, attributes=[self.visit(a) for a in node.attributes])