PyStellarDB 0.13.5__py2.py3-none-any.whl → 1.0__py2.py3-none-any.whl

{PyStellarDB-0.13.5.dist-info → PyStellarDB-1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: PyStellarDB
-Version: 0.13.5
+Version: 1.0
 Summary: Python interface to StellarDB
 Home-page: https://github.com/WarpCloud/PyStellarDB
 Author: Zhiping Wang
@@ -10,25 +10,26 @@ Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Topic :: Database :: Front-Ends
-Requires-Python: >=2.7
+Requires-Python: >=3.6
 License-File: LICENSE
 Requires-Dist: future
 Requires-Dist: python-dateutil
-Requires-Dist: pyhive
-Requires-Dist: sasl
-Requires-Dist: thrift
-Requires-Dist: thrift-sasl >=0.3.0
-Provides-Extra: hive
-Requires-Dist: sasl >=0.2.1 ; extra == 'hive'
-Requires-Dist: thrift >=0.10.0 ; extra == 'hive'
-Provides-Extra: kerberos
-Requires-Dist: requests-kerberos >=0.12.0 ; extra == 'kerberos'
-Provides-Extra: presto
-Requires-Dist: requests >=1.0.0 ; extra == 'presto'
+Requires-Dist: pyhive[hive_pure_sasl]
+Requires-Dist: thrift>=0.10.0
 Provides-Extra: pyspark
-Requires-Dist: pyspark >=2.4.0 ; extra == 'pyspark'
-Provides-Extra: sqlalchemy
-Requires-Dist: sqlalchemy >=1.3.0 ; extra == 'sqlalchemy'
+Requires-Dist: pyspark>=2.4.0; extra == "pyspark"
+Provides-Extra: kerberos
+Requires-Dist: kerberos>=1.3.0; extra == "kerberos"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: home-page
+Dynamic: license
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 PyStellarDB
 ===========
@@ -60,6 +61,7 @@ PLAIN Mode (No security is configured)
 ---------------------------------------
 .. code-block:: python
 
+    """version < 1.0 """
     from pystellardb import stellar_hive
 
     conn = stellar_hive.StellarConnection(host="localhost", port=10000, graph_name='pokemon')
@@ -70,11 +72,19 @@ PLAIN Mode (No security is configured)
 
     print cur.fetchall()
 
+    """version >= 1.0 """
+    from pystellardb import Connection, Graph
+    conn = stellar_hive.StellarConnection(host="localhost", port=10000)
+    graph = Graph('pokemon', conn)
+    query_result = graph.execute('match p = (a)-[f]->(b) return a,f,b limit 1')
+    for row in query_result:
+        print(row[0].toJSON(), row[1].toJSON(), row[2].toJSON())
 
 LDAP Mode
 ---------
 .. code-block:: python
 
+    """version < 1.0 """
     from pystellardb import stellar_hive
 
     conn = stellar_hive.StellarConnection(host="localhost", port=10000, username='hive', password='123456', auth='LDAP', graph_name='pokemon')
@@ -85,6 +95,14 @@ LDAP Mode
 
     print cur.fetchall()
 
+    """version >= 1.0 """
+    from pystellardb import Connection, Graph
+    conn = stellar_hive.StellarConnection(host="localhost", port=10000, username='hive', password='123456', auth='LDAP')
+    graph = Graph('pokemon', conn)
+    query_result = graph.execute('match p = (a)-[f]->(b) return a,f,b limit 1')
+    for row in query_result:
+        print(row[0].toJSON(), row[1].toJSON(), row[2].toJSON())
+
 
 Kerberos Mode
 -------------
@@ -96,6 +114,12 @@ Kerberos Mode
     # In Linux: kinit -kt FILE_PATH_OF_KEYTABL PRINCIPAL_NAME
     # In Mac: kinit -t FILE_PATH_OF_KEYTABL -f PRINCIPAL_NAME
 
+    # Run with Kerberos path environment variables
+    # ENV KRB5_CONFIG=/etc/krb5.conf
+    # ENV KRB5_CLIENT_KTNAME=/etc/krb5.keytab
+    # ENV KRB5_KTNAME=/etc/krb5.keytab
+
+    """version < 1.0 """
     from pystellardb import stellar_hive
 
     conn = stellar_hive.StellarConnection(host="localhost", port=10000, kerberos_service_name='hive', auth='KERBEROS', graph_name='pokemon')
@@ -106,6 +130,13 @@ Kerberos Mode
 
     print cur.fetchall()
 
+    """version >= 1.0 """
+    from pystellardb import Connection, Graph
+    conn = stellar_hive.StellarConnection(host="localhost", port=10000, kerberos_service_name='hive', auth='KERBEROS')
+    graph = Graph('pokemon', conn)
+    query_result = graph.execute('match p = (a)-[f]->(b) return a,f,b limit 1')
+    for row in query_result:
+        print(row[0].toJSON(), row[1].toJSON(), row[2].toJSON())
 
 Execute Hive Query
 ------------------
@@ -123,6 +154,7 @@ Execute Graph Query and change to a PySpark RDD object
 ------------------------------------------------------
 .. code-block:: python
 
+    """version < 1.0 """
     from pyspark import SparkContext
     from pystellardb import stellar_hive
     
@@ -140,6 +172,22 @@ Execute Graph Query and change to a PySpark RDD object
 
     rdd.map(lambda x: (x[0].toJSON(), x[1].toJSON(), x[2].toJSON())).foreach(f)
 
+    """version >= 1.0 """
+    from pyspark import SparkContext
+    from pystellardb import Connection, Graph
+    
+    sc = SparkContext("local", "Demo App")
+
+
+    conn = stellar_hive.StellarConnection(host="localhost", port=10000)
+    graph = Graph('pokemon', conn)
+    query_result = graph.execute('match p = (a)-[f]->(b) return a,f,b limit 1')
+    rdd = query_result.toRDD(sc)
+
+    def f(x): print(x)
+
+    rdd.map(lambda x: (x[0].toJSON(), x[1].toJSON(), x[2].toJSON())).foreach(f)
+
     # Every line of this query is in format of Tuple(VertexObject, EdgeObject, VertexObject)
     # Vertex and Edge object has a function of toJSON() which can print the object in JSON format
 
@@ -171,24 +219,24 @@ Dependencies
 Required:
 ------------
 
-- Python 2.7+ / Python 3
+- Python 3.6+
 
-System SASL
-------------
+System SASL(Depricated since 1.0):
+----------------------------------
 
 Ubuntu:
 
 .. code-block:: bash
 
     apt-get install libsasl2-dev libsasl2-2 libsasl2-modules-gssapi-mit
-    apt-get install python-dev gcc              #Update python and gcc if needed
+    apt-get install python3-dev gcc              #Update python and gcc if needed
 
 RHEL/CentOS:
 
 .. code-block:: bash
 
     yum install cyrus-sasl-md5 cyrus-sasl-plain cyrus-sasl-gssapi cyrus-sasl-devel
-    yum install gcc-c++ python-devel.x86_64     #Update python and gcc if needed
+    yum install gcc-c++ python3-devel.x86_64     #Update python and gcc if needed
 
     # if pip3 install fails with a message like 'Can't connect to HTTPS URL because the SSL module is not available'
     # you may need to update ssl & reinstall python

PyStellarDB-1.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+pystellardb/__init__.py,sha256=LHYhg_z_zEX3YLetkbxcZrlJDjQYN5jS9W7RlUck7Wg,285
+pystellardb/_version.py,sha256=bsy7XZ8ffg4i9t7-lI7_arjuL7Y5XBvn7aTm4u1Dirw,495
+pystellardb/graph_types.py,sha256=XbtPebhaRV7OimW-GxFJGSOcp-OnzWrlIe9M7Xv3sbU,14032
+pystellardb/sasl_compat.py,sha256=HLS5PC2T4-E9FCknFZDVGSRpGdrOnKgrrHeJNMtGNPU,1322
+pystellardb/stellar_hive.py,sha256=-6sAQtWbAPhMFUetBAAjkrIqImCYs8hpUvDmCTfdk7I,14964
+pystellardb/stellar_rdd.py,sha256=TYwsWYeCxfOliGq1kV3ArNXdye55cKWZF7s9M9nDdt4,1324
+pystellardb/v2/__init__.py,sha256=12jK5F3bAmNnvxUpoUhgSY4laBfiNiU2GQVM_PkZ8hA,142
+pystellardb/v2/connection.py,sha256=Ma1AXo70QNuoaeJkL-D5rAQCV79knfWFFmpi3Ej2aCE,10762
+pystellardb/v2/database.py,sha256=iI0XF5sgyPiRXeXDJOJsfYrC5K46671k3liWotL-7HA,3147
+PyStellarDB-1.0.dist-info/LICENSE,sha256=1qDFxrywejs7xNBfOr6T-7lOuqDgSNIES77kTYege3w,560
+PyStellarDB-1.0.dist-info/METADATA,sha256=1MDEIEp7fmQKO6vLN3x071v2lB_uH2rIrzxoWXkoAWA,11221
+PyStellarDB-1.0.dist-info/WHEEL,sha256=9Hm2OB-j1QcCUq9Jguht7ayGIIZBRTdOXD1qg9cCgPM,109
+PyStellarDB-1.0.dist-info/top_level.txt,sha256=DRk-SeGVCdVAzv2CwFmdu75Yo7DgjUA3Hpu-9l8qPuU,12
+PyStellarDB-1.0.dist-info/RECORD,,

{PyStellarDB-0.13.5.dist-info → PyStellarDB-1.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (70.1.1)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py2-none-any
 Tag: py3-none-any

pystellardb/__init__.py

@@ -1,4 +1,17 @@
 
 from ._version import get_versions
+from .v2 import Connection, Graph, CypherResult
+
+
 __version__ = get_versions()['version']
 del get_versions
+
+__all__ = [
+    '__version__',
+    'stellar_hive',
+    'stellar_rdd',
+    'graph_types',
+    'Connection',
+    'Graph',
+    'CypherResult',
+]

pystellardb/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-17T19:30:54+0800",
+ "date": "2025-05-06T17:47:56+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "edb05dd3d51f0a5af260e6cd442163058a81db20",
- "version": "0.13.5"
+ "full-revisionid": "76d9b0b752e5e453e3046d5864750c4e3ad85b36",
+ "version": "1.0"
 }
 '''  # END VERSION_JSON
 

pystellardb/sasl_compat.py

@@ -0,0 +1,49 @@
+from __future__ import absolute_import
+
+from puresasl.client import SASLClient, SASLError
+from contextlib import contextmanager
+
+@contextmanager
+def error_catcher(self, Exc = Exception):
+    try:
+        self.error = None
+        yield
+    except Exc as e:
+        self.error = str(e)
+
+
+class PureSASLClient(SASLClient):
+    def __init__(self, *args, **kwargs):
+        self.error = None
+        super(PureSASLClient, self).__init__(*args, **kwargs)
+
+    def start(self, mechanism):
+        with error_catcher(self, SASLError):
+            if isinstance(mechanism, list):
+                self.choose_mechanism(mechanism)
+            else:
+                self.choose_mechanism([mechanism])
+            return True, self.mechanism, self.process()
+        # else
+        return False, mechanism, None
+
+    def encode(self, incoming):
+        with error_catcher(self):
+            return True, self.unwrap(incoming)
+        # else
+        return False, None
+
+    def decode(self, outgoing):
+        with error_catcher(self):
+            return True, self.wrap(outgoing)
+        # else
+        return False, None
+
+    def step(self, challenge=None):
+        with error_catcher(self):
+            return True, self.process(challenge)
+        # else
+        return False, None
+
+    def getError(self):
+        return self.error

pystellardb/stellar_hive.py

@@ -9,8 +9,9 @@ from __future__ import absolute_import
 from __future__ import unicode_literals
 
 from pyhive import hive
-from pystellardb import graph_types
-from pystellardb import stellar_rdd
+from . import graph_types
+from . import stellar_rdd
+from .sasl_compat import PureSASLClient
 
 from TCLIService import TCLIService
 from TCLIService import ttypes
@@ -102,12 +103,10 @@ class StellarConnection(object):
                 self._transport = thrift.transport.TTransport.TBufferedTransport(
                     socket)
             elif auth in ('LDAP', 'KERBEROS', 'NONE', 'CUSTOM'):
-                # Defer import so package dependency is optional
-                import sasl
                 import thrift_sasl
 
                 if auth == 'KERBEROS':
-                    # KERBEROS mode in hive.server2.authentication is GSSAPI in sasl library
+                    # KERBEROS mode in hive.server2.authentication is GSSAPI in SASL
                     sasl_auth = 'GSSAPI'
                 else:
                     sasl_auth = 'PLAIN'
@@ -116,20 +115,15 @@ class StellarConnection(object):
                         password = 'x'
 
                 def sasl_factory():
-                    sasl_client = sasl.Client()
-                    sasl_client.setAttr('host', host)
                     if sasl_auth == 'GSSAPI':
-                        sasl_client.setAttr('service', kerberos_service_name)
+                        sasl_client = PureSASLClient(host, mechanism="GSSAPI", service=kerberos_service_name)
                     elif sasl_auth == 'PLAIN':
-                        sasl_client.setAttr('username', username)
-                        sasl_client.setAttr('password', password)
+                        sasl_client = PureSASLClient(host, mechanism="PLAIN", username=username, password=password)
                     else:
-                        raise AssertionError
-                    sasl_client.init()
+                        raise AssertionError("Unsupported SASL mechanism")
                     return sasl_client
 
-                self._transport = thrift_sasl.TSaslClientTransport(
-                    sasl_factory, sasl_auth, socket)
+                self._transport = thrift_sasl.TSaslClientTransport(sasl_factory, sasl_auth, socket)
             else:
                 # All HS2 config options:
                 # https://cwiki.apache.org/confluence/display/Hive/Setting+Up+HiveServer2#SettingUpHiveServer2-Configuration
@@ -174,9 +168,12 @@ class StellarConnection(object):
                         schemaInJson)
                     
                     # get schema from data
-                    cursor.execute('manipulate graph {} get_schema_from_data'.format(graph_name))
-                    self._graph_schema_from_data = cursor.fetchone()[0]
-                    self._graph_schema_from_data = json.loads(self._graph_schema_from_data)
+                    try:
+                        cursor.execute('manipulate graph {} get_schema_from_data'.format(graph_name))
+                        self._graph_schema_from_data = cursor.fetchone()[0]
+                        self._graph_schema_from_data = json.loads(self._graph_schema_from_data)
+                    except:
+                        pass
             else:
                 assert response.serverProtocolVersion == protocol_version, \
                     "Unable to handle protocol version {}".format(response.serverProtocolVersion)

pystellardb/v2/__init__.py

@@ -0,0 +1,8 @@
+from .connection import Connection, CypherResult
+from .database import Graph
+
+__all__ = [
+    "Graph",
+    "Connection",
+    "CypherResult",
+]

pystellardb/v2/connection.py

@@ -0,0 +1,276 @@
+from __future__ import absolute_import
+from __future__ import unicode_literals
+
+from types import TracebackType
+from typing import Any, Union
+from pyhive import hive
+from pyhive.exc import *
+from TCLIService import TCLIService
+from TCLIService import ttypes
+
+import thrift.protocol.TBinaryProtocol
+import thrift.transport.TSocket
+import thrift.transport.TTransport
+import getpass
+import logging
+import json
+import sys
+
+from pystellardb.stellar_rdd import transformToRDD
+from pystellardb.graph_types import Vertex, Edge, Path, GraphSchema
+from pystellardb.sasl_compat import PureSASLClient
+
+__logger = logging.getLogger(__name__)
+
+class CypherResult(hive.Cursor):
+    """
+    QueryResult class for handling query results.
+    """
+
+    def __init__(self, connection: Any, graph_schema: GraphSchema, arraysize: int = 1000):
+        super(CypherResult, self).__init__(connection, arraysize)
+        self._graph_schema = graph_schema
+        self._column_comments = []
+
+    @property
+    def description(self):
+        """This read-only attribute is a sequence of 7-item sequences.
+
+        Each of these sequences contains information describing one result column:
+
+        - name
+        - type_code
+        - display_size (None in current implementation)
+        - internal_size (None in current implementation)
+        - precision (None in current implementation)
+        - scale (None in current implementation)
+        - null_ok (always True in current implementation)
+
+        This attribute will be ``None`` for operations that do not return rows or if the cursor has
+        not had an operation invoked via the :py:meth:`execute` method yet.
+
+        The ``type_code`` can be interpreted by comparing it to the Type Objects specified in the
+        section below.
+        """
+        if self._operationHandle is None or not self._operationHandle.hasResultSet:
+            return None
+        if self._description is None:
+            req = ttypes.TGetResultSetMetadataReq(self._operationHandle)
+            response = self._connection.client.GetResultSetMetadata(req)
+            hive._check_status(response)
+            columns = response.schema.columns
+            self._description = []
+            # If it's a cypher query, column comment is not null
+            self._column_comments = [
+                col.comment for col in response.schema.columns
+                if col.comment is not None
+            ]
+
+            for col in columns:
+                primary_type_entry = col.typeDesc.types[0]
+                if primary_type_entry.primitiveEntry is None:
+                    # All fancy stuff maps to string
+                    type_code = ttypes.TTypeId._VALUES_TO_NAMES[
+                        ttypes.TTypeId.STRING_TYPE]
+                else:
+                    type_id = primary_type_entry.primitiveEntry.type
+                    type_code = ttypes.TTypeId._VALUES_TO_NAMES[type_id]
+                self._description.append(
+                    (col.columnName.decode('utf-8')
+                     if sys.version_info[0] == 2 else col.columnName,
+                     type_code.decode('utf-8') if sys.version_info[0] == 2 else
+                     type_code, None, None, None, None, True))
+        return self._description
+    
+    def fetchone(self):
+        row = super(CypherResult, self).fetchone()
+
+        if row is None:
+            return None
+
+        parsed_row = []
+        for i in range(0, len(self._column_comments)):
+            parsed_row.append(
+                self._convertData(self._column_comments[i], row[i]))
+
+        return tuple(parsed_row)
+    
+    def _convertData(self, type, data):
+        """Convert Crux type to Readable type"""
+        if type == 'boolean':
+            return bool(data)
+        elif type == 'int':
+            return int(data)
+        elif type == 'long':
+            return int(data)
+        elif type == 'float' or type == 'double':
+            return float(data)
+        elif type == 'CruxType:Node' or type == 'GraphNode':
+            return Vertex.parseVertexFromJson(data)
+        elif type == 'CruxType:Relation' or type == 'GraphRelation':
+            return Edge.parseEdgeFromJson(self._graph_schema, data)
+        elif type == 'CruxType:Path':
+            return Path.parsePathFromJson(self._graph_schema, data)
+        elif type.startswith('CruxType:List'):
+            return self._parseList(type, data)
+        elif type.startswith('CruxType:Map'):
+            return self._parseMap(type, data)
+        else:
+            return data
+
+    def _parseList(self, type, data):
+        """Parse 'CruxType:List' type"""
+        parsed_data = json.loads(data)
+        newType = type[len('CruxType:List') + 1:type.find('>')]
+
+        return [self._convertData(newType, json.dumps(entry)) for entry in parsed_data]
+
+    def _parseMap(self, type, data):
+        """Parse 'CruxType:Map' type"""
+        parsed_data = json.loads(data)
+        newTypes = type[len('CruxType:Map') + 1:-2].split(',')
+
+        result = {}
+
+        for entry in parsed_data.keys():
+            key = self._convertData(newTypes[0], entry)
+            result[key] = self._convertData(newTypes[1], parsed_data[entry])
+
+        return result
+
+    def toRDD(self, sc, parallelism=1):
+        """
+        Transform to RDD
+        param sc: SparkContext
+        param parallelism: RDD parallelism
+        """
+        return transformToRDD(self, sc, parallelism)
+
+class Connection(object):
+    """
+    Connection class for connecting to the Stellar database.
+    """
+
+    def __init__(
+            self, 
+            host: str, 
+            port: int, 
+            auth: Union[str, None] = None, 
+            username: Union[str, None] = None, 
+            password: Union[str, None] = None,
+            kerberos_service_name: Union[str, None] = None,):
+        """Connect to HiveServer2
+
+        :param host: What host HiveServer2 runs on
+        :param port: What port HiveServer2 runs on. Defaults to 10000.
+        :param auth: The value of hive.server2.authentication used by HiveServer2. Defaults to ``NONE``.
+        :param username: Use with auth='LDAP' only
+        :param password: Use with auth='LDAP' only
+        :param kerberos_service_name: Use with auth='KERBEROS' only
+
+        The way to support LDAP and GSSAPI is originated from cloudera/Impyla:
+        https://github.com/cloudera/impyla/blob/255b07ed973d47a3395214ed92d35ec0615ebf62
+        /impala/_thrift_api.py#L152-L160
+        """
+        username = username or getpass.getuser()
+
+        if (password is not None) != (auth in ('LDAP', 'CUSTOM')):
+            raise ValueError(
+                "Password should be set if and only if in LDAP or CUSTOM mode; "
+                "Remove password or use one of those modes")
+        if (kerberos_service_name is not None) != (auth == 'KERBEROS'):
+            raise ValueError(
+                "kerberos_service_name should be set if and only if in KERBEROS mode"
+            )
+        if port is None:
+            port = 10000
+        if auth is None:
+            auth = 'NONE'
+        socket = thrift.transport.TSocket.TSocket(host, port)
+        if auth == 'NOSASL':
+            # NOSASL corresponds to hive.server2.authentication=NOSASL in hive-site.xml
+            self._transport = thrift.transport.TTransport.TBufferedTransport(socket)
+        elif auth in ('LDAP', 'KERBEROS', 'NONE', 'CUSTOM'):
+            import thrift_sasl
+
+            if auth == 'KERBEROS':
+                # KERBEROS mode in hive.server2.authentication is GSSAPI in SASL
+                sasl_auth = 'GSSAPI'
+            else:
+                sasl_auth = 'PLAIN'
+                if password is None:
+                    # Password doesn't matter in NONE mode, just needs to be nonempty.
+                    password = 'x'
+
+            def sasl_factory():
+                if sasl_auth == 'GSSAPI':
+                    sasl_client = PureSASLClient(host, mechanism="GSSAPI", service=kerberos_service_name)
+                elif sasl_auth == 'PLAIN':
+                    sasl_client = PureSASLClient(host, mechanism="PLAIN", username=username, password=password)
+                else:
+                    raise AssertionError("Unsupported SASL mechanism")
+                return sasl_client
+
+            self._transport = thrift_sasl.TSaslClientTransport(sasl_factory, sasl_auth, socket)
+        else:
+            # All HS2 config options:
+            # https://cwiki.apache.org/confluence/display/Hive/Setting+Up+HiveServer2#SettingUpHiveServer2-Configuration
+            # PAM currently left to end user via thrift_transport option.
+            raise NotImplementedError(
+                "Only NONE, NOSASL, LDAP, KERBEROS, CUSTOM "
+                "authentication are supported, got {}".format(auth))
+        protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol(self._transport)
+        self._client = TCLIService.Client(protocol)
+        # oldest version that still contains features we care about
+        # "V6 uses binary type for binary payload (was string) and uses columnar result set"
+        protocol_version = ttypes.TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V6
+        try:
+            self._transport.open()
+            open_session_req = ttypes.TOpenSessionReq(
+                client_protocol=protocol_version,
+                configuration={},
+                username=username,
+            )
+            response = self._client.OpenSession(open_session_req)
+            hive._check_status(response)
+            assert response.sessionHandle is not None, "Expected a session from OpenSession"
+            self._sessionHandle = response.sessionHandle
+            assert response.serverProtocolVersion == protocol_version, \
+                "Unable to handle protocol version {}".format(response.serverProtocolVersion)
+            
+        except:
+            self._transport.close()
+            raise
+
+    def __enter__(self):
+        """Transport should already be opened by __init__"""
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Call close"""
+        self.close()
+
+    @property
+    def client(self):
+        return self._client
+
+    @property
+    def sessionHandle(self):
+        return self._sessionHandle
+
+    def close(self):
+        req = ttypes.TCloseSessionReq(sessionHandle=self._sessionHandle)
+        response = self._client.CloseSession(req)
+        self._transport.close()
+        hive._check_status(response)
+
+    def execute(self, operation: str, graph_schema: Union[GraphSchema, None] = None) -> CypherResult:
+        """
+        Execute a query on the database.
+
+        :param operation: The query to execute
+        """
+        cursor = CypherResult(self, graph_schema)
+        cursor.execute(operation)
+        return cursor
+        

pystellardb/v2/database.py

@@ -0,0 +1,110 @@
+"""
+Define database class for operations on a single graph.
+"""
+
+from __future__ import absolute_import
+from types import TracebackType
+import logging
+import json
+from typing import Any, Union
+
+from ..graph_types import GraphSchema
+from .connection import Connection, CypherResult
+
+_logger = logging.getLogger(__name__)
+
+class Graph(object):
+    """
+    Database class for operations on a single graph.
+    """
+
+    def __init__(self, graph_name: str, connection: Connection):
+        self.graph_name = graph_name
+        self.connection = connection
+        self.graph_schema: GraphSchema = None
+        self._init_connection()
+
+    def __enter__(self):
+        """
+        Enter the database context.
+        """
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Exit the database context.
+        """
+        if self.connection:
+            self.connection.close()
+        self.connection = None
+
+    def _init_connection(self):
+        """
+        Initialize the connection to the database.
+        """
+        
+        self.connection.execute('config query.lang cypher')
+        
+        """Try to bind cursor with graph name."""
+        try:
+            self.connection.execute(f'use graph {self.graph_name}')
+            self._get_graph_schema()
+        except Exception as e:
+            _logger.debug(f"graph {self.graph_name} not found")
+
+    def _get_graph_schema(self):
+        """
+        Get the graph schema.
+        """
+        query_result = self.connection.execute('DESCRIBE GRAPH {} RAW'.format(self.graph_name))
+        schemaInJson = query_result.fetchone()[0]
+        query_result.close()
+        self.graph_schema = GraphSchema.parseSchemaFromJson(schemaInJson)
+        
+        # get schema from data
+        try:
+            query_result = self.connection.execute('manipulate graph {} get_schema_from_data'.format(self.graph_name))
+            self.graph_labels = query_result.fetchone()[0]
+            query_result.close()
+            self.graph_labels = json.loads(self.graph_labels)
+        except:
+            pass
+
+    def execute(self, query: str, *args, **kwargs) -> CypherResult:
+        """
+        Execute a query on the database.
+        """
+        query_result = self.connection.execute(query, self.graph_schema)
+
+        if (query.strip().lower().startswith('create graph')):
+            self._get_graph_schema()
+            
+        return query_result
+
+    def create(self, cypher: str):
+        """
+        Create a graph.
+        """
+        create_result = self.connection.execute(cypher)
+        create_result.close()
+        self._get_graph_schema()
+    
+    def drop(self):
+        """
+        Drop the graph.
+        """
+        drop_result = self.connection.execute('DROP GRAPH {}'.format(self.graph_name))
+        drop_result.close()
+        self.graph_schema = None
+
+    def get_graph_schema(self) -> Union[GraphSchema, None]:
+        """
+        Get the graph schema.
+        """
+        return self.graph_schema
+    
+    def get_labels(self) -> Any:
+        """
+        Get the labels of the graph.
+        """
+        return self.graph_labels

PyStellarDB-0.13.5.dist-info/RECORD

@@ -1,10 +0,0 @@
-pystellardb/__init__.py,sha256=JOl41NviMN-qDV0Z8ZPmhNIxvgyauGGJHdB4A-8MhqM,93
-pystellardb/_version.py,sha256=cwBvUjWVT20dAp04I_QQFlk81VsfhUWKUfT2FKtmxhk,498
-pystellardb/graph_types.py,sha256=XbtPebhaRV7OimW-GxFJGSOcp-OnzWrlIe9M7Xv3sbU,14032
-pystellardb/stellar_hive.py,sha256=xVnONjG03CLNwW3dymR2ZqXiEUvcMlIvPIOxzyoDFfE,15067
-pystellardb/stellar_rdd.py,sha256=TYwsWYeCxfOliGq1kV3ArNXdye55cKWZF7s9M9nDdt4,1324
-PyStellarDB-0.13.5.dist-info/LICENSE,sha256=1qDFxrywejs7xNBfOr6T-7lOuqDgSNIES77kTYege3w,560
-PyStellarDB-0.13.5.dist-info/METADATA,sha256=Lg8U6us8NJ-hlFQEwaZDsfGZOgvPKnChM7lY6HtTt3s,9390
-PyStellarDB-0.13.5.dist-info/WHEEL,sha256=_4XEmVmaBFWtekSGrbfOGNjC2I5lUr0lZSRblBllIFA,109
-PyStellarDB-0.13.5.dist-info/top_level.txt,sha256=DRk-SeGVCdVAzv2CwFmdu75Yo7DgjUA3Hpu-9l8qPuU,12
-PyStellarDB-0.13.5.dist-info/RECORD,,