From 42da7f1896a4884b1cdc595626b0bf006b105e22 Mon Sep 17 00:00:00 2001
From: timgl <tim@posthog.com>
Date: Fri, 26 Jul 2024 11:42:27 -0700
Subject: [PATCH] fix

---
 posthog/hogql/property.py           | 10 ++++++----
 posthog/hogql/test/test_property.py | 25 +++++++++++++------------
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/posthog/hogql/property.py b/posthog/hogql/property.py
index 98cb2408db8f3..afceaf8f3edf8 100644
--- a/posthog/hogql/property.py
+++ b/posthog/hogql/property.py
@@ -620,11 +620,13 @@ def selector_to_expr(selector_string: str):
                 useful_elements.append(ast.Constant(value=part.data["tag_name"]))
 
         if "attr_id" in part.data:
-            exprs.append(
-                parse_expr(
-                    "indexOf(elements_chain_ids, {value}) > 0", {"value": ast.Constant(value=part.data["attr_id"])}
-                )
+            id_expr = parse_expr(
+                "indexOf(elements_chain_ids, {value}) > 0", {"value": ast.Constant(value=part.data["attr_id"])}
             )
+            if len(selector.parts) == 1 and len(part.data.keys()) == 1:
+                # OPTIMIZATION: if there's only one selector part and that only filters on an ID, we don't need to also query elements_chain separately
+                return id_expr
+            exprs.append(id_expr)
     if len(useful_elements) > 0:
         exprs.append(
             parse_expr(
diff --git a/posthog/hogql/test/test_property.py b/posthog/hogql/test/test_property.py
index bbf683d9b4019..02aa63b9cb856 100644
--- a/posthog/hogql/test/test_property.py
+++ b/posthog/hogql/test/test_property.py
@@ -501,31 +501,37 @@ def test_selector_to_expr(self):
             ),
         )
         self.assertEqual(
-            self._selector_to_expr("#withid"),
+            self._selector_to_expr("a#withid"),
             clear_locations(
                 parse_expr(
-                    """{regex} and indexOf(elements_chain_ids, 'withid') > 0""",
+                    """{regex} and indexOf(elements_chain_ids, 'withid') > 0 and arrayCount(x -> x IN ['a'], elements_chain_elements) > 0""",
                     {
                         "regex": elements_chain_match(
-                            '(^|;).*?attr_id="withid".*?([-_a-zA-Z0-9\\.:"= ]*?)?($|;|:([^;^\\s]*(;|$|\\s)))'
+                            '(^|;)a.*?attr_id="withid".*?([-_a-zA-Z0-9\\.:"= ]*?)?($|;|:([^;^\\s]*(;|$|\\s)))'
                         )
                     },
                 )
             ),
         )
+
         self.assertEqual(
-            self._selector_to_expr("#with-dashed-id"),
+            self._selector_to_expr("a#with-dashed-id"),
             clear_locations(
                 parse_expr(
-                    """{regex} and indexOf(elements_chain_ids, 'with-dashed-id') > 0""",
+                    """{regex} and indexOf(elements_chain_ids, 'with-dashed-id') > 0 and arrayCount(x -> x IN ['a'], elements_chain_elements) > 0""",
                     {
                         "regex": elements_chain_match(
-                            '(^|;).*?attr_id="with\\-dashed\\-id".*?([-_a-zA-Z0-9\\.:"= ]*?)?($|;|:([^;^\\s]*(;|$|\\s)))'
+                            '(^|;)a.*?attr_id="with\\-dashed\\-id".*?([-_a-zA-Z0-9\\.:"= ]*?)?($|;|:([^;^\\s]*(;|$|\\s)))'
                         )
                     },
                 )
             ),
         )
+        # test optimization
+        self.assertEqual(
+            self._selector_to_expr("#with-dashed-id"),
+            clear_locations(parse_expr("""indexOf(elements_chain_ids, 'with-dashed-id') > 0""")),
+        )
         self.assertEqual(
             self._selector_to_expr("#with-dashed-id"),
             self._selector_to_expr("[id='with-dashed-id']"),
@@ -534,12 +540,7 @@ def test_selector_to_expr(self):
             self._selector_to_expr("#with\\slashed\\id"),
             clear_locations(
                 parse_expr(
-                    "{regex} and indexOf(elements_chain_ids, 'with\\\\slashed\\\\id') > 0",
-                    {
-                        "regex": elements_chain_match(
-                            '(^|;).*?attr_id="with\\\\slashed\\\\id".*?([-_a-zA-Z0-9\\.:"= ]*?)?($|;|:([^;^\\s]*(;|$|\\s)))'
-                        )
-                    },
+                    "indexOf(elements_chain_ids, 'with\\\\slashed\\\\id') > 0",
                 )
             ),
         )