SK-2954: Rename insert records->values, drop Upsert class for a dict shape, add types

BaseInsertRequest's shared field is now named values (matching PDB's
terminology) instead of records, and table/values are required rather than
defaulted; each variant's InsertRequest forwards them explicitly. Removes
flowvault's Upsert class in favor of a plain dict (now typed via a TypedDict)
to match the rest of flowvault's dict-based request shape, and adds return
type hints to every method on v2's VaultController to match insert's.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: saileshwar-skyflow

common/tests/vault/data/test_base_insert.py

@@ -4,16 +4,18 @@
 
 
 class TestBaseInsertRequest(unittest.TestCase):
-    def test_defaults(self):
-        request = BaseInsertRequest()
-        self.assertIsNone(request.table)
-        self.assertIsNone(request.records)
+    def test_table_and_values_are_required(self):
+        with self.assertRaises(TypeError):
+            BaseInsertRequest()
+
+    def test_upsert_defaults_to_none(self):
+        request = BaseInsertRequest(table="t1", values=[{"values": {"a": 1}}])
         self.assertIsNone(request.upsert)
 
     def test_construction(self):
-        request = BaseInsertRequest(table="t1", records=[{"values": {"a": 1}}], upsert="upsert_val")
+        request = BaseInsertRequest(table="t1", values=[{"values": {"a": 1}}], upsert="upsert_val")
         self.assertEqual(request.table, "t1")
-        self.assertEqual(request.records, [{"values": {"a": 1}}])
+        self.assertEqual(request.values, [{"values": {"a": 1}}])
         self.assertEqual(request.upsert, "upsert_val")
 
 

common/vault/base_vault_controller.py

@@ -60,8 +60,3 @@ def _validate_field_values(self, values):
                     self._skyflow_messages.Error.EMPTY_KEY_IN_INSERT_DATA.value,
                     _INVALID_INPUT_ERROR_CODE,
                 )
-            # if value is None or (isinstance(value, str) and not value.strip()):
-            #     raise SkyflowError(
-            #         self._skyflow_messages.Error.EMPTY_VALUE_IN_INSERT_DATA.value,
-            #         _INVALID_INPUT_ERROR_CODE,
-            #     )

common/vault/data/_base_insert_request.py

@@ -1,6 +1,7 @@
-class BaseInsertRequest:
+from typing import Union
 
-    def __init__(self, table=None, records=None, upsert=None):
+class BaseInsertRequest:
+    def __init__(self, table: str, values: list, upsert: Union[str, dict] = None):
         self.table = table
-        self.records = records
-        self.upsert = upsert
+        self.values = values
+        self.upsert = upsert

flowvault/skyflow_flowvault/utils/_skyflow_messages.py

@@ -20,9 +20,9 @@ class Error(Enum):
         INVALID_RECORDS_TYPE_IN_INSERT = f"{error_prefix} Insert failed. 'records' must be a list of dicts."
         INVALID_RECORD_DATA_IN_INSERT = f"{error_prefix} Insert failed. Each record's 'values' must be a non-empty dict."
         INVALID_TABLE_NAME_IN_INSERT = f"{error_prefix} Insert failed. 'table' must be a non-empty string."
-        INVALID_UPSERT_TYPE_IN_INSERT = f"{error_prefix} Insert failed. 'upsert' must be an Upsert instance."
-        INVALID_UPSERT_UNIQUE_COLUMNS_IN_INSERT = f"{error_prefix} Insert failed. Upsert.unique_columns must be a non-empty list of strings."
-        INVALID_UPSERT_UPDATE_TYPE_IN_INSERT = f"{error_prefix} Insert failed. Upsert.update_type must be an UpsertType value."
+        INVALID_UPSERT_TYPE_IN_INSERT = f"{error_prefix} Insert failed. 'upsert' must be a dict."
+        INVALID_UPSERT_UNIQUE_COLUMNS_IN_INSERT = f"{error_prefix} Insert failed. Upsert's 'unique_columns' must be a non-empty list of strings."
+        INVALID_UPSERT_UPDATE_TYPE_IN_INSERT = f"{error_prefix} Insert failed. Upsert's 'update_type' must be an UpsertType value."
         TOO_MANY_RECORDS_IN_INSERT = f"{error_prefix} Insert failed. A single insert request cannot contain more than 10000 records."
         TABLE_NAME_IN_BOTH_PLACES_IN_INSERT = (
             f"{error_prefix} Insert failed. 'table' cannot be set on InsertRequest at the same "

flowvault/skyflow_flowvault/utils/validations/_validations.py

@@ -8,9 +8,9 @@
 )
 from skyflow_flowvault.utils import SkyflowMessages
 from skyflow_flowvault.utils.enums import UpsertType
-from skyflow_flowvault.vault.data import Upsert
 
 VALID_INSERT_RECORD_KEYS = ["values", "table", "upsert"]
+VALID_UPSERT_KEYS = ["update_type", "unique_columns"]
 
 invalid_input_error_code = CommonMessages.ErrorCodes.INVALID_INPUT.value
 
@@ -19,56 +19,59 @@
 # identical to v2's, confirmed, so both variants now share one implementation.
 
 
-def _validate_upsert(upsert):
+def _validate_upsert(logger, upsert):
     if upsert is None:
         return
-    if not isinstance(upsert, Upsert):
+    if not isinstance(upsert, dict):
         raise SkyflowError(SkyflowMessages.Error.INVALID_UPSERT_TYPE_IN_INSERT.value, invalid_input_error_code)
-    if (not isinstance(upsert.unique_columns, list) or not upsert.unique_columns
-            or not all(isinstance(c, str) for c in upsert.unique_columns)):
+    validate_keys(logger, upsert, VALID_UPSERT_KEYS)
+    unique_columns = upsert.get("unique_columns")
+    if (not isinstance(unique_columns, list) or not unique_columns
+            or not all(isinstance(c, str) for c in unique_columns)):
         raise SkyflowError(SkyflowMessages.Error.INVALID_UPSERT_UNIQUE_COLUMNS_IN_INSERT.value, invalid_input_error_code)
-    if upsert.update_type is not None and not isinstance(upsert.update_type, UpsertType):
+    update_type = upsert.get("update_type")
+    if update_type is not None and not isinstance(update_type, UpsertType):
         raise SkyflowError(SkyflowMessages.Error.INVALID_UPSERT_UPDATE_TYPE_IN_INSERT.value, invalid_input_error_code)
 
 
 MAX_INSERT_RECORDS = 10000  # matches Java's v3 Validations.validateInsertRequest (hardcoded, not configurable)
 
 
 def validate_insert_request(logger, request):
-    if not isinstance(request.records, list) or not all(isinstance(r, dict) for r in request.records):
+    if not isinstance(request.values, list) or not all(isinstance(r, dict) for r in request.values):
         raise SkyflowError(SkyflowMessages.Error.INVALID_RECORDS_TYPE_IN_INSERT.value, invalid_input_error_code)
 
-    if not request.records:
+    if not request.values:
         raise SkyflowError(SkyflowMessages.Error.EMPTY_RECORDS_IN_INSERT.value, invalid_input_error_code)
 
-    if len(request.records) > MAX_INSERT_RECORDS:
+    if len(request.values) > MAX_INSERT_RECORDS:
         raise SkyflowError(SkyflowMessages.Error.TOO_MANY_RECORDS_IN_INSERT.value, invalid_input_error_code)
 
     # request.table/record["table"] format and record["values"] emptiness/key/value validity are
     # checked by the controller via the shared BaseVaultController._validate_table_name_if_present()
     # / _validate_field_values() -- not here, to avoid duplicating that logic.
 
-    _validate_upsert(request.upsert)
+    _validate_upsert(logger, request.upsert)
 
-    for record in request.records:
+    for record in request.values:
         validate_keys(logger, record, VALID_INSERT_RECORD_KEYS)
-        _validate_upsert(record.get("upsert"))
+        _validate_upsert(logger, record.get("upsert"))
 
     # table must be set in exactly one place -- request-level (every record) or per-record (no
     # partial mix) -- and upsert must live at that same place (mirrors Java's v3 Validations).
     table_at_request_level = request.table is not None
 
     if table_at_request_level:
-        for record in request.records:
+        for record in request.values:
             if record.get("table") is not None:
                 raise SkyflowError(SkyflowMessages.Error.TABLE_NAME_IN_BOTH_PLACES_IN_INSERT.value, invalid_input_error_code)
     else:
-        for record in request.records:
+        for record in request.values:
             if record.get("table") is None:
                 raise SkyflowError(SkyflowMessages.Error.TABLE_NAME_MISSING_IN_INSERT.value, invalid_input_error_code)
 
     if table_at_request_level:
-        for record in request.records:
+        for record in request.values:
             if record.get("upsert") is not None:
                 raise SkyflowError(SkyflowMessages.Error.RECORD_LEVEL_UPSERT_NOT_ALLOWED_IN_INSERT.value, invalid_input_error_code)
     else:

flowvault/skyflow_flowvault/vault/controller/_vault.py

@@ -23,20 +23,20 @@ def insert(self, request: InsertRequest) -> InsertResponse:
         log_info(SkyflowMessages.Info.VALIDATE_INSERT_REQUEST.value, self._vault_client.get_logger())
         validate_insert_request(self._vault_client.get_logger(), request)
         self._validate_table_name_if_present(request.table)
-        for record in request.records:
+        for record in request.values:
             self._validate_table_name_if_present(record.get("table"))
             self._validate_field_values(record.get("values"))
         log_info(SkyflowMessages.Info.INSERT_REQUEST_RESOLVED.value, self._vault_client.get_logger())
         self._vault_client.initialize_client_configuration()
 
         insert_api = self._vault_client.get_insert_api()
 
-        needs_per_record_table = any(r.get("table") is not None for r in request.records)
-        needs_per_record_upsert = any(r.get("upsert") is not None for r in request.records)
+        needs_per_record_table = any(r.get("table") is not None for r in request.values)
+        needs_per_record_upsert = any(r.get("upsert") is not None for r in request.values)
 
         wire_records = [
             self.__build_wire_record(record, request, needs_per_record_table, needs_per_record_upsert)
-            for record in request.records
+            for record in request.values
         ]
 
         try:
@@ -57,13 +57,11 @@ def insert(self, request: InsertRequest) -> InsertResponse:
             inserted_fields, errors = self.__split_success_and_errors(raw_response.data.records or [], 0, request_id)
         except Exception as e:
             log_error_log(SkyflowMessages.ErrorLogs.INSERT_RECORDS_REJECTED.value, self._vault_client.get_logger())
-            inserted_fields, errors = [], self.__errors_from_exception(e, request.records, 0)
+            inserted_fields, errors = [], self.__errors_from_exception(e, request.values, 0)
 
         log_info(SkyflowMessages.Info.INSERT_SUCCESS.value, self._vault_client.get_logger())
         return InsertResponse(inserted_fields=inserted_fields, errors=errors if errors else None)
 
-    # Not built out this round (insert-only) -- stubs exist so this class stays instantiable
-    # under BaseVaultController's abstract contract.
     def get(self, request):
         raise NotImplementedError("VaultController.get is not implemented yet")
 
@@ -100,9 +98,10 @@ def __build_headers(self):
     def __to_v1_upsert(self, upsert):
         if upsert is None:
             return None
+        update_type = upsert.get("update_type")
         return V1Upsert(
-            update_type=upsert.update_type.value if upsert.update_type else None,
-            unique_columns=upsert.unique_columns,
+            update_type=update_type.value if update_type else None,
+            unique_columns=upsert.get("unique_columns"),
         )
 
     def __extract_request_id(self, headers):

flowvault/skyflow_flowvault/vault/data/_insert_request.py

@@ -1,9 +1,7 @@
 from common.vault.data import BaseInsertRequest
+from skyflow_flowvault.vault.data._upsert import Upsert
 
 
 class InsertRequest(BaseInsertRequest):
-    """table/upsert are request-level defaults; individual records (plain dicts shaped
-    {"values": {...}, "table": ..., "upsert": ...}) may override either."""
-
-    def __init__(self, records, table=None, upsert=None):
-        super().__init__(table, records=records, upsert=upsert)
+    def __init__(self, values: list, table: str = None, upsert: Upsert = None):
+        super().__init__(table, values, upsert=upsert)

flowvault/skyflow_flowvault/vault/data/_upsert.py

@@ -1,6 +1,6 @@
-class Upsert:
-    """Mirrors the wire type V1Upsert. update_type is a skyflow_flowvault.utils.enums.UpsertType value."""
+from typing import Optional, TypedDict
+from skyflow_flowvault.utils.enums import UpsertType
 
-    def __init__(self, update_type=None, unique_columns=None):
-        self.update_type = update_type
-        self.unique_columns = unique_columns
+class Upsert(TypedDict, total=False):
+    update_type: Optional[UpsertType]
+    unique_columns: list