tracdap-runtime 0.6.5__py3-none-any.whl → 0.6.6__py3-none-any.whl

tracdap/rt/_exec/context.py

@@ -55,8 +55,6 @@ class TracContextImpl(_api.TracContext):
             Output views will contain schemas but no data.
     """
 
-    __DEFAULT_TEMPORAL_OBJECTS = False
-
     def __init__(self,
                  model_def: _meta.ModelDefinition,
                  model_class: _api.TracModel.__class__,
@@ -134,53 +132,22 @@ class TracContextImpl(_api.TracContext):
         else:
             return copy.deepcopy(data_view.trac_schema)
 
-    def get_table(self, dataset_name: str, framework, **kwargs) -> _eapi._DATA_FRAMEWORK:  # noqa
-
-        # Support the experimental API data framework syntax
-
-        if framework == _eapi.PANDAS:
-            return self.get_pandas_table(dataset_name, **kwargs)
-        elif framework == _eapi.POLARS:
-            return self.get_polars_table(dataset_name)
-        else:
-            raise _ex.ERuntimeValidation(f"Unsupported data framework [{framework}]")
-
-    def get_pandas_table(self, dataset_name: str, use_temporal_objects: tp.Optional[bool] = None) \
-            -> "_data.pandas.DataFrame":
-
-        _val.require_package("pandas", _data.pandas)
-        _val.validate_signature(self.get_pandas_table, dataset_name, use_temporal_objects)
+    def get_table(self, dataset_name: str, framework: _eapi.DataFramework[_eapi.DATA_API], **framework_args) -> _eapi.DATA_API:
 
-        data_view, schema = self.__get_data_view(dataset_name)
-        part_key = _data.DataPartKey.for_root()
-
-        if use_temporal_objects is None:
-            use_temporal_objects = self.__DEFAULT_TEMPORAL_OBJECTS
-
-        return _data.DataMapping.view_to_pandas(data_view, part_key, schema, use_temporal_objects)
-
-    def get_polars_table(self, dataset_name: str) -> "_data.polars.DataFrame":
-
-        _val.require_package("polars", _data.polars)
-        _val.validate_signature(self.get_polars_table, dataset_name)
-
-        data_view, schema = self.__get_data_view(dataset_name)
-        part_key = _data.DataPartKey.for_root()
-
-        return _data.DataMapping.view_to_polars(data_view, part_key, schema)
-
-    def __get_data_view(self, dataset_name: str):
-
-        _val.validate_signature(self.__get_data_view, dataset_name)
+        _val.validate_signature(self.get_table, dataset_name, framework)
+        _val.require_package(framework.protocol_name, framework.api_type)
 
         self.__val.check_dataset_valid_identifier(dataset_name)
         self.__val.check_dataset_defined_in_model(dataset_name)
         self.__val.check_dataset_available_in_context(dataset_name)
+        self.__val.check_data_framework_args(framework, framework_args)
 
         static_schema = self.__get_static_schema(self.__model_def, dataset_name)
         data_view = self.__local_ctx.get(dataset_name)
         part_key = _data.DataPartKey.for_root()
 
+        converter = _data.DataConverter.for_framework(framework, **framework_args)
+
         self.__val.check_context_object_type(dataset_name, data_view, _data.DataView)
         self.__val.check_dataset_schema_defined(dataset_name, data_view)
         self.__val.check_dataset_part_present(dataset_name, data_view, part_key)
@@ -193,7 +160,18 @@ class TracContextImpl(_api.TracContext):
         else:
             schema = data_view.arrow_schema
 
-        return data_view, schema
+        table = _data.DataMapping.view_to_arrow(data_view, part_key)
+
+        # Data conformance is applied automatically inside the converter, if schema != None
+        return converter.from_internal(table, schema)
+
+    def get_pandas_table(self, dataset_name: str, use_temporal_objects: tp.Optional[bool] = None)  -> "_data.pandas.DataFrame":
+
+        return self.get_table(dataset_name, _eapi.PANDAS, use_temporal_objects=use_temporal_objects)
+
+    def get_polars_table(self, dataset_name: str) -> "_data.polars.DataFrame":
+
+        return self.get_table(dataset_name, _eapi.POLARS)
 
     def put_schema(self, dataset_name: str, schema: _meta.SchemaDefinition):
 
@@ -225,57 +203,28 @@ class TracContextImpl(_api.TracContext):
 
         self.__local_ctx[dataset_name] = updated_view
 
-    def put_table(self, dataset_name: str, dataset: _eapi._DATA_FRAMEWORK, **kwargs):  # noqa
-
-        # Support the experimental API data framework syntax
-
-        if _data.pandas and isinstance(dataset, _data.pandas.DataFrame):
-            self.put_pandas_table(dataset_name, dataset)
-        elif _data.polars and isinstance(dataset, _data.polars.DataFrame):
-            self.put_polars_table(dataset_name, dataset)
-        else:
-            raise _ex.ERuntimeValidation(f"Unsupported data framework[{type(dataset)}]")
-
-    def put_pandas_table(self, dataset_name: str, dataset: "_data.pandas.DataFrame"):
-
-        _val.require_package("pandas", _data.pandas)
-        _val.validate_signature(self.put_pandas_table, dataset_name, dataset)
-
-        part_key = _data.DataPartKey.for_root()
-        data_view, schema = self.__put_data_view(dataset_name, part_key, dataset, _data.pandas.DataFrame)
+    def put_table(
+            self, dataset_name: str, dataset: _eapi.DATA_API,
+            framework: tp.Optional[_eapi.DataFramework[_eapi.DATA_API]] = None,
+            **framework_args):
 
-        # Data conformance is applied inside these conversion functions
+        _val.validate_signature(self.put_table, dataset_name, dataset, framework)
 
-        updated_item = _data.DataMapping.pandas_to_item(dataset, schema)
-        updated_view = _data.DataMapping.add_item_to_view(data_view, part_key, updated_item)
+        if framework is None:
+            framework = _data.DataConverter.get_framework(dataset)
 
-        self.__local_ctx[dataset_name] = updated_view
-
-    def put_polars_table(self, dataset_name: str, dataset: "_data.polars.DataFrame"):
-
-        _val.require_package("polars", _data.polars)
-        _val.validate_signature(self.put_polars_table, dataset_name, dataset)
-
-        part_key = _data.DataPartKey.for_root()
-        data_view, schema = self.__put_data_view(dataset_name, part_key, dataset, _data.polars.DataFrame)
-
-        # Data conformance is applied inside these conversion functions
-
-        updated_item = _data.DataMapping.polars_to_item(dataset, schema)
-        updated_view = _data.DataMapping.add_item_to_view(data_view, part_key, updated_item)
-
-        self.__local_ctx[dataset_name] = updated_view
-
-    def __put_data_view(self, dataset_name: str, part_key: _data.DataPartKey, dataset: tp.Any, framework: type):
-
-        _val.validate_signature(self.__put_data_view, dataset_name, part_key, dataset, framework)
+        _val.require_package(framework.protocol_name, framework.api_type)
 
         self.__val.check_dataset_valid_identifier(dataset_name)
         self.__val.check_dataset_is_model_output(dataset_name)
-        self.__val.check_provided_dataset_type(dataset, framework)
+        self.__val.check_provided_dataset_type(dataset, framework.api_type)
+        self.__val.check_data_framework_args(framework, framework_args)
 
         static_schema = self.__get_static_schema(self.__model_def, dataset_name)
         data_view = self.__local_ctx.get(dataset_name)
+        part_key = _data.DataPartKey.for_root()
+
+        converter = _data.DataConverter.for_framework(framework)
 
         if data_view is None:
             if static_schema is not None:
@@ -294,7 +243,21 @@ class TracContextImpl(_api.TracContext):
         else:
             schema = data_view.arrow_schema
 
-        return data_view, schema
+        # Data conformance is applied automatically inside the converter, if schema != None
+        table = converter.to_internal(dataset, schema)
+        item = _data.DataItem(schema, table)
+
+        updated_view = _data.DataMapping.add_item_to_view(data_view, part_key, item)
+
+        self.__local_ctx[dataset_name] = updated_view
+
+    def put_pandas_table(self, dataset_name: str, dataset: "_data.pandas.DataFrame"):
+
+        self.put_table(dataset_name, dataset, _eapi.PANDAS)
+
+    def put_polars_table(self, dataset_name: str, dataset: "_data.polars.DataFrame"):
+
+        self.put_table(dataset_name, dataset, _eapi.POLARS)
 
     def log(self) -> logging.Logger:
 
@@ -335,7 +298,7 @@ class TracDataContextImpl(TracContextImpl, _eapi.TracDataContext):
     def __init__(
             self, model_def: _meta.ModelDefinition, model_class: _api.TracModel.__class__,
             local_ctx: tp.Dict[str, tp.Any], dynamic_outputs: tp.List[str],
-            storage_map: tp.Dict[str, tp.Union[_eapi.TracFileStorage]],
+            storage_map: tp.Dict[str, tp.Union[_eapi.TracFileStorage, _eapi.TracDataStorage]],
             checkout_directory: pathlib.Path = None):
 
         super().__init__(model_def, model_class, local_ctx, dynamic_outputs, checkout_directory)
@@ -358,8 +321,27 @@ class TracDataContextImpl(TracContextImpl, _eapi.TracDataContext):
 
         return self.__storage_map[storage_key]
 
-    def get_data_storage(self, storage_key: str) -> None:
-        raise _ex.ERuntimeValidation("Data storage API not available yet")
+    def get_data_storage(
+            self, storage_key: str,
+            framework: _eapi.DataFramework[_eapi.DATA_API],
+            **framework_args) -> _eapi.TracDataStorage[_eapi.DATA_API]:
+
+        _val.validate_signature(self.get_file_storage, storage_key)
+
+        self.__val.check_storage_valid_identifier(storage_key)
+        self.__val.check_storage_available(self.__storage_map, storage_key)
+        self.__val.check_storage_type(self.__storage_map, storage_key, _eapi.TracDataStorage)
+        self.__val.check_data_framework_args(framework, framework_args)
+
+        storage = self.__storage_map[storage_key]
+        converter = _data.DataConverter.for_framework(framework, **framework_args)
+
+        # Create a shallow copy of the storage impl with a converter for the requested data framework
+        # At some point we will need a storage factory class, bc the internal data API can also be different
+        storage = copy.copy(storage)
+        storage._TracDataStorageImpl__converter = converter
+
+        return storage
 
     def add_data_import(self, dataset_name: str):
 
@@ -372,15 +354,30 @@ class TracDataContextImpl(TracContextImpl, _eapi.TracDataContext):
         self.__local_ctx[dataset_name] = _data.DataView.create_empty()
         self.__dynamic_outputs.append(dataset_name)
 
-    def set_source_metadata(self, dataset_name: str, storage_key: str, source_info: _eapi.FileStat):
+    def set_source_metadata(self, dataset_name: str, storage_key: str, source_info: tp.Union[_eapi.FileStat, str]):
 
-        _val.validate_signature(self.add_data_import, dataset_name, storage_key, source_info)
+        _val.validate_signature(self.set_source_metadata, dataset_name, storage_key, source_info)
+
+        self.__val.check_dataset_valid_identifier(dataset_name)
+        self.__val.check_dataset_available_in_context(dataset_name)
+        self.__val.check_storage_valid_identifier(storage_key)
+        self.__val.check_storage_available(self.__storage_map, storage_key)
+
+        storage = self.__storage_map[storage_key]
+
+        if isinstance(storage, _eapi.TracFileStorage):
+            if not isinstance(source_info, _eapi.FileStat):
+                self.__val.report_public_error(f"Expected storage_info to be a FileStat, [{storage_key}] refers to file storage")
+
+        if isinstance(storage, _eapi.TracDataStorage):
+            if not isinstance(source_info, str):
+                self.__val.report_public_error(f"Expected storage_info to be a table name, [{storage_key}] refers to dadta storage")
 
         pass  # Not implemented yet, only required when imports are sent back to the platform
 
     def set_attribute(self, dataset_name: str, attribute_name: str, value: tp.Any):
 
-        _val.validate_signature(self.add_data_import, dataset_name, attribute_name, value)
+        _val.validate_signature(self.set_attribute, dataset_name, attribute_name, value)
 
         pass  # Not implemented yet, only required when imports are sent back to the platform
 
@@ -531,13 +528,132 @@ class TracFileStorageImpl(_eapi.TracFileStorage):
         super().write_bytes(storage_path, data)
 
 
+class TracDataStorageImpl(_eapi.TracDataStorage[_eapi.DATA_API]):
+
+    def __init__(
+            self, storage_key: str, storage_impl: _storage.IDataStorageBase[_data.T_INTERNAL_DATA, _data.T_INTERNAL_SCHEMA],
+            data_converter: _data.DataConverter[_eapi.DATA_API, _data.T_INTERNAL_DATA, _data.T_INTERNAL_SCHEMA],
+            write_access: bool, checkout_directory):
+
+        self.__storage_key = storage_key
+        self.__converter = data_converter
+
+        self.__has_table = lambda tn: storage_impl.has_table(tn)
+        self.__list_tables = lambda: storage_impl.list_tables()
+        self.__read_table = lambda tn: storage_impl.read_table(tn)
+        self.__native_read_query = lambda q, ps: storage_impl.native_read_query(q, **ps)
+
+        if write_access:
+            self.__create_table = lambda tn, s: storage_impl.create_table(tn, s)
+            self.__write_table = lambda tn, ds: storage_impl.write_table(tn, ds)
+        else:
+            self.__create_table = None
+            self.__write_table = None
+
+        self.__log = _util.logger_for_object(self)
+        self.__val = TracStorageValidator(self.__log, checkout_directory, self.__storage_key)
+
+    def has_table(self, table_name: str) -> bool:
+
+        _val.validate_signature(self.has_table, table_name)
+
+        self.__val.check_operation_available(self.has_table, self.__has_table)
+        self.__val.check_table_name_is_valid(table_name)
+        self.__val.check_storage_path_is_valid(table_name)
+
+        try:
+            return self.__has_table(table_name)
+        except _ex.EStorageRequest as e:
+            self.__val.report_public_error(e)
+
+    def list_tables(self) -> tp.List[str]:
+
+        _val.validate_signature(self.list_tables)
+
+        self.__val.check_operation_available(self.list_tables, self.__list_tables)
+
+        try:
+            return self.__list_tables()
+        except _ex.EStorageRequest as e:
+            self.__val.report_public_error(e)
+
+    def create_table(self, table_name: str, schema: _api.SchemaDefinition):
+
+        _val.validate_signature(self.create_table, table_name, schema)
+
+        self.__val.check_operation_available(self.create_table, self.__create_table)
+        self.__val.check_table_name_is_valid(table_name)
+        self.__val.check_storage_path_is_valid(table_name)
+
+        arrow_schema = _data.DataMapping.trac_to_arrow_schema(schema)
+
+        try:
+            self.__create_table(table_name, arrow_schema)
+        except _ex.EStorageRequest as e:
+            self.__val.report_public_error(e)
+
+    def read_table(self, table_name: str) -> _eapi.DATA_API:
+
+        _val.validate_signature(self.read_table, table_name)
+
+        self.__val.check_operation_available(self.read_table, self.__read_table)
+        self.__val.check_table_name_is_valid(table_name)
+        self.__val.check_table_name_not_reserved(table_name)
+
+        try:
+            raw_data = self.__read_table(table_name)
+            return self.__converter.from_internal(raw_data)
+
+        except _ex.EStorageRequest as e:
+            self.__val.report_public_error(e)
+
+    def native_read_query(self, query: str, **parameters) -> _eapi.DATA_API:
+
+        _val.validate_signature(self.native_read_query, query, **parameters)
+
+        self.__val.check_operation_available(self.native_read_query, self.__native_read_query)
+
+        # TODO: validate query and parameters
+        # Some validation is performed by the impl
+
+        try:
+            raw_data = self.__native_read_query(query, **parameters)
+            return self.__converter.from_internal(raw_data)
+
+        except _ex.EStorageRequest as e:
+            self.__val.report_public_error(e)
+
+    def write_table(self, table_name: str, dataset: _eapi.DATA_API):
+
+        _val.validate_signature(self.write_table, table_name, dataset)
+
+        self.__val.check_operation_available(self.read_table, self.__read_table)
+        self.__val.check_table_name_is_valid(table_name)
+        self.__val.check_table_name_not_reserved(table_name)
+        self.__val.check_provided_dataset_type(dataset, self.__converter.framework.api_type)
+
+        try:
+            raw_data = self.__converter.to_internal(dataset)
+            self.__write_table(table_name, raw_data)
+
+        except _ex.EStorageRequest as e:
+            self.__val.report_public_error(e)
+
+
 class TracContextErrorReporter:
 
+    _VALID_IDENTIFIER = re.compile("^[a-zA-Z_]\\w*$",)
+    _RESERVED_IDENTIFIER = re.compile("^(trac_|_)\\w*")
+
     def __init__(self, log: logging.Logger, checkout_directory: pathlib.Path):
 
         self.__log = log
         self.__checkout_directory = checkout_directory
 
+    def report_public_error(self, exception: Exception):
+
+        self._report_error(str(exception), exception)
+
     def _report_error(self, message, cause: Exception = None):
 
         full_stack = traceback.extract_stack()
@@ -554,11 +670,18 @@ class TracContextErrorReporter:
         else:
             raise _ex.ERuntimeValidation(message)
 
+    @staticmethod
+    def _type_name(type_: type):
 
-class TracContextValidator(TracContextErrorReporter):
+        module = type_.__module__
+
+        if module is None or module == str.__class__.__module__ or module == tp.__name__:
+            return _val.type_name(type_, False)
+        else:
+            return _val.type_name(type_, True)
 
-    __VALID_IDENTIFIER = re.compile("^[a-zA-Z_]\\w*$",)
-    __RESERVED_IDENTIFIER = re.compile("^(trac_|_)\\w*")
+
+class TracContextValidator(TracContextErrorReporter):
 
     def __init__(
             self, log: logging.Logger,
@@ -578,7 +701,7 @@ class TracContextValidator(TracContextErrorReporter):
         if param_name is None:
             self._report_error(f"Parameter name is null")
 
-        if not self.__VALID_IDENTIFIER.match(param_name):
+        if not self._VALID_IDENTIFIER.match(param_name):
             self._report_error(f"Parameter name {param_name} is not a valid identifier")
 
     def check_param_defined_in_model(self, param_name: str):
@@ -596,7 +719,7 @@ class TracContextValidator(TracContextErrorReporter):
         if dataset_name is None:
             self._report_error(f"Dataset name is null")
 
-        if not self.__VALID_IDENTIFIER.match(dataset_name):
+        if not self._VALID_IDENTIFIER.match(dataset_name):
             self._report_error(f"Dataset name {dataset_name} is not a valid identifier")
 
     def check_dataset_not_defined_in_model(self, dataset_name: str):
@@ -710,12 +833,39 @@ class TracContextValidator(TracContextErrorReporter):
                 f"The object referenced by [{item_name}] in the current context has the wrong type" +
                 f" (expected {expected_type_name}, got {actual_type_name})")
 
+    def check_data_framework_args(self, framework: _eapi.DataFramework, framework_args: tp.Dict[str, tp.Any]):
+
+        expected_args = _data.DataConverter.get_framework_args(framework)
+        unexpected_args = list(filter(lambda arg: arg not in expected_args, framework_args.keys()))
+
+        if any(unexpected_args):
+            unknown_args = ", ".join(unexpected_args)
+            self._report_error(f"Using [{framework}], some arguments were not recognized: [{unknown_args}]")
+
+        for arg_name, arg_type in expected_args.items():
+
+            arg_value = framework_args.get(arg_name)
+
+            if _val.check_type(arg_type, arg_value):
+                continue
+
+            if arg_value is None:
+                self._report_error(f"Using [{framework}], required argument [{arg_name}] is missing")
+
+            else:
+                expected_type_name = self._type_name(arg_type)
+                actual_type_name = self._type_name(type(arg_value))
+
+                self._report_error(
+                    f"Using [{framework}], argument [{arg_name}] has the wrong type" +
+                    f" (expected {expected_type_name}, got {actual_type_name})")
+
     def check_storage_valid_identifier(self, storage_key):
 
         if storage_key is None:
             self._report_error(f"Storage key is null")
 
-        if not self.__VALID_IDENTIFIER.match(storage_key):
+        if not self._VALID_IDENTIFIER.match(storage_key):
             self._report_error(f"Storage key {storage_key} is not a valid identifier")
 
     def check_storage_available(self, storage_map: tp.Dict, storage_key: str):
@@ -737,16 +887,6 @@ class TracContextValidator(TracContextErrorReporter):
             else:
                 self._report_error(f"Storage key [{storage_key}] refers to file storage, not data storage")
 
-    @staticmethod
-    def _type_name(type_: type):
-
-        module = type_.__module__
-
-        if module is None or module == str.__class__.__module__:
-            return type_.__qualname__
-
-        return module + '.' + type_.__name__
-
 
 class TracStorageValidator(TracContextErrorReporter):
 
@@ -777,3 +917,30 @@ class TracStorageValidator(TracContextErrorReporter):
 
         if _val.StorageValidator.storage_path_is_empty(storage_path):
             self._report_error(f"Storage path [{storage_path}] is not allowed")
+
+    def check_table_name_is_valid(self, table_name: str):
+
+        if table_name is None:
+            self._report_error(f"Table name is null")
+
+        if not self._VALID_IDENTIFIER.match(table_name):
+            self._report_error(f"Table name {table_name} is not a valid identifier")
+
+    def check_table_name_not_reserved(self, table_name: str):
+
+        if self._RESERVED_IDENTIFIER.match(table_name):
+            self._report_error(f"Table name {table_name} is a reserved identifier")
+
+    def check_provided_dataset_type(self, dataset: tp.Any, expected_type: type):
+
+        if dataset is None:
+            self._report_error(f"Provided dataset is null")
+
+        if not isinstance(dataset, expected_type):
+
+            expected_type_name = self._type_name(expected_type)
+            actual_type_name = self._type_name(type(dataset))
+
+            self._report_error(
+                f"Provided dataset is the wrong type" +
+                f" (expected {expected_type_name}, got {actual_type_name})")