frogml-core 0.0.71__py3-none-any.whl → 0.0.73__py3-none-any.whl

frogml_proto/qwak/service_discovery/service_discovery_location_service_pb2_grpc.py

@@ -0,0 +1,231 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+
+from frogml_proto.qwak.service_discovery import service_discovery_location_service_pb2 as qwak_dot_service__discovery_dot_service__discovery__location__service__pb2
+
+
+class LocationDiscoveryServiceStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.GetOfflineServingUrl = channel.unary_unary(
+                '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetOfflineServingUrl',
+                request_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+                response_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+                )
+        self.GetDistributionManagerUrl = channel.unary_unary(
+                '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetDistributionManagerUrl',
+                request_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+                response_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+                )
+        self.GetAnalyticsEngineUrl = channel.unary_unary(
+                '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetAnalyticsEngineUrl',
+                request_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+                response_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+                )
+        self.GetMetricsGatewayUrl = channel.unary_unary(
+                '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetMetricsGatewayUrl',
+                request_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+                response_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+                )
+        self.GetFeaturesOperatorUrl = channel.unary_unary(
+                '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetFeaturesOperatorUrl',
+                request_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+                response_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+                )
+        self.GetHostingGatewayUrl = channel.unary_unary(
+                '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetHostingGatewayUrl',
+                request_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+                response_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+                )
+
+
+class LocationDiscoveryServiceServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def GetOfflineServingUrl(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def GetDistributionManagerUrl(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def GetAnalyticsEngineUrl(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def GetMetricsGatewayUrl(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def GetFeaturesOperatorUrl(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def GetHostingGatewayUrl(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_LocationDiscoveryServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'GetOfflineServingUrl': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetOfflineServingUrl,
+                    request_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.FromString,
+                    response_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.SerializeToString,
+            ),
+            'GetDistributionManagerUrl': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetDistributionManagerUrl,
+                    request_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.FromString,
+                    response_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.SerializeToString,
+            ),
+            'GetAnalyticsEngineUrl': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetAnalyticsEngineUrl,
+                    request_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.FromString,
+                    response_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.SerializeToString,
+            ),
+            'GetMetricsGatewayUrl': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetMetricsGatewayUrl,
+                    request_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.FromString,
+                    response_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.SerializeToString,
+            ),
+            'GetFeaturesOperatorUrl': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetFeaturesOperatorUrl,
+                    request_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.FromString,
+                    response_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.SerializeToString,
+            ),
+            'GetHostingGatewayUrl': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetHostingGatewayUrl,
+                    request_deserializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.FromString,
+                    response_serializer=qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'qwak.service_discovery.location.v0.LocationDiscoveryService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+ # This class is part of an EXPERIMENTAL API.
+class LocationDiscoveryService(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def GetOfflineServingUrl(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetOfflineServingUrl',
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def GetDistributionManagerUrl(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetDistributionManagerUrl',
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def GetAnalyticsEngineUrl(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetAnalyticsEngineUrl',
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def GetMetricsGatewayUrl(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetMetricsGatewayUrl',
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def GetFeaturesOperatorUrl(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetFeaturesOperatorUrl',
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def GetHostingGatewayUrl(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/qwak.service_discovery.location.v0.LocationDiscoveryService/GetHostingGatewayUrl',
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequest.SerializeToString,
+            qwak_dot_service__discovery_dot_service__discovery__location__service__pb2.GetServingUrlRequestResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

frogml_services_mock/mocks/frogml_mocks.py

@@ -64,6 +64,9 @@ from frogml_services_mock.mocks.job_registry_service_api import (
 from frogml_services_mock.mocks.kube_captain_service_api import (
     KubeCaptainServiceApiMock,
 )
+from frogml_services_mock.mocks.location_discovery_service_api import (
+    LocationDiscoveryServiceApiMock,
+)
 from frogml_services_mock.mocks.logging_service import LoggingServiceApiMock
 from frogml_services_mock.mocks.model_management_service import (
     ModelsManagementServiceMock,
@@ -131,3 +134,4 @@ class FrogmlMocks:
     prompt_manager_service: PromptManagerServiceMock
     model_version_manager_service: ModelVersionManagerServiceMock
     repository_service: RepositoryServiceMock
+    location_discovery_service: LocationDiscoveryServiceApiMock

frogml_services_mock/mocks/location_discovery_service_api.py

@@ -0,0 +1,104 @@
+from typing import Dict, Optional
+
+import grpc
+from frogml_proto.qwak.service_discovery.service_discovery_location_pb2 import (
+    ServiceLocationDescriptor,
+)
+from frogml_proto.qwak.service_discovery.service_discovery_location_service_pb2 import (
+    GetServingUrlRequestResponse,
+)
+from frogml_proto.qwak.service_discovery.service_discovery_location_service_pb2_grpc import (
+    LocationDiscoveryServiceServicer,
+)
+
+
+class LocationDiscoveryServiceApiMock(LocationDiscoveryServiceServicer):
+    """
+    Mock implementation of the LocationDiscoveryService for testing SDK behavior.
+    Allows setting mock responses and optional error codes for each endpoint.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._responses: Dict[str, Optional[ServiceLocationDescriptor]] = {}
+        self._error_codes: Dict[str, grpc.StatusCode] = {}
+
+    def _set_mock(
+        self,
+        key: str,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._responses[key] = response
+        self._error_codes[key] = error_code
+
+    def _handle(
+        self, key: str, context: grpc.ServicerContext
+    ) -> GetServingUrlRequestResponse:
+        response = self._responses.get(key)
+        if response:
+            return GetServingUrlRequestResponse(location=response)
+        context.set_code(self._error_codes.get(key, grpc.StatusCode.NOT_FOUND))
+        context.set_details(f"No mock response set for {key}")
+        return GetServingUrlRequestResponse()
+
+    # Setters
+    def set_get_offline_serving_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("offline", response, error_code)
+
+    def set_get_distribution_manager_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("distribution", response, error_code)
+
+    def set_get_analytics_engine_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("analytics", response, error_code)
+
+    def set_get_metrics_gateway_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("metrics", response, error_code)
+
+    def set_get_features_operator_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("features", response, error_code)
+
+    def set_get_hosting_gateway_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("hosting", response, error_code)
+
+    def GetOfflineServingUrl(self, request, context):
+        return self._handle("offline", context)
+
+    def GetDistributionManagerUrl(self, request, context):
+        return self._handle("distribution", context)
+
+    def GetAnalyticsEngineUrl(self, request, context):
+        return self._handle("analytics", context)
+
+    def GetMetricsGatewayUrl(self, request, context):
+        return self._handle("metrics", context)
+
+    def GetFeaturesOperatorUrl(self, request, context):
+        return self._handle("features", context)
+
+    def GetHostingGatewayUrl(self, request, context):
+        return self._handle("hosting", context)

frogml_services_mock/services_mock.py

@@ -3,7 +3,7 @@ from typing import Any, Generator, List, Tuple
 
 import grpc
 import pytest
-
+from frogml_core.inner.di_configuration import FrogmlContainer
 from frogml_proto.jfml.model_version.v1.model_version_manager_service_pb2_grpc import (
     add_ModelVersionManagerServiceServicer_to_server,
 )
@@ -113,6 +113,9 @@ from frogml_proto.qwak.secret_service.secret_service_pb2_grpc import (
 from frogml_proto.qwak.self_service.user.v1.user_service_pb2_grpc import (
     add_UserServiceServicer_to_server,
 )
+from frogml_proto.qwak.service_discovery.service_discovery_location_service_pb2_grpc import (
+    add_LocationDiscoveryServiceServicer_to_server,
+)
 from frogml_proto.qwak.vectors.v1.collection.collection_service_pb2_grpc import (
     add_VectorCollectionServiceServicer_to_server,
 )
@@ -122,7 +125,6 @@ from frogml_proto.qwak.vectors.v1.vector_service_pb2_grpc import (
 from frogml_proto.qwak.workspace.workspace_service_pb2_grpc import (
     add_WorkspaceManagementServiceServicer_to_server,
 )
-from frogml_core.inner.di_configuration import FrogmlContainer
 from frogml_services_mock.mocks.alert_manager_service_api import (
     AlertManagerServiceApiMock,
 )
@@ -187,6 +189,9 @@ from frogml_services_mock.mocks.job_registry_service_api import (
 from frogml_services_mock.mocks.kube_captain_service_api import (
     KubeCaptainServiceApiMock,
 )
+from frogml_services_mock.mocks.location_discovery_service_api import (
+    LocationDiscoveryServiceApiMock,
+)
 from frogml_services_mock.mocks.logging_service import LoggingServiceApiMock
 from frogml_services_mock.mocks.model_management_service import (
     ModelsManagementServiceMock,
@@ -249,16 +254,17 @@ def frogml_container():
         feature_store,
         file_versioning,
         instance_template,
+        jfrog_gateway,
         kube_deployment_captain,
+        location_discovery,
         logging_client,
         model_management,
+        model_version_manager,
         project,
         secret_service,
         user_application_instance,
         vector_store,
         workspace_manager,
-        model_version_manager,
-        jfrog_gateway,
     )
     from frogml_core.clients.administration import (
         authentication,
@@ -303,6 +309,7 @@ def frogml_container():
             prompt_manager_client,
             model_version_manager,
             jfrog_gateway,
+            location_discovery,
         ]
     )
 
@@ -527,6 +534,11 @@ def attach_servicers(free_port, server):
                 RepositoryServiceMock,
                 add_RepositoryServiceServicer_to_server,
             ),
+            (
+                "location_discovery_service",
+                LocationDiscoveryServiceApiMock,
+                add_LocationDiscoveryServiceServicer_to_server,
+            ),
             ("port", free_port, None),
         ],
     )

frogml_core/feature_store/offline/_query_engine.py

@@ -1,32 +0,0 @@
-import uuid
-from abc import ABC, abstractmethod
-
-import pandas as pd
-
-
-class BaseQueryEngine(ABC):
-    @abstractmethod
-    def upload_table(self, df: pd.DataFrame):
-        pass
-
-    @abstractmethod
-    def cleanup(self):
-        pass
-
-    @abstractmethod
-    def run_query(self, query: str):
-        pass
-
-    @abstractmethod
-    def read_pandas_from_query(self, query: str, parse_dates=None):
-        pass
-
-    @staticmethod
-    @abstractmethod
-    def get_quotes():
-        pass
-
-    class JoinTableSpec:
-        def __init__(self, join_tables_db_name: str, quotes: str):
-            self.table_name = str(uuid.uuid4())
-            self.join_table_full_path = f"{quotes}{join_tables_db_name}{quotes}.{quotes}{self.table_name}{quotes}"

frogml_core/feature_store/offline/athena/athena_query_engine.py

@@ -1,154 +0,0 @@
-import time
-import uuid
-
-import pandas as pd
-from google.protobuf.duration_pb2 import Duration
-
-from frogml_proto.qwak.ecosystem.v0.ecosystem_runtime_service_pb2 import (
-    GetCloudCredentialsParameters,
-    GetCloudCredentialsRequest,
-    OfflineFeatureStoreClient,
-    PermissionSet,
-)
-from frogml_core.clients.administration.eco_system.client import EcosystemClient
-from frogml_core.exceptions import FrogmlException
-from frogml_core.feature_store.offline._query_engine import BaseQueryEngine
-
-RECONNECT_THRESHOLD_SEC = 300
-
-
-class AthenaQueryEngine(BaseQueryEngine):
-    def __init__(self):
-        eco_client = EcosystemClient()
-        self.bucket, environment_id = self._get_env_details(eco_client)
-
-        self.staging_folder_prefix = (
-            f"{environment_id}/tmp/offline_fs/{str(uuid.uuid4())}"  # nosec B108
-        )
-        self.temp_join_table_base_folder = (
-            f"s3://{self.bucket}/{self.staging_folder_prefix}"
-        )
-
-        self.conn, self.expiration_time = self._init_connection()
-        self.cursor = self.conn.cursor()
-
-        self.join_table_specs = []
-
-        self.join_tables_db_name = f"qwak_temp_data_{environment_id.replace('-', '_')}"
-        self.cursor.execute(f"CREATE DATABASE IF NOT EXISTS {self.join_tables_db_name}")
-
-    @staticmethod
-    def _get_env_details(eco_client):
-        environment_configuration = eco_client.get_environment_configuration()
-
-        return (
-            environment_configuration.configuration.object_storage_bucket,
-            environment_configuration.id,
-        )
-
-    def _init_connection(self):
-        try:
-            # obtain credentials through STS
-            eco_client = EcosystemClient()
-            cloud_credentials_response = eco_client.get_cloud_credentials(
-                request=GetCloudCredentialsRequest(
-                    parameters=GetCloudCredentialsParameters(
-                        duration=Duration(seconds=60 * 60, nanos=0),
-                        permission_set=PermissionSet(
-                            offline_feature_store_client=OfflineFeatureStoreClient()
-                        ),
-                    )
-                )
-            )
-
-            aws_credentials = (
-                cloud_credentials_response.cloud_credentials.aws_temporary_credentials
-            )
-
-            try:
-                from pyathena import connect
-                from pyathena.pandas.cursor import PandasCursor
-            except ImportError:
-                raise FrogmlException(
-                    """
-                    Missing 'pyathena' dependency required for fetching data from the offline store.
-                    Please pip install pyathena
-                """
-                )
-
-            conn = connect(
-                s3_staging_dir=self.temp_join_table_base_folder,
-                aws_access_key_id=aws_credentials.access_key_id,
-                aws_secret_access_key=aws_credentials.secret_access_key,
-                aws_session_token=aws_credentials.session_token,
-                region_name=aws_credentials.region,
-                cursor_class=PandasCursor,
-            )
-
-            return (
-                conn,
-                aws_credentials.expiration_time.seconds,
-            )
-
-        except FrogmlException as e:
-            raise e
-
-        except Exception as e:
-            raise FrogmlException(
-                f"Got an error trying to retrieve credentials to query the offline store "
-                f"in the cloud, error is: {e}"
-            )
-
-    def upload_table(self, df: pd.DataFrame):
-        join_table_spec = super().JoinTableSpec(
-            self.join_tables_db_name, AthenaQueryEngine.get_quotes()
-        )
-        self.join_table_specs.append(join_table_spec)
-
-        from pyathena.pandas.util import to_sql
-
-        to_sql(
-            df,
-            join_table_spec.table_name,
-            self.conn,
-            f"{self.temp_join_table_base_folder}/{join_table_spec.table_name}/",
-            schema=self.join_tables_db_name,
-            index=False,
-            if_exists="replace",
-        )
-
-        return join_table_spec.join_table_full_path
-
-    def run_query(self, query: str):
-        self._check_reconnection()
-        return self.cursor.execute(query).fetchall()
-
-    def read_pandas_from_query(self, query: str, parse_dates=None):
-        self._check_reconnection()
-        return pd.read_sql(
-            query,
-            self.conn,
-            parse_dates=parse_dates,
-        )
-
-    def _check_reconnection(self):
-        if self.expiration_time - time.time() < RECONNECT_THRESHOLD_SEC:
-            self.conn, self.expiration_time = self._init_connection()
-            self.cursor = self.conn.cursor()
-
-    def cleanup(self):
-        self._check_reconnection()
-        for join_table_spec in self.join_table_specs:
-            self.cursor.execute(
-                f"""DROP TABLE {join_table_spec.join_table_full_path.replace('"', '`')}"""
-            )
-
-        self.join_table_specs = []
-
-        s3 = self.conn.session.resource("s3")
-        bucket = s3.Bucket(self.bucket)
-        bucket.objects.filter(Prefix=self.staging_folder_prefix).delete()
-
-    @staticmethod
-    def get_quotes():
-        return '"'