From 3bfa0ea7621165b72b891d2f0350aa38a13c18d0 Mon Sep 17 00:00:00 2001 From: Russell Spitzer Date: Tue, 28 Feb 2017 14:08:52 -0800 Subject: [PATCH] SPARKC-476: SessionProxy Should proxy all Runtime Interfaces Previously Session proxy would only proxy the Session interface. Here we look up all available interfaces and proxy all of them. --- .../datastax/spark/connector/cql/SessionProxy.scala | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/spark-cassandra-connector/src/main/scala/com/datastax/spark/connector/cql/SessionProxy.scala b/spark-cassandra-connector/src/main/scala/com/datastax/spark/connector/cql/SessionProxy.scala index a25a0246a..2c5ef57b8 100644 --- a/spark-cassandra-connector/src/main/scala/com/datastax/spark/connector/cql/SessionProxy.scala +++ b/spark-cassandra-connector/src/main/scala/com/datastax/spark/connector/cql/SessionProxy.scala @@ -4,6 +4,9 @@ import java.lang.reflect.{InvocationHandler, InvocationTargetException, Method, import com.datastax.driver.core.{RegularStatement, Session, SimpleStatement} +import org.apache.commons.lang3.ClassUtils + + /** Wraps a `Session` and intercepts: * - `close` method to invoke `afterClose` handler * - `prepare` methods to cache `PreparedStatement` objects. */ @@ -56,10 +59,14 @@ object SessionProxy { /** Creates a new `SessionProxy` delegating to the given `Session`. * Additionally registers a callback on `Session#close` method. + * * @param afterClose code to be invoked after the session has been closed */ - def wrapWithCloseAction(session: Session)(afterClose: Session => Any): Session = + def wrapWithCloseAction(session: Session)(afterClose: Session => Any): Session = { + val listInterfaces = ClassUtils.getAllInterfaces(session.getClass) + val availableInterfaces = listInterfaces.toArray[Class[_]](new Array[Class[_]](listInterfaces.size)) Proxy.newProxyInstance( session.getClass.getClassLoader, - Array(classOf[Session]), + availableInterfaces, new SessionProxy(session, afterClose)).asInstanceOf[Session] + } } \ No newline at end of file