Update Python model and types

Author: yujiyokoo

languages/python/jupyter_notebook/cs_models.py

@@ -9,19 +9,19 @@
 import os
 
 class CsTypeDecorator(TypeDecorator):
-    def __init__(self, table_name, column_name):
+    def __init__(self, table, column):
         super().__init__()
-        self.table_name = table_name
-        self.column_name = column_name
+        self.table = table
+        self.column = column
 
     def process_bind_param(self, value, dialect):
         if value is not None:
             value_dict = {
                 "k": "pt",
                 "p": str(value),
                 "i": {
-                    "t": self.table_name,
-                    "c": self.column_name
+                    "t": self.table,
+                    "c": self.column
                 },
                 "v": 1,
                 "q": None
@@ -106,4 +106,12 @@ def __init__(self, e_utf8_str=None, e_jsonb=None, e_int=None, e_float=None, e_da
         self.encrypted_boolean = e_bool
 
     def __repr__(self):
-        return f"<Example(id={self.id}, encrypted_utf8_str={self.encrypted_utf8_str}, encrypted_jsonb={self.encrypted_jsonb}, encrypted_int={self.encrypted_int}, encrypted_float={self.encrypted_float}, encrypted_date={self.encrypted_date}, encrypted_boolean={self.encrypted_boolean})>"
+        return "<Example(" \
+            f"id={self.id}, " \
+            f"encrypted_utf8_str={self.encrypted_utf8_str}, " \
+            f"encrypted_jsonb={self.encrypted_jsonb}, " \
+            f"encrypted_int={self.encrypted_int}, " \
+            f"encrypted_float={self.encrypted_float}, " \
+            f"encrypted_date={self.encrypted_date}, " \
+            f"encrypted_boolean={self.encrypted_boolean}" \
+            ")>"

languages/python/jupyter_notebook/cs_models_test.py

@@ -4,9 +4,14 @@
 from cs_models import *
 
 class TestExampleModel(unittest.TestCase):
+    pg_password = os.getenv('PGPASSWORD', 'postgres')
+    pg_user = os.getenv('PGUSER', 'postgres')
+    pg_host = os.getenv('PGHOST', 'localhost')
+    pg_port = os.getenv('PGPORT', '6432')
+    pg_db = os.getenv('PGDATABASE', 'cs_test_db')
+
     def setUp(self):
-        # TODO: configure database URL in environment variable and use a test db (not getting_started)
-        self.engine = create_engine('postgresql://postgres:postgres@localhost:6432/cipherstash_getting_started')
+        self.engine = create_engine(f'postgresql://{self.pg_user}:{self.pg_password}@{self.pg_host}:{self.pg_port}/{self.pg_db}')
         Session = sessionmaker(bind=self.engine)
         self.session = Session()
         BaseModel.metadata.create_all(self.engine)
@@ -41,6 +46,13 @@ def test_encrypted_utf8_str(self):
     def test_encrypted_jsonb(self):
         found = self.session.query(Example).filter(Example.id == self.example.id).one()
         self.assertEqual(found.encrypted_jsonb,  {"key": "value"})
+
+    def test_example_prints_value(self):
+        self.example.id = 1
+        self.assertEqual(
+            str(self.example),
+            "<Example(id=1, encrypted_utf8_str=str, encrypted_jsonb={'key': 'value'}, encrypted_int=1, encrypted_float=1.1, encrypted_date=2024-01-01, encrypted_boolean=True)>"
+        )
 
 if __name__ == '__main__':
     unittest.main()

languages/python/jupyter_notebook/cs_types.py

@@ -2,17 +2,18 @@
 from pprint import pprint
 from datetime import datetime
 import json
+from enum import Enum
 
 class CsValue:
     def __init__(self, v, t: str, c: str):
         self.value = v
         self.table = t
         self.column = c
 
-    def to_db_format(self, query_type=None):
+    def to_db_format(self, query_type = None):
         data = {
             "k": "pt",
-            "p": self.value_in_db_format(query_type),
+            "p": self._value_in_db_format(query_type),
             "i": {
               "t": str(self.table),
               "c": str(self.column)
@@ -24,62 +25,59 @@ def to_db_format(self, query_type=None):
 
     @classmethod
     def from_parsed_json(cls, parsed):
-        return cls.value_from_db_format(parsed["p"])
+        return cls._value_from_db_format(parsed["p"])
 
 class CsInt(CsValue):
-    def value_in_db_format(self, query_type = None):
+    def _value_in_db_format(self, query_type):
         return str(self.value)
 
     @classmethod
-    def value_from_db_format(cls, s: str):
+    def _value_from_db_format(cls, s: str):
         return int(s)
 
 class CsBool(CsValue):
-    def value_in_db_format(self, query_type = None):
+    def _value_in_db_format(self, query_type):
         return str(self.value).lower()
 
     @classmethod
-    def value_from_db_format(cls, s: str):
+    def _value_from_db_format(cls, s: str):
         return s.lower() == 'true'
 
 class CsDate(CsValue):
-    def value_in_db_format(self, query_type = None):
+    def _value_in_db_format(self, query_type):
         return self.value.isoformat()
 
     @classmethod
-    def value_from_db_format(cls, s: str):
+    def _value_from_db_format(cls, s: str):
         return datetime.fromisoformat(s).date()
 
 class CsFloat(CsValue):
-    def value_in_db_format(self, query_type = None):
+    def _value_in_db_format(self, query_type):
         return str(self.value)
 
     @classmethod
-    def value_from_db_format(cls, s: str):
+    def _value_from_db_format(cls, s: str):
         return float(s)
 
 class CsText(CsValue):
-    def value_in_db_format(self, query_type = None):
+    def _value_in_db_format(self, query_type):
         return self.value
 
     @classmethod
-    def value_from_db_format(cls, s: str):
+    def _value_from_db_format(cls, s: str):
         return s
 
 class CsJsonb(CsValue):
-    def value_in_db_format(self, query_type):
+    def _value_in_db_format(self, query_type):
         if query_type == "ejson_path":
             return self.value
         else:
             return json.dumps(self.value)
 
     @classmethod
-    def value_from_db_format(cls, s: str):
+    def _value_from_db_format(cls, s: str):
         return json.loads(s)
 
-def id_map(x):
-    return x
-
 class CsRow:
     column_function_mapping = {
         'encrypted_int': CsInt.from_parsed_json,
@@ -90,9 +88,15 @@ class CsRow:
         'encrypted_jsonb': CsText.from_parsed_json
     }
 
+    @staticmethod
+    def id_map(x):
+        return x
+
     def __init__(self, row):
         self.row = {}
         for k, v in row.items():
-            self.row[k] = None if v == None else self.column_function_mapping.get(k, id_map)(v)
-
-
+            if v == None:
+                self.row[k] = None
+            else:
+                mapping = self.column_function_mapping.get(k, self.id_map)
+                self.row[k] = mapping(v)

languages/python/jupyter_notebook/cs_types_test.py

@@ -4,6 +4,9 @@
 from cs_types import *
 
 class EqlTest(unittest.TestCase):
+    def setUp(self):
+        self.template_dict = json.loads('{"k": "pt", "p": "1", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}')
+
     def test(self):
         self.assertTrue(True)
 
@@ -14,94 +17,113 @@ def test_to_db_format(self):
         )
 
     def test_from_parsed_json_uses_p_value(self):
-        parsed = json.loads('{"k": "pt", "p": "1", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}')
+        self.template_dict["p"] = "1"
         self.assertEqual(
-            CsInt.from_parsed_json(parsed),
+            CsInt.from_parsed_json(self.template_dict),
             1
         )
 
-    def test_cs_int_prints_value(self):
-        cs_int = CsInt(1, "table", "column")
+    def test_cs_int_to_db_format(self):
+        cs_int = CsInt(123, "table", "column")
         self.assertEqual(
-            cs_int.value_in_db_format(),
-            "1"
+            '{"k": "pt", "p": "123", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}',
+            cs_int.to_db_format()
         )
 
-    def test_ces_int_makes_int(self):
+    def test_cs_int_from_parsed_json(self):
+        self.template_dict["p"] = "123"
         self.assertEqual(
-            CsInt.value_from_db_format("1"),
-            1
+            CsInt.from_parsed_json(self.template_dict),
+            123
         )
 
-    def test_cs_bool_prints_value_in_lower_case(self):
+    def test_cs_bool_to_db_format_true(self):
         cs_bool = CsBool(True, "table", "column")
         self.assertEqual(
-            cs_bool.value_in_db_format(),
-            "true"
+            '{"k": "pt", "p": "true", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}',
+            cs_bool.to_db_format()
+        )
+    
+    def test_cs_bool_to_db_format_false(self):
+        cs_bool = CsBool(False, "table", "column")
+        self.assertEqual(
+            '{"k": "pt", "p": "false", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}',
+            cs_bool.to_db_format()
         )
 
-    def test_cs_bool_returns_bool(self):
+    def test_cs_bool_from_parsed_json_true(self):
+        self.template_dict["p"] = "true"
         self.assertEqual(
-            CsBool.value_from_db_format("true"),
+            CsBool.from_parsed_json(self.template_dict),
             True
         )
 
-    def test_cs_date_prints_value(self):
+    def test_cs_bool_from_parsed_json_false(self):
+        self.template_dict["p"] = "false"
+        self.assertEqual(
+            CsBool.from_parsed_json(self.template_dict),
+            False
+        )
+
+    def test_cs_date_to_db_format(self):
         cs_date = CsDate(date(2024, 11, 1), "table", "column")
         self.assertEqual(
-            cs_date.value_in_db_format(),
-            "2024-11-01"
+            '{"k": "pt", "p": "2024-11-01", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}',
+            cs_date.to_db_format()
         )
 
-    def test_cs_date_returns_datetime(self):
+    def test_cs_date_from_parsed_json(self):
+        self.template_dict["p"] = "2024-11-01"
         self.assertEqual(
-            CsDate.value_from_db_format("2024-11-01"),
+            CsDate.from_parsed_json(self.template_dict),
             date(2024, 11, 1)
         )
 
-    def test_cs_float_prints_value(self):
+    def test_cs_float_to_db_format(self):
         cs_float = CsFloat(1.1, "table", "column")
         self.assertEqual(
-            cs_float.value_in_db_format(),
-            "1.1"
+            '{"k": "pt", "p": "1.1", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}',
+            cs_float.to_db_format()
         )
 
-    def test_cs_float_returns_float(self):
+    def test_cs_float_from_parsed_json(self):
+        self.template_dict["p"] = "1.1"
         self.assertEqual(
-            CsFloat.value_from_db_format("1.1"),
+            CsFloat.from_parsed_json(self.template_dict),
             1.1
         )
 
-    def test_cs_text_prints_value(self):
+    def test_cs_text_to_db_format(self):
         cs_text = CsText("text", "table", "column")
         self.assertEqual(
-            cs_text.value_in_db_format(),
-            "text"
+            '{"k": "pt", "p": "text", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}',
+            cs_text.to_db_format()
         )
 
-    def test_cs_text_returns_value(self):
+    def test_cs_text_from_parsed_json(self):
+        self.template_dict["p"] = "text"
         self.assertEqual(
-            CsText.value_from_db_format("text"),
+            CsText.from_parsed_json(self.template_dict),
             "text"
         )
 
     def test_cs_jsonb_prints_json_string(self):
         cs_jsonb = CsJsonb({"a": 1}, "table", "column")
         self.assertEqual(
-            cs_jsonb.value_in_db_format("ste_vec"),
+            cs_jsonb._value_in_db_format("ste_vec"),
             '{"a": 1}'
         )
 
     def test_cs_jsonb_prints_value_for_ejson_path(self):
         cs_jsonb = CsJsonb("$.a.b", "table", "column")
         self.assertEqual(
-            cs_jsonb.value_in_db_format("ejson_path"),
+            cs_jsonb._value_in_db_format("ejson_path"),
             '$.a.b'
         )
 
     def test_cs_jsonb_returns_value(self):
         self.assertEqual(
-            CsJsonb.value_from_db_format('{"a": 1}'),
+            CsJsonb._value_from_db_format('{"a": 1}'),
             {"a": 1}
         )