def test_create_session(mocker, server): mocker.patch.object(Session, 'from_json') client = LivyClient(server) session = client.create_session(SessionKind.PYSPARK) assert session == Session.from_json.return_value Session.from_json.assert_called_once_with(MOCK_SESSION_JSON)
def test_verify(requests_mock, mocker, verify): requests_mock.get("http://example.com/sessions", json={"sessions": []}) mocker.patch.object(Session, "from_json") client = LivyClient("http://example.com", verify=verify) client.list_sessions() [request] = requests_mock.request_history assert request.verify is verify
def test_list_sessions(mocker, server): mocker.patch.object(Session, 'from_json') client = LivyClient(server) sessions = client.list_sessions() assert sessions == [Session.from_json.return_value] Session.from_json.assert_called_once_with(MOCK_SESSION_JSON)
def test_get_session(mocker, server, auth): mocker.patch.object(Session, "from_json") client = LivyClient(server, auth) session = client.get_session(MOCK_SESSION_ID) assert session == Session.from_json.return_value Session.from_json.assert_called_once_with(MOCK_SESSION_JSON)
def test_list_statements(mocker, server, auth): mocker.patch.object(Statement, "from_json") client = LivyClient(server, auth) statements = client.list_statements(MOCK_SESSION_ID) assert statements == [Statement.from_json.return_value] Statement.from_json.assert_called_once_with(MOCK_SESSION_ID, MOCK_STATEMENT_JSON)
def test_get_statement(mocker, server): mocker.patch.object(Statement, 'from_json') client = LivyClient(server) statement = client.get_statement(MOCK_SESSION_ID, MOCK_STATEMENT_ID) assert statement == Statement.from_json.return_value Statement.from_json.assert_called_once_with(MOCK_SESSION_ID, MOCK_STATEMENT_JSON)
def test_list_sessions(requests_mock, mocker): requests_mock.get("http://example.com/sessions", json={"sessions": [MOCK_SESSION_JSON]}) mocker.patch.object(Session, "from_json") client = LivyClient("http://example.com") sessions = client.list_sessions() assert sessions == [Session.from_json.return_value] Session.from_json.assert_called_once_with(MOCK_SESSION_JSON)
def test_delete_session(requests_mock): requests_mock.delete( f"http://example.com/sessions/{MOCK_SESSION_ID}", json={"msg": "deleted"}, ) client = LivyClient("http://example.com") client.delete_session(MOCK_SESSION_ID) assert requests_mock.called
def __init__( self, url: str, batch_id: int, auth: Auth = None, verify: Verify = True, requests_session: requests.Session = None, ) -> None: self.client = LivyClient(url, auth, verify, requests_session) self.batch_id = batch_id
def test_create_statement(mocker, server, auth): mocker.patch.object(Statement, "from_json") client = LivyClient(server, auth) statement = client.create_statement(MOCK_SESSION_ID, MOCK_CODE, StatementKind.PYSPARK) assert statement == Statement.from_json.return_value Statement.from_json.assert_called_once_with(MOCK_SESSION_ID, MOCK_STATEMENT_JSON)
def test_get_batch_log(requests_mock, mocker): requests_mock.get( f"http://example.com/batches/{MOCK_BATCH_ID}/log", json=MOCK_BATCH_LOG_JSON, ) mocker.patch.object(BatchLog, "from_json") client = LivyClient("http://example.com") batch = client.get_batch_log(MOCK_BATCH_ID) assert batch == BatchLog.from_json.return_value BatchLog.from_json.assert_called_once_with(MOCK_BATCH_LOG_JSON)
def __init__(self, url: str, kind: SessionKind = SessionKind.PYSPARK, spark_conf: Dict[str, Any] = None, echo: bool = True, check: bool = True) -> None: self.client = LivyClient(url) self.kind = kind self.echo = echo self.check = check self.session_id: Optional[int] = None self.spark_conf = spark_conf
def test_get_session(requests_mock, mocker): requests_mock.get( f"http://example.com/sessions/{MOCK_SESSION_ID}", json=MOCK_SESSION_JSON, ) mocker.patch.object(Session, "from_json") client = LivyClient("http://example.com") session = client.get_session(MOCK_SESSION_ID) assert session == Session.from_json.return_value Session.from_json.assert_called_once_with(MOCK_SESSION_JSON)
def test_create_session(mocker, server, auth): mocker.patch.object(Session, "from_json") client = LivyClient(server, auth) session = client.create_session( SessionKind.PYSPARK, proxy_user=MOCK_PROXY_USER, spark_conf=MOCK_SPARK_CONF, ) assert session == Session.from_json.return_value Session.from_json.assert_called_once_with(MOCK_SESSION_JSON)
def test_auth(requests_mock, mocker): requests_mock.get("http://example.com/sessions", json={"sessions": []}) mocker.patch.object(Session, "from_json") def dummy_auth(request): request.headers["Authorization"] = "dummy-token" return request client = LivyClient("http://example.com", auth=dummy_auth) client.list_sessions() [request] = requests_mock.request_history assert request.headers["Authorization"] == "dummy-token"
def test_list_statements(requests_mock, mocker): requests_mock.get( f"http://example.com/sessions/{MOCK_SESSION_ID}/statements", json={"statements": [MOCK_STATEMENT_JSON]}, ) mocker.patch.object(Statement, "from_json") client = LivyClient("http://example.com") statements = client.list_statements(MOCK_SESSION_ID) assert statements == [Statement.from_json.return_value] Statement.from_json.assert_called_once_with(MOCK_SESSION_ID, MOCK_STATEMENT_JSON)
def test_get_statement(requests_mock, mocker): requests_mock.get( f"http://example.com/sessions/{MOCK_SESSION_ID}" + f"/statements/{MOCK_STATEMENT_ID}", json=MOCK_STATEMENT_JSON, ) mocker.patch.object(Statement, "from_json") client = LivyClient("http://example.com") statement = client.get_statement(MOCK_SESSION_ID, MOCK_STATEMENT_ID) assert statement == Statement.from_json.return_value Statement.from_json.assert_called_once_with(MOCK_SESSION_ID, MOCK_STATEMENT_JSON)
def __init__( self, url: str, session_id: int, auth: Auth = None, verify: Verify = True, requests_session: requests.Session = None, kind: SessionKind = SessionKind.PYSPARK, echo: bool = True, check: bool = True, ) -> None: self.client = LivyClient(url, auth, verify, requests_session) self.session_id = session_id self.kind = kind self.echo = echo self.check = check
def test_create_batch(requests_mock, mocker): requests_mock.get("http://example.com/version", json={"version": "0.5.0-incubating"}) requests_mock.post("http://example.com/batches", json=MOCK_BATCH_JSON) mocker.patch.object(Batch, "from_json") client = LivyClient("http://example.com") batch = client.create_batch( file=MOCK_BATCH_FILE, class_name=MOCK_BATCH_CLASSNAME, args=MOCK_BATCH_ARGS, proxy_user=MOCK_PROXY_USER, jars=MOCK_JARS, py_files=MOCK_PY_FILES, files=MOCK_FILES, driver_memory=MOCK_DRIVER_MEMORY, driver_cores=MOCK_DRIVER_CORES, executor_memory=MOCK_EXECUTOR_MEMORY, executor_cores=MOCK_EXECUTOR_CORES, num_executors=MOCK_NUM_EXECUTORS, archives=MOCK_ARCHIVES, queue=MOCK_QUEUE, name=MOCK_NAME, spark_conf=MOCK_SPARK_CONF, ) assert batch == Batch.from_json.return_value Batch.from_json.assert_called_once_with(MOCK_BATCH_JSON) assert requests_mock.last_request.json() == { "file": MOCK_BATCH_FILE, "proxyUser": MOCK_PROXY_USER, "className": MOCK_BATCH_CLASSNAME, "args": MOCK_BATCH_ARGS, "jars": MOCK_JARS, "pyFiles": MOCK_PY_FILES, "files": MOCK_FILES, "driverMemory": MOCK_DRIVER_MEMORY, "driverCores": MOCK_DRIVER_CORES, "executorMemory": MOCK_EXECUTOR_MEMORY, "executorCores": MOCK_EXECUTOR_CORES, "numExecutors": MOCK_NUM_EXECUTORS, "archives": MOCK_ARCHIVES, "queue": MOCK_QUEUE, "name": MOCK_NAME, "conf": MOCK_SPARK_CONF, }
def test_custom_requests_session(mocker): mocker.patch.object(Session, "from_json") mock_requests_session = mocker.Mock() mock_response = mocker.Mock() mock_response.json.return_value = {"sessions": []} mock_requests_session.request.return_value = mock_response client = LivyClient("http://example.com", requests_session=mock_requests_session) client.list_sessions() mock_requests_session.request.assert_called_once() # Check that a custom session does not get closed client.close() mock_requests_session.close.assert_not_called()
def __init__( self, url: str, auth: Auth = None, kind: SessionKind = SessionKind.PYSPARK, proxy_user: str = None, spark_conf: Dict[str, Any] = None, echo: bool = True, check: bool = True, ) -> None: self.client = LivyClient(url, auth) self.kind = kind self.proxy_user = proxy_user self.spark_conf = spark_conf self.echo = echo self.check = check self.session_id: Optional[int] = None
def test_create_session(requests_mock, mocker): requests_mock.get("http://example.com/version", json={"version": "0.5.0-incubating"}) requests_mock.post("http://example.com/sessions", json=MOCK_SESSION_JSON) mocker.patch.object(Session, "from_json") client = LivyClient("http://example.com") session = client.create_session( SessionKind.PYSPARK, proxy_user=MOCK_PROXY_USER, jars=MOCK_JARS, py_files=MOCK_PY_FILES, files=MOCK_FILES, driver_memory=MOCK_DRIVER_MEMORY, driver_cores=MOCK_DRIVER_CORES, executor_memory=MOCK_EXECUTOR_MEMORY, executor_cores=MOCK_EXECUTOR_CORES, num_executors=MOCK_NUM_EXECUTORS, archives=MOCK_ARCHIVES, queue=MOCK_QUEUE, name=MOCK_NAME, spark_conf=MOCK_SPARK_CONF, heartbeat_timeout=MOCK_HEARTBEAT_TIMEOUT, ) assert session == Session.from_json.return_value Session.from_json.assert_called_once_with(MOCK_SESSION_JSON) assert requests_mock.last_request.json() == { "kind": "pyspark", "proxyUser": MOCK_PROXY_USER, "conf": MOCK_SPARK_CONF, "heartbeatTimeoutInSecond": MOCK_HEARTBEAT_TIMEOUT, "jars": MOCK_JARS, "pyFiles": MOCK_PY_FILES, "files": MOCK_FILES, "driverMemory": MOCK_DRIVER_MEMORY, "driverCores": MOCK_DRIVER_CORES, "executorMemory": MOCK_EXECUTOR_MEMORY, "executorCores": MOCK_EXECUTOR_CORES, "numExecutors": MOCK_NUM_EXECUTORS, "archives": MOCK_ARCHIVES, "queue": MOCK_QUEUE, "name": MOCK_NAME, }
def test_create_statement(requests_mock, mocker): requests_mock.get("http://example.com/version", json={"version": "0.5.0-incubating"}) requests_mock.post( f"http://example.com/sessions/{MOCK_SESSION_ID}/statements", json=MOCK_STATEMENT_JSON, ) mocker.patch.object(Statement, "from_json") client = LivyClient("http://example.com") statement = client.create_statement(MOCK_SESSION_ID, MOCK_CODE, StatementKind.PYSPARK) assert statement == Statement.from_json.return_value Statement.from_json.assert_called_once_with(MOCK_SESSION_ID, MOCK_STATEMENT_JSON) assert requests_mock.last_request.json() == { "code": MOCK_CODE, "kind": "pyspark", }
def __init__( self, url: str, auth: Auth = None, verify: Verify = True, kind: SessionKind = SessionKind.PYSPARK, proxy_user: str = None, jars: List[str] = None, py_files: List[str] = None, files: List[str] = None, driver_memory: str = None, driver_cores: int = None, executor_memory: str = None, executor_cores: int = None, num_executors: int = None, archives: List[str] = None, queue: str = None, name: str = None, spark_conf: Dict[str, Any] = None, echo: bool = True, check: bool = True, headers: dict = {}, ) -> None: self.client = LivyClient(url, auth, verify=verify, headers=headers) self.kind = kind self.proxy_user = proxy_user self.jars = jars self.py_files = py_files self.files = files self.driver_memory = driver_memory self.driver_cores = driver_cores self.executor_memory = executor_memory self.executor_cores = executor_cores self.num_executors = num_executors self.archives = archives self.queue = queue self.name = name self.spark_conf = spark_conf self.echo = echo self.check = check self.session_id: Optional[int] = None
def __init__( self, url: str, file: str, auth: Auth = None, class_name: str = None, args: List[str] = None, proxy_user: str = None, jars: List[str] = None, py_files: List[str] = None, files: List[str] = None, driver_memory: str = None, driver_cores: int = None, executor_memory: str = None, executor_cores: int = None, num_executors: int = None, archives: List[str] = None, queue: str = None, name: str = None, spark_conf: Dict[str, Any] = None, ) -> None: self.client = LivyClient(url, auth) self.file = file self.class_name = class_name self.args = args self.proxy_user = proxy_user self.jars = jars self.py_files = py_files self.files = files self.driver_memory = driver_memory self.driver_cores = driver_cores self.executor_memory = executor_memory self.executor_cores = executor_cores self.num_executors = num_executors self.archives = archives self.queue = queue self.name = name self.spark_conf = spark_conf self.batch_id: Optional[int] = None
def create( cls, url: str, auth: Auth = None, verify: Verify = True, requests_session: requests.Session = None, kind: SessionKind = SessionKind.PYSPARK, proxy_user: str = None, jars: List[str] = None, py_files: List[str] = None, files: List[str] = None, driver_memory: str = None, driver_cores: int = None, executor_memory: str = None, executor_cores: int = None, num_executors: int = None, archives: List[str] = None, queue: str = None, name: str = None, spark_conf: Dict[str, Any] = None, heartbeat_timeout: int = None, echo: bool = True, check: bool = True, ) -> "LivySession": """Create a new Livy session. The py_files, files, jars and archives arguments are lists of URLs, e.g. ["s3://bucket/object", "hdfs://path/to/file", ...] and must be reachable by the Spark driver process. If the provided URL has no scheme, it's considered to be relative to the default file system configured in the Livy server. URLs in the py_files argument are copied to a temporary staging area and inserted into Python's sys.path ahead of the standard library paths. This allows you to import .py, .zip and .egg files in Python. URLs for jars, py_files, files and archives arguments are all copied to the same working directory on the Spark cluster. The driver_memory and executor_memory arguments have the same format as JVM memory strings with a size unit suffix ("k", "m", "g" or "t") (e.g. 512m, 2g). See https://spark.apache.org/docs/latest/configuration.html for more information on Spark configuration properties. :param url: The URL of the Livy server. :param auth: A requests-compatible auth object to use when making requests. :param verify: Either a boolean, in which case it controls whether we verify the server’s TLS certificate, or a string, in which case it must be a path to a CA bundle to use. Defaults to ``True``. :param requests_session: A specific ``requests.Session`` to use, allowing advanced customisation. The caller is responsible for closing the session. :param kind: The kind of session to create. :param proxy_user: User to impersonate when starting the session. :param jars: URLs of jars to be used in this session. :param py_files: URLs of Python files to be used in this session. :param files: URLs of files to be used in this session. :param driver_memory: Amount of memory to use for the driver process (e.g. '512m'). :param driver_cores: Number of cores to use for the driver process. :param executor_memory: Amount of memory to use per executor process (e.g. '512m'). :param executor_cores: Number of cores to use for each executor. :param num_executors: Number of executors to launch for this session. :param archives: URLs of archives to be used in this session. :param queue: The name of the YARN queue to which submitted. :param name: The name of this session. :param spark_conf: Spark configuration properties. :param heartbeat_timeout: Optional Timeout in seconds to which session be automatically orphaned if no heartbeat is received. :param echo: Whether to echo output printed in the remote session. Defaults to ``True``. :param check: Whether to raise an exception when a statement in the remote session fails. Defaults to ``True``. """ client = LivyClient(url, auth, verify, requests_session) session = client.create_session( kind, proxy_user, jars, py_files, files, driver_memory, driver_cores, executor_memory, executor_cores, num_executors, archives, queue, name, spark_conf, heartbeat_timeout, ) client.close() return cls( url, session.session_id, auth, verify, requests_session, kind, echo, check, )
def test_delete_session(mocker, server): client = LivyClient(server) client.delete_session(MOCK_SESSION_ID)
from livy.client import LivyClient from livy.models import SessionKind livy_endpoint = 'http://localhost:8998' jars = ['s3://my-bucket/Person.jar'] client = LivyClient(livy_endpoint) session = client.create_session(SessionKind.PYSPARK, jars=jars)