SK-2954: Add ISkyflow/IVaultController/IVaultClient interfaces, type insert()

Splits BaseSkyflow, BaseVaultController, and BaseVaultClient into a pure
interface (ISkyflow/IVaultController/IVaultClient, declaring the contract via
ABC + abstractmethod) plus a base class implementing the shared logic, so
future variant-specific overrides have a clear contract to satisfy. Also adds
type hints to insert() at every layer (BaseInsertRequest/BaseInsertResponse in
the abstract method, each variant's own InsertRequest/InsertResponse in their
concrete override), and renames base_vault.py to base_vault_controller.py to
match its class name.

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

Author: saileshwar-skyflow

common/client/base_skyflow.py

@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from collections import OrderedDict
 from functools import partial
 
@@ -13,7 +14,79 @@
 )
 
 
-class Skyflow:
+class ISkyflow(ABC):
+    @classmethod
+    @abstractmethod
+    def builder(cls):
+        raise NotImplementedError
+
+    @abstractmethod
+    def add_vault_config(self, config):
+        raise NotImplementedError
+
+    @abstractmethod
+    def remove_vault_config(self, vault_id):
+        raise NotImplementedError
+
+    @abstractmethod
+    def update_vault_config(self, config):
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_vault_config(self, vault_id):
+        raise NotImplementedError
+
+    @abstractmethod
+    def add_connection_config(self, config):
+        raise NotImplementedError
+
+    @abstractmethod
+    def remove_connection_config(self, connection_id):
+        raise NotImplementedError
+
+    @abstractmethod
+    def update_connection_config(self, config):
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_connection_config(self, connection_id):
+        raise NotImplementedError
+
+    @abstractmethod
+    def add_skyflow_credentials(self, credentials):
+        raise NotImplementedError
+
+    @abstractmethod
+    def update_skyflow_credentials(self, credentials):
+        raise NotImplementedError
+
+    @abstractmethod
+    def set_log_level(self, log_level):
+        raise NotImplementedError
+
+    @abstractmethod
+    def update_log_level(self, log_level):
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_log_level(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def vault(self, vault_id=None):
+        raise NotImplementedError
+
+    @abstractmethod
+    def connection(self, connection_id=None):
+        raise NotImplementedError
+
+    @abstractmethod
+    def detect(self, vault_id=None):
+        raise NotImplementedError
+
+
+class BaseSkyflow(ISkyflow):
+
     def __init__(self, builder):
         self.__builder = builder
         log_info(self.__builder._skyflow_messages.Info.CLIENT_INITIALIZED.value, self.__builder.get_logger())
@@ -87,9 +160,10 @@ def detect(self, vault_id=None):
         vault_config = self.__builder.get_vault_config(vault_id)
         return vault_config.get(OptionField.DETECT_CONTROLLER)
 
-    class Builder:
+    class Builder(ABC):
         # -- hooks, filled in per-variant by make_skyflow_class() -- left None here so using
-        # this template directly (rather than through make_skyflow_class()) fails fast.
+        # this template directly (rather than through make_skyflow_class()) fails fast, with a
+        # clear message (see _REQUIRED_HOOKS check in __init__ below).
         _vault_client_cls = None
         _vault_controller_cls = None
         _connection_cls = None
@@ -106,7 +180,21 @@ class Builder:
         _validate_credentials = None
         _set_active_log_level = None
 
+        # Connection/Detect support and their validators are legitimately optional per variant --
+        # everything else must be supplied by make_skyflow_class() before this Builder is usable.
+        _REQUIRED_HOOKS = (
+            '_vault_client_cls', '_vault_controller_cls', '_logger_cls', '_default_log_level',
+            '_skyflow_messages', '_skyflow_cls', '_validate_vault_config',
+            '_validate_update_vault_config', '_validate_log_level', '_validate_credentials',
+        )
+
         def __init__(self):
+            missing = [hook for hook in self._REQUIRED_HOOKS if getattr(self, hook) is None]
+            if missing:
+                raise NotImplementedError(
+                    "BaseSkyflow.Builder is an interface template -- build a concrete Skyflow "
+                    f"class via make_skyflow_class() instead of using it directly. Missing: {', '.join(missing)}"
+                )
             self.__vault_configs = OrderedDict()
             self.__vault_list = list()
             self.__connection_configs = OrderedDict()
@@ -326,7 +414,7 @@ def make_skyflow_class(*, vault_client_cls, vault_controller_cls, skyflow_messag
         '_validate_credentials': staticmethod(validate_credentials),
         '_set_active_log_level': staticmethod(set_active_log_level) if set_active_log_level else None,
     }
-    variant_builder = type('Builder', (Skyflow.Builder,), builder_attrs)
-    variant_skyflow = type('Skyflow', (Skyflow,), {'Builder': variant_builder})
+    variant_builder = type('Builder', (BaseSkyflow.Builder,), builder_attrs)
+    variant_skyflow = type('Skyflow', (BaseSkyflow,), {'Builder': variant_builder})
     variant_builder._skyflow_cls = variant_skyflow
     return variant_skyflow

common/tests/client/test_base_skyflow.py

@@ -3,7 +3,7 @@
 from common.errors import SkyflowError
 from common.utils import LogLevel, SkyflowMessages
 from common.utils.logger import Logger
-from common.client.base_skyflow import make_skyflow_class
+from common.client.base_skyflow import BaseSkyflow, make_skyflow_class
 
 
 class FakeVaultClient:
@@ -185,5 +185,17 @@ def test_make_skyflow_class_requires_connection_validators_when_connection_cls_g
             )
 
 
+class TestBaseSkyflowInterface(unittest.TestCase):
+    def test_using_the_template_builder_directly_raises_with_a_clear_message(self):
+        with self.assertRaises(NotImplementedError) as ctx:
+            BaseSkyflow.Builder()
+        self.assertIn("make_skyflow_class()", str(ctx.exception))
+        self.assertIn("_vault_client_cls", str(ctx.exception))
+
+    def test_make_skyflow_class_produced_builder_constructs_fine(self):
+        Skyflow = make_fake_skyflow()
+        self.assertIsInstance(Skyflow.builder(), BaseSkyflow.Builder)
+
+
 if __name__ == "__main__":
     unittest.main()

common/tests/vault/test_base_vault.py …ests/vault/test_base_vault_controller.pycommon/tests/vault/test_base_vault.py renamed to common/tests/vault/test_base_vault_controller.py common/tests/vault/test_base_vault.py renamed to common/tests/vault/test_base_vault_controller.py

@@ -2,7 +2,7 @@
 
 from common.errors import SkyflowError
 from common.utils import SkyflowMessages
-from common.vault.base_vault import BaseVaultController
+from common.vault.base_vault_controller import BaseVaultController
 
 
 class DummyVaultController(BaseVaultController):

common/vault/base_vault_client.py

@@ -6,11 +6,17 @@
 from common.utils.constants import OptionField, CredentialField, ConfigField
 
 
-class BaseVaultClient(ABC):
-    """Shared credential resolution, vault-URL resolution, and bearer-token fetch/cache/expiry
-    logic. Uses single-underscore attributes deliberately -- double-underscore would name-mangle
-    per-subclass and break cross-class state sharing."""
+class IVaultClient(ABC):
+    @abstractmethod
+    def resolve_vault_url(self, cluster_id, env, vault_id, logger=None):
+        raise NotImplementedError
+
+    @abstractmethod
+    def initialize_api_client(self, vault_url, bearer_token):
+        raise NotImplementedError
+
 
+class BaseVaultClient(IVaultClient):
     def __init__(self, config):
         self._config = config
         self._common_skyflow_credentials = None
@@ -52,19 +58,6 @@ def initialize_client_configuration(self):
         if needs_reinit:
             self.initialize_api_client(self._vault_url, bearer_token)
 
-    @abstractmethod
-    def resolve_vault_url(self, cluster_id, env, vault_id, logger=None):
-        """Per-variant hook: different vault types are hosted on different subdomains for the
-        same cluster_id/env (v2: vault.skyflowapis.<tld>; v3: skyvault.skyflowapis.<tld>)."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def initialize_api_client(self, vault_url, bearer_token):
-        """Construct the variant's generated API client into self._api_client. v2 bakes
-        bearer_token into a refreshable callable; v3's client has no token param at all, so auth
-        is injected per-call instead (see get_current_bearer_token)."""
-        raise NotImplementedError
-
     def get_current_bearer_token(self):
         return self._bearer_token
 

common/vault/base_vault.py common/vault/base_vault_controller.pycommon/vault/base_vault.py renamed to common/vault/base_vault_controller.py

@@ -2,11 +2,39 @@
 
 from common.errors import SkyflowError
 from common.utils import SkyflowMessages as _CommonSkyflowMessages
+from common.vault.data import BaseInsertRequest, BaseInsertResponse
 
 _INVALID_INPUT_ERROR_CODE = _CommonSkyflowMessages.ErrorCodes.INVALID_INPUT.value
 
 
-class BaseVaultController(ABC):
+class IVaultController(ABC):
+
+    @abstractmethod
+    def insert(self, request: BaseInsertRequest) -> BaseInsertResponse:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get(self, request):
+        raise NotImplementedError
+
+    @abstractmethod
+    def update(self, request):
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete(self, request):
+        raise NotImplementedError
+
+    @abstractmethod
+    def query(self, request):
+        raise NotImplementedError
+
+    @abstractmethod
+    def detokenize(self, request):
+        raise NotImplementedError
+
+
+class BaseVaultController(IVaultController):
 
     _skyflow_messages = None
 
@@ -37,27 +65,3 @@ def _validate_field_values(self, values):
             #         self._skyflow_messages.Error.EMPTY_VALUE_IN_INSERT_DATA.value,
             #         _INVALID_INPUT_ERROR_CODE,
             #     )
-
-    @abstractmethod
-    def insert(self, request):
-        raise NotImplementedError
-
-    @abstractmethod
-    def get(self, request):
-        raise NotImplementedError
-
-    @abstractmethod
-    def update(self, request):
-        raise NotImplementedError
-
-    @abstractmethod
-    def delete(self, request):
-        raise NotImplementedError
-
-    @abstractmethod
-    def query(self, request):
-        raise NotImplementedError
-
-    @abstractmethod
-    def detokenize(self, request):
-        raise NotImplementedError

flowvault/skyflow_flowvault/vault/controller/_vault.py

@@ -3,12 +3,12 @@
 from common.utils import SkyflowMessages as CommonMessages
 from common.utils.constants import SKY_META_DATA_HEADER
 from common.utils.logger import log_info, log_error_log
-from common.vault.base_vault import BaseVaultController
+from common.vault.base_vault_controller import BaseVaultController
 from skyflow_flowvault.generated.rest import V1InsertRecordData, V1Upsert
 from skyflow_flowvault.generated.rest.core import ApiError
 from skyflow_flowvault.utils import SkyflowMessages, get_metrics
 from skyflow_flowvault.utils.validations import validate_insert_request
-from skyflow_flowvault.vault.data import InsertResponse
+from skyflow_flowvault.vault.data import InsertRequest, InsertResponse
 
 REQUEST_ID_HEADER = "x-request-id"
 
@@ -19,7 +19,7 @@ class VaultController(BaseVaultController):
     def __init__(self, vault_client):
         super().__init__(vault_client)
 
-    def insert(self, request):
+    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)

flowvault/tests/utils/validations/test__validations.py

@@ -84,7 +84,7 @@ def test_table_missing_from_one_record_raises(self):
     # dict, are now validated by the controller via the shared
     # BaseVaultController._validate_field_values() -- see test__vault.py's
     # test_insert_raises_on_empty_key/_on_empty_value/_on_non_dict_values/_on_empty_values, and
-    # common/tests/vault/test_base_vault.py for the shared helper's own unit tests.
+    # common/tests/vault/test_base_vault_controller.py for the shared helper's own unit tests.
 
     def test_falsy_non_string_values_are_valid(self):
         """0, False, [], {} are all legitimate values -- only None/empty-string should raise
@@ -123,7 +123,7 @@ def test_records_must_not_be_empty(self):
     # Table name format (non-empty string if provided) is now validated by the controller via
     # the shared BaseVaultController._validate_table_name_if_present() -- see
     # test__vault.py's test_insert_raises_on_invalid_table_name and
-    # common/tests/vault/test_base_vault.py for the shared helper's own unit tests.
+    # common/tests/vault/test_base_vault_controller.py for the shared helper's own unit tests.
 
     def test_table_is_optional_when_every_record_has_its_own(self):
         request = InsertRequest(records=[dict(values={"a": 1}, table="t2")])

flowvault/tests/vault/controller/test__vault.py

@@ -64,7 +64,7 @@ def test_insert_raises_for_invalid_request(self):
 
     # ------------------------------------------------------------------ #
     # shared BaseVaultController validation helpers, exercised end-to-end via insert()
-    # (unit-tested in isolation in common/tests/vault/test_base_vault.py)
+    # (unit-tested in isolation in common/tests/vault/test_base_vault_controller.py)
     # ------------------------------------------------------------------ #
 
     def test_insert_raises_on_empty_key(self):

v2/skyflow/vault/controller/__init__.py

@@ -3,6 +3,6 @@
 from ._detect import Detect
 
 # Public backward-compatible name -- existing consumers do `from skyflow.vault.controller import
-# Vault`; VaultController is the canonical internal name (extends common.vault.base_vault's
+# Vault`; VaultController is the canonical internal name (extends common.vault.base_vault_controller's
 # BaseVaultController), but the old public name must keep resolving to the exact same class.
 Vault = VaultController

v2/skyflow/vault/controller/_vault.py

@@ -2,7 +2,7 @@
 import json
 import os
 from typing import Optional
-from common.vault.base_vault import BaseVaultController
+from common.vault.base_vault_controller import BaseVaultController
 from skyflow.generated.rest import V1FieldRecords, V1BatchRecord, V1TokenizeRecordRequest, \
     V1DetokenizeRecordRequest
 from skyflow.generated.rest.core.file import File
@@ -15,7 +15,7 @@
 from skyflow.utils.logger import log_info, log_error_log
 from skyflow.utils.validations import validate_insert_request, validate_delete_request, validate_query_request, \
     validate_get_request, validate_update_request, validate_detokenize_request, validate_tokenize_request, validate_file_upload_request
-from skyflow.vault.data import InsertRequest, UpdateRequest, DeleteRequest, GetRequest, QueryRequest, FileUploadRequest, FileUploadResponse
+from skyflow.vault.data import InsertRequest, InsertResponse, UpdateRequest, DeleteRequest, GetRequest, QueryRequest, FileUploadRequest, FileUploadResponse
 from skyflow.vault.tokens import DetokenizeRequest, TokenizeRequest
 
 class VaultController(BaseVaultController):
@@ -94,7 +94,7 @@ def __get_file_for_file_upload(self, request: FileUploadRequest) -> Optional[Fil
     def __get_headers(self):
         return {SKY_META_DATA_HEADER: json.dumps(get_metrics())}
 
-    def insert(self, request: InsertRequest):
+    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)
         for item in request.values: