initial commit of pysplicecontext.

Author: Amrit Baveja

setup.py

@@ -0,0 +1,18 @@
+from setuptools import setup, find_packages
+dependencies = [
+"atomicwrites==1.1.5",
+"attrs==18.1.0",
+"more-itertools==4.2.0",
+"pluggy==0.6.0",
+"py==1.5.3",
+"py4j==0.10.7",
+"pyspark==2.3.1",
+"pytest==3.6.1",
+"six==1.11.0"
+]
+setup(
+    name="splicemachine",
+    version="0.2.2",
+    install_requires=dependencies,
+    packages=['splicemachine'],
+)

splicemachine/__init__.py

@@ -0,0 +1,104 @@
+from __future__ import print_function
+
+from pyspark.sql import DataFrame
+from py4j.java_gateway import java_import
+
+
+class PySpliceContext:
+    """
+    This class implements a SpliceMachineContext object (similar to the SparkContext object)
+    """
+
+    def __init__(self, JDBC_URL, sparkSQLContext, _unitTesting=False):
+        """
+        :param JDBC_URL: (string) The JDBC URL Connection String for your Splice Machine Cluster
+        :param sparkSQLContext: (sparkContext) A SparkContext Object for executing Spark Queries
+        """
+        self.jdbcurl = JDBC_URL
+        self._unitTesting = _unitTesting
+
+        if not _unitTesting:  # Private Internal Argument to Override Using JVM
+            self.sparkSQLContext = sparkSQLContext
+            self.jvm = self.sparkSQLContext._sc._jvm
+            java_import(self.jvm, "com.splicemachine.spark.splicemachine.*")
+            java_import(self.jvm, "org.apache.spark.sql.execution.datasources.jdbc.{JDBCOptions, JdbcUtils}")
+            java_import(self.jvm, "scala.collection.JavaConverters._")
+            self.context = self.jvm.com.splicemachine.spark.splicemachine.SplicemachineContext(self.jdbcurl)
+
+        else:
+            from .utils import FakeJContext
+            self.sparkSQLContext = sparkSQLContext
+            self.jvm = ''
+            self.context = FakeJContext(self.jdbcurl)
+
+    def getConnection(self):
+        """
+        Return a connection to the database
+        """
+        return self.context.getConnection()
+
+    def tableExists(self, schemaTableName):
+        """
+        Check whether or not a table exists
+
+        :param schemaTableName: (string) Table Name
+        """
+        return self.context.tableExists(schemaTableName)
+
+    def dropTable(self, schemaTableName):  # works
+        """
+        Drop a specified table.
+
+        :param schemaTableName (optional): (string) schemaName.tableName
+        """
+        return self.context.dropTable(schemaTableName)
+
+    def df(self, sql):
+        """
+        Return a Spark Dataframe from the results of a Splice Machine SQL Query
+
+        :param sql: (string) SQL Query (eg. SELECT * FROM table1 WHERE column2 > 3)
+        :return: A Spark DataFrame containing the results
+        """
+        if self._unitTesting:
+            return self.context.df(sql)
+        return DataFrame(self.context.df(sql), self.sparkSQLContext)
+
+    def insert(self, dataFrame, schemaTableName):
+        """
+        Insert a RDD into a table (schema.table).  The schema is required since RDD's do not have schema.
+
+        :param dataFrame: (RDD) The dataFrame you would like to insert
+        :param schemaTableName: (string) The table in which you would like to insert the RDD
+        """
+        return self.context.insert(dataFrame._jdf, schemaTableName)
+
+    def delete(self, dataFrame, schemaTableName):
+        """
+        Delete records in a dataframe based on joining by primary keys from the data frame.
+        Be careful with column naming and case sensitivity.
+
+        :param dataFrame: (RDD) The dataFrame you would like to delete
+        :param schemaTableName: (string) Splice Machine Table
+        """
+        return self.context.delete(dataFrame._jdf, schemaTableName)
+
+    def update(self, dataFrame, schemaTableName):
+        """
+        Update data from a dataframe for a specified schemaTableName (schema.table).
+        The keys are required for the update and any other columns provided will be updated in the rows.
+
+        :param dataFrame: (RDD) The dataFrame you would like to update
+        :param schemaTableName: (string) Splice Machine Table
+        :return:
+        """
+        return self.context.update(dataFrame._jdf, schemaTableName)
+
+    def getSchema(self, schemaTableName):
+        """
+        Return the schema via JDBC.
+
+        :param schemaTableName: (string) Table name
+        """
+        return self.context.getSchema(schemaTableName)
+

splicemachine/context.py

@@ -0,0 +1,61 @@
+import pyspark
+from .context import PySpliceContext
+from .utils import fakeDf
+
+conf = pyspark.SparkConf().setAppName('Unit Test Python SpliceContext API')
+sc = pyspark.SparkContext(conf=conf)
+sqlContext = pyspark.sql.SQLContext(sc)
+spliceContext = PySpliceContext('', sqlContext, _unitTesting=True)
+
+
+class TestContext:
+    def test_getConnection(self):
+        out = spliceContext.getConnection()
+        assert out['event'] == 'get connection'
+
+    def test_tableExists(self):
+        out = spliceContext.tableExists('schema1.table1')
+        assert out['event'] == 'table exists'
+        assert out['schemaTableName'] == 'schema1.table1'
+        assert out['schemaName'] == 'schema1'
+        assert out['tableName'] == 'table1'
+
+    def test_dropTable(self):
+        out = spliceContext.dropTable('schema2.table3')
+        assert out['event'] == 'drop table'
+        assert out['schemaTableName'] == 'schema2.table3'
+        assert out['schemaName'] == 'schema2'
+        assert out['tableName'] == 'table3'
+
+    def test_df(self):
+        out = spliceContext.df('SELECT * FROM table1')
+        assert out['sql'] == 'SELECT * FROM table1'
+        assert out['event'] == 'df'
+
+    def test_insert(self):
+        out = spliceContext.insert(fakeDf(), 'schema.table94')
+        assert out['tableName'] == 'table94'
+        assert out['schemaTableName'] == 'schema.table94'
+        assert out['schemaName'] == 'schema'
+        assert out['event'] == 'insert'
+
+    def test_delete(self):
+        out = spliceContext.delete(fakeDf(), 'schema4.table4')
+        assert out['tableName'] == 'table4'
+        assert out['schemaTableName'] == 'schema4.table4'
+        assert out['schemaName'] == 'schema4'
+        assert out['event'] == 'delete'
+
+    def test_update(self):
+        out = spliceContext.update(fakeDf(), 'schema0.table390')
+        assert out['tableName'] == 'table390'
+        assert out['schemaTableName'] == 'schema0.table390'
+        assert out['schemaName'] == 'schema0'
+        assert out['event'] == 'update'
+
+    def test_getSchema(self):
+        out = spliceContext.getSchema('schema41.table12')
+        assert out['event'] == 'getSchema'
+        assert out['schemaTableName'] == 'schema41.table12'
+        assert out['schemaName'] == 'schema41'
+        assert out['tableName'] == 'table12'

splicemachine/test_context.py

@@ -0,0 +1,56 @@
+def _generateOperationsTable(**kwargs):
+    """
+    Usage: _generateOperationsTable(event='get connection')
+    --> {'event': 'get connection'}
+    """
+    return kwargs
+
+
+class fakeDf(object):
+    def __init__(self):
+        self._jdf = ''
+
+
+class FakeJContext(object):
+    """
+    This class is a Fake Representation of the Scala SpliceMachineContext API for unit testing
+    """
+
+    def __init__(self, JDBC_URL):
+        print("Class Initialized")
+
+    def getConnection(self):
+        return _generateOperationsTable(event='get connection')
+
+    def tableExists(self, schemaTableName):
+        schemaName, tableName = schemaTableName.split('.')
+        return _generateOperationsTable(event='table exists', schemaTableName=schemaTableName, schemaName=schemaName,
+                                        tableName=tableName)
+
+    def dropTable(self, schemaTableName):
+        schemaName, tableName = schemaTableName.split('.')
+        return _generateOperationsTable(event='drop table', schemaTableName=schemaTableName, schemaName=schemaName,
+                                        tableName=tableName)
+
+    def df(self, sql):
+        return _generateOperationsTable(event='df', sql=sql)
+
+    def insert(self, dataFrame, schemaTableName):
+        schemaName, tableName = schemaTableName.split('.')
+        return _generateOperationsTable(event='insert', schemaTableName=schemaTableName, schemaName=schemaName,
+                                        tableName=tableName, dataFrame=dataFrame)
+
+    def delete(self, dataFrame, schemaTableName):
+        schemaName, tableName = schemaTableName.split('.')
+        return _generateOperationsTable(event='delete', schemaTableName=schemaTableName, schemaName=schemaName,
+                                        tableName=tableName, dataFrame=dataFrame)
+
+    def update(self, dataFrame, schemaTableName):
+        schemaName, tableName = schemaTableName.split('.')
+        return _generateOperationsTable(event='update', schemaTableName=schemaTableName, schemaName=schemaName,
+                                        tableName=tableName, dataFrame=dataFrame)
+
+    def getSchema(self, schemaTableName):
+        schemaName, tableName = schemaTableName.split('.')
+        return _generateOperationsTable(event='getSchema', schemaTableName=schemaTableName, schemaName=schemaName,
+                                        tableName=tableName)