projections (#4)

roman-right · web-flow · commit ec7b2caf58bb · 2023-04-05T22:14:42.000+02:00
* feature | projections

* remove | prints

* update | types map

* update | init imports

* update | tests gh action
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -8,12 +8,11 @@ jobs:
       matrix:
         python-version: [ 3.8, 3.9, 3.10.5 ]
         poetry-version: [ 1.2.0 ]
-        os: [ ubuntu-18.04 ]
-    runs-on: ${{ matrix.os }}
+    runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
       - name: Build And Run DB
-        run: scripts/build_and_run.sh
+        run: docker run -d --init -p 10001:10001 -t romanright/annadb:0.2.0
       - uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
diff --git a/annadb/__init__.py b/annadb/__init__.py
@@ -1,18 +1,38 @@
-__version__ = "0.1.3"
+__version__ = "0.1.4"
 
 __all__ = [
     # Entities
     "Connection",
     # utils
     "to_str",
     # Operators
+    # Update operators
     "Inc",
     "Set",
+    # Find operators
+    # Comparison
+    "Eq",
+    "Neq",
+    "Gt",
+    "Lt",
+    "Gte",
+    "Lte",
+    # Logical
+    "And",
+    "Or",
+    "Not",
+    # Projection operators
+    "Keep",
+    "keep",
     # Fields
     "root",
+    # Constants
+    "null",
 ]
 
 from annadb.connection import Connection
 from annadb.dump import to_str
-from annadb.query.path import root
+from annadb.query.find.operators import Eq, Neq, Gt, Lt, Gte, Lte, And, Or, Not
+from annadb.query.project.operators import Keep
 from annadb.query.update.operators import Inc, Set
+from annadb.data_types.constants import root, null, keep
diff --git a/annadb/constants.py b/annadb/constants.py
@@ -21,6 +21,7 @@
 SORT_QUERY: str = "sort"
 OFFSET_QUERY: str = "offset"
 LIMIT_QUERY: str = "limit"
+PROJECT_QUERY: str = "project"
 
 # FIND OPERATORS
 EQ_OPERATOR: str = "eq"
@@ -43,6 +44,9 @@
 ASC_OPERATOR: str = "asc"
 DESC_OPERATOR: str = "desc"
 
+# PROJECT OPERATORS
+KEEP_OPERATOR: str = "keep"
+
 # RESPONSE
 RESPONSE_OBJECTS: str = "objects"
 RESPONSE_IDS: str = "ids"
diff --git a/annadb/data_types/constants.py b/annadb/data_types/constants.py
@@ -1,5 +1,7 @@
+from annadb import Keep
 from annadb.data_types.primitive import Null
 from annadb.query.path import Path
 
 root = Path(None)
 null = Null()
+keep = Keep()
diff --git a/annadb/data_types/types_map.py b/annadb/data_types/types_map.py
@@ -43,6 +43,8 @@
     QUERY_SET,
     LIMIT_QUERY,
     OFFSET_QUERY,
+    KEEP_OPERATOR,
+    PROJECT_QUERY,
 )
 from annadb.data_types.map import Map
 from annadb.data_types.primitive import (
@@ -58,6 +60,7 @@
 from annadb.data_types.vector import Vector
 from annadb.query.find.operators import Eq, Neq, Gt, Gte, Lt, Lte, And, Or, Not
 from annadb.query.path import Path
+from annadb.query.project.operators import Keep
 from annadb.query.query_set import QuerySet
 from annadb.query.sort.operators import Asc, Desc
 from annadb.query.types import (
@@ -68,7 +71,7 @@
     Delete,
     Sort,
     Limit,
-    Offset,
+    Offset, Project,
 )
 from annadb.query.update.operators import Set, Inc
 from annadb.response import (
@@ -106,6 +109,7 @@
     SORT_QUERY: Sort,
     LIMIT_QUERY: Limit,
     OFFSET_QUERY: Offset,
+    PROJECT_QUERY: Project,
     # FIND OPERATORS
     EQ_OPERATOR: Eq,
     NEQ_OPERATOR: Neq,
@@ -122,6 +126,8 @@
     # SORT OPERATORS
     ASC_OPERATOR: Asc,
     DESC_OPERATOR: Desc,
+    # PROJECTION OPERATORS
+    KEEP_OPERATOR: Keep,
     # RESPONSE
     RESPONSE_OBJECTS: Objects,
     RESPONSE_IDS: IDs,
diff --git a/annadb/query/base.py b/annadb/query/base.py
@@ -58,3 +58,10 @@ def update(self, *args):
         return UpdateQuery(
             *args, query_set=self.query_set, collection=self.collection
         )
+
+    def project(self, projection: dict):
+        from annadb.query.project.query import ProjectQuery
+
+        return ProjectQuery(
+            projection, query_set=self.query_set, collection=self.collection
+        )
diff --git a/annadb/query/project/__init__.py b/annadb/query/project/__init__.py
diff --git a/annadb/query/project/operators.py b/annadb/query/project/operators.py
@@ -0,0 +1,10 @@
+from annadb.constants import KEEP_OPERATOR
+from annadb.data_types.primitive import PrimitiveBase
+
+
+class Keep(PrimitiveBase):
+    prefix = KEEP_OPERATOR
+    instance_type = "keep"
+
+    def __init__(self, _=None):
+        super(Keep, self).__init__(None)
diff --git a/annadb/query/project/query.py b/annadb/query/project/query.py
@@ -0,0 +1,20 @@
+import typing
+
+from annadb.query.base import BaseQuery
+from annadb.query.query_set import QuerySet
+from annadb.query.types import Project
+
+if typing.TYPE_CHECKING:
+    from annadb.collection import Collection
+
+
+class ProjectQuery(BaseQuery):
+    def __init__(
+            self,
+            project_map: typing.Optional[dict],
+            query_set: QuerySet,
+            collection: "Collection"
+    ):
+        self.collection = collection
+        self.query_set = query_set
+        self.query_set.push(Project(data=project_map))
diff --git a/annadb/query/types.py b/annadb/query/types.py
@@ -6,8 +6,9 @@
     UPDATE_QUERY,
     OFFSET_QUERY,
     LIMIT_QUERY,
-    DELETE_QUERY,
+    DELETE_QUERY, PROJECT_QUERY,
 )
+from annadb.data_types.map import MapUnique
 from annadb.data_types.modifier import ModifierBase
 from annadb.data_types.primitive import PrimitiveBase
 from annadb.data_types.vector import VectorBase
@@ -46,3 +47,7 @@ class Delete(PrimitiveBase):
 
     def __init__(self, _=None):
         super(Delete, self).__init__(None)
+
+
+class Project(MapUnique):
+    prefix = PROJECT_QUERY
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "annadb"
-version = "0.1.3"
+version = "0.1.4"
 description = "AnnaDB driver and shell"
 authors = ["Roman Right <roman-right@protonmail.com>"]
 
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -59,6 +59,8 @@ def objects(collection):
             "blink": "test" if i % 2 == 0 else {"smth": i},
             "blink2": {"a": 100, "c": "d"} if i % 2 == 0 else {},
             "ts": UTS(12345),
+            "d": {"smth": i, "smth2": 2},
+            "l": [9,8,7,6,5,4,3]
         }
         for i in range(10)
     ]
diff --git a/tests/e2e/test_project.py b/tests/e2e/test_project.py
@@ -0,0 +1,96 @@
+from annadb import to_str
+from annadb.console import HTMLConsole
+
+
+class TestProject:
+    def test_keep(self, conn, objects):
+        resp = conn.send_query(
+            """
+            collection|test|:q[
+                find[],
+                sort[asc(value|name|)],
+                project{
+                    s|name|:keep,
+                    s|d|:keep,
+                }
+            ]
+            """
+        )
+        for i, (k, v) in enumerate(resp[0]["data"].items()):
+            assert set(v._value.keys()) == {"name", "d"}
+            assert v["name"] == f"test_{i}"
+            assert v["d"] == {"smth": i, "smth2": 2}
+
+    def test_set_by_path(self, conn, objects):
+        resp = conn.send_query(
+            """
+            collection|test|:q[
+                find[],
+                sort[asc(value|name|)],
+                project{
+                    s|name|:value|smth|
+                }
+            ]
+            """
+        )
+        for i, (k, v) in enumerate(resp[0]["data"].items()):
+            assert set(v._value.keys()) == {"name"}
+            assert v["name"] == "TEST"
+
+    def test_set_value_primitive(self, conn, objects):
+        resp = conn.send_query(
+            """
+            collection|test|:q[
+                find[],
+                sort[asc(value|name|)],
+                project{
+                    s|name|:s|NEW VALUE|
+                }
+            ]
+            """
+        )
+        for i, (k, v) in enumerate(resp[0]["data"].items()):
+            assert set(v._value.keys()) == {"name"}
+            assert v["name"] == "NEW VALUE"
+
+    def test_set_value_map(self, conn, objects):
+        resp = conn.send_query(
+            """
+            collection|test|:q[
+                find[],
+                sort[asc(value|name|)],
+                project{
+                    s|name|:m{
+                        s|test|:value|smth|,
+                    },
+                    s|d|:m{
+                        s|smth2|:keep
+                    }
+                }
+            ]
+            """
+        )
+        for i, (k, v) in enumerate(resp[0]["data"].items()):
+            assert set(v._value.keys()) == {"name", "d"}
+            assert v["name"] == {"test": "TEST"}
+            assert v["d"] == {"smth2": 2}
+
+    def test_set_value_vector(self, conn, objects):
+        resp = conn.send_query(
+            """
+            collection|test|:q[
+                find[],
+                sort[asc(value|name|)],
+                project{
+                    s|name|:v[
+                        value|smth|,
+                    ],
+                    s|l|:v[s|TEST|,keep,keep]
+                }
+            ]
+            """
+        )
+        for i, (k, v) in enumerate(resp[0]["data"].items()):
+            assert set(v._value.keys()) == {"name", "l"}
+            assert v["name"] == ["TEST"]
+            assert v["l"] == ["TEST", 8, 7]
diff --git a/tests/query/test_delete.py b/tests/query/test_delete.py
@@ -1,5 +1,3 @@
-from pprint import pprint
-
 from annadb.data_types.constants import root
 
 
@@ -28,5 +26,4 @@ def test_delete_one(self, collection, primitives):
         resp_delete = collection.get(primitives.data[0]).delete().run()
         assert resp_delete.meta.count == 1
         resp = collection.all().run()
-        pprint(resp.data.items())
         assert len(resp.data) == primitives.meta.count - 1
diff --git a/tests/query/test_project.py b/tests/query/test_project.py
@@ -0,0 +1,56 @@
+from annadb import root, keep
+
+
+class TestProject:
+    def test_keep(self, collection, objects):
+        resp = collection.find().sort(+root.name).project(
+            {"name": keep, "d": keep}
+        ).run()
+        for i, (k, v) in enumerate(resp.data.items()):
+            assert set(v._value.keys()) == {"name", "d"}
+            assert v["name"] == f"test_{i}"
+            assert v["d"] == {"smth": i, "smth2": 2}
+
+    def test_set_by_path(self, collection, objects):
+        resp = collection.find().sort(+root.name).project(
+            {"name": root.smth}
+        ).run()
+        for i, (k, v) in enumerate(resp.data.items()):
+            assert set(v._value.keys()) == {"name"}
+            assert v["name"] == "TEST"
+
+    def test_set_value_primitive(self, collection, objects):
+        resp = collection.find().sort(+root.name).project(
+            {"name": "NEW VALUE"}
+        ).run()
+        for i, (k, v) in enumerate(resp.data.items()):
+            assert set(v._value.keys()) == {"name"}
+            assert v["name"] == "NEW VALUE"
+
+    def test_set_value_map(self, collection, objects):
+        resp = collection.find().sort(+root.name).project(
+            {
+                "name": {
+                    "test": root.smth
+                },
+                "d": {
+                    "smth2": keep
+                }
+            }
+        ).run()
+        for i, (k, v) in enumerate(resp.data.items()):
+            assert set(v._value.keys()) == {"name", "d"}
+            assert v["name"] == {"test": "TEST"}
+            assert v["d"] == {"smth2": 2}
+
+    def test_set_value_vector(self, collection, objects):
+        resp = collection.find().sort(+root.name).project(
+            {
+                "name": [root.smth],
+                "l": ["TEST", keep, keep]
+            }
+        ).run()
+        for i, (k, v) in enumerate(resp.data.items()):
+            assert set(v._value.keys()) == {"name", "l"}
+            assert v["name"] == ["TEST"]
+            assert v["l"] == ["TEST", 8, 7]
diff --git a/tests/query/test_update.py b/tests/query/test_update.py
@@ -42,7 +42,6 @@ def test_rewrite_object(self, collection):
         ).run()
         resp = collection.find().run()
         for k, v in resp.data.items():
-            print(v.items())
             assert v["a"] == "b"
 
 
diff --git a/tests/test_annadb.py b/tests/test_annadb.py
@@ -2,4 +2,4 @@
 
 
 def test_version():
-    assert __version__ == "0.1.3"
+    assert __version__ == "0.1.4"
