Add an impersonation_target parm to drill+sadrill URLs. (#74)

When present, this parameter will be converted to a userName
property in POSTs made to /query.json.

Author: jnturton

.editorconfig

@@ -0,0 +1,14 @@
+root = true
+
+[*]
+indent_style = space
+indent_size = 2
+end_of_line = lf
+charset = utf-8
+trim_trailing_whitespace = true
+insert_final_newline = true
+
+[*.py]
+indent_size = 4
+max_line_length = 80
+

CHANGELOG.md

@@ -1,12 +1,20 @@
 ## [Unreleased]
 
+## [1.1.2] - 2022-03-14
+
+### Changed
+
+- Add an impersonation_target parm to drill+sadrill URLs. When present,
+  this parameter will be converted to a userName property in POSTs made to
+  /query.json.
+
 ## [1.1.1] - 2021-07-28
 
 ### Fixed
 
-- Backwards compatibility with Drill < 1.19, limited to returning
-  all data values as strings. Users not able to upgrade to >= 1.19
-  must implement their own typecasting or use sqlalchemy-drill 0.3.
+- Backwards compatibility with Drill < 1.19, limited to returning all data
+  values as strings. Users not able to upgrade to >= 1.19 must implement their
+  own typecasting or use sqlalchemy-drill 0.3.
 
 ## [1.1.0] - 2021-07-21
 

README.md

@@ -22,9 +22,9 @@ Alternatively, you can download the latest release from github and install from
 python3 -m pip install git+https://github.com/JohnOmernik/sqlalchemy-drill.git
 ```
 
-## Usage
+## Usage with REST
 
-To use Drill with SQLAlchemy you will need to craft a connection string in the format below:
+Drill's REST API can execute queries with results streamed to JSON returned over chunked HTTP for Drill >= 1.19, otherwise with results buffered and then returned in a conventional HTTP response.  A SQLAlchemy URL to connect to Drill over REST looks like the following.
 
 ```
 drill+sadrill://<username>:<password>@<host>:<port>/<storage_plugin>?use_ssl=True
@@ -36,7 +36,20 @@ To connect to Drill running on a local machine running in embedded mode you can
 drill+sadrill://localhost:8047/dfs?use_ssl=False
 ```
 
+### Supported URL query parameters
+
+| Parameter                 | Type    | Description                                                    |
+| ------------------------- | ------- | -------------------------------------------------------------- |
+| use_ssl                   | boolean | Whether to connect to Drill using HTTPS                        |
+| verify_ssl                | boolean | Whether to verify the server's TLS certificate                 |
+| impersonation_target\[1\] | string  | Username of a Drill user to be impersonated by this connection |
+
+[1] Requires a build of Drill that incorporates the fix for DRILL-8168.
+
+### Trailing metadata
+
 Query result metadata returned by the Drill REST API is stored in the `result_md` field of the DB-API Cursor object.  Note that any trailing metadata, i.e. metadata which comes after result row data, will only be populated after you have iterated through all of the returned rows.  If you need this trailing metadata you can make the cursor object reachable after it has been completely iterated by obtaining a reference to it beforehand, as follows.
+
 ```python
 r = engine.execute('select current_timestamp')
 r.cursor.result_md  # access metadata, but only leading metadata
@@ -46,11 +59,13 @@ cur.result_md       # access metadata, including trailing metadata
 del cur             # optionally delete the reference when done
 ```
 
-### Changes in Drill 1.19 affecting drill+sadrill
+### Drill < 1.19
 
 In versions of Drill earlier than 1.19, all data values are serialised to JSON strings and column type metadata comes after the data itself.  As a result, for these versions of Drill, the drill+sadrill dialect returns every data value as a string.  To convert non-string data to its native type you need to typecast it yourself.
 
-In Drill 1.19 the REST API began making use of numeric types in JSON for numbers and times, the latter via a UNIX time representation.  As a result, the drill+sadrill dialect is able to return appropriate types for numbers and times when used with Drill >= 1.19.
+### Drill >= 1.19
+
+In Drill 1.19 the REST API began making use of numeric types in JSON for numbers and times, the latter via a UNIX time representation while column type metadata was moved ahead of the result data.  Because of this, the drill+sadrill dialect is able to return appropriate types for numbers and times when used with Drill >= 1.19.
 
 ## Usage with JDBC
 

setup.py

@@ -29,7 +29,7 @@
     long_description = f.read()
 
 setup(name='sqlalchemy_drill',
-      version='1.1.1',
+      version='1.1.2',
       description="Apache Drill for SQLAlchemy",
       long_description=long_description,
       long_description_content_type="text/markdown",
@@ -64,7 +64,7 @@
       license='MIT',
       url='https://github.com/JohnOmernik/sqlalchemy-drill',
       download_url='https://github.com/JohnOmernik/sqlalchemy-drill/archive/'
-      '1.1.1.tar.gz',
+      '1.1.2.tar.gz',
       packages=find_packages(),
       include_package_data=True,
       tests_require=['nose >= 0.11'],

sqlalchemy_drill/__init__.py

@@ -19,7 +19,7 @@
 # OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
-__version__ = '1.1.1'
+__version__ = '1.1.2'
 from sqlalchemy.dialects import registry
 
 registry.register("drill", "sqlalchemy_drill.sadrill", "DrillDialect_sadrill")

sqlalchemy_drill/drilldbapi/_drilldbapi.py

@@ -331,13 +331,15 @@ def __init__(self,
                  host: str,
                  port: int,
                  proto: str,
+                 impersonation_target: str,
                  session: Session):
         if session is None:
             raise ProgrammingError('A Requests session is required.', None)
 
         self._base_url = f'{proto}{host}:{port}'
         self._session = session
         self._connected = True
+        self._impersonation_target = impersonation_target
 
         logger.debug('queries Drill\'s version number...')
         resp = self.submit_query(
@@ -361,22 +363,31 @@ def __init__(self,
 
     def submit_query(self, query: str):
         payload = api_globals._PAYLOAD.copy()
+        payload['userName'] = self._impersonation_target
+
         # TODO: autoLimit, defaultSchema
         payload['query'] = query
 
         logger.debug('sends an HTTP POST with payload')
         logger.debug(payload)
 
-        return self._session.post(
+        resp = self._session.post(
             f'{self._base_url}/query.json',
             data=dumps(payload),
             headers=api_globals._HEADER,
             timeout=None,
             stream=True
         )
 
-    # Decorator for methods which require connection
+        if resp.status_code == 200:
+            return resp
+        else:
+            raise DatabaseError(
+                resp.json().get('errorMessage', None),
+                resp.status_code
+            )
 
+    # Decorator for methods which require connection
     def connected(func):
 
         def func_wrapper(self, *args, **kwargs):
@@ -419,24 +430,24 @@ def connect(host: str,
             drilluser: str = None,
             drillpass: str = None,
             verify_ssl: bool = False,
-            ca_certs: bool = None,
+            impersonation_target: str = None
             ) -> Connection:
 
     session = Session()
 
-    if verify_ssl is False:
-        session.verify = False
-    else:
-        if ca_certs is not None:
-            session.verify = ca_certs
-        else:
-            session.verify = True
-
+    session.verify = verify_ssl
     proto = 'https://' if use_ssl in [True, 'True', 'true'] else 'http://'
     base_url = f'{proto}{host}:{port}'
 
+    logging.info(
+        f'will log in with user {drilluser} and impersonation target '
+        f'{impersonation_target}'
+    )
+
     if drilluser is None:
         payload = api_globals._PAYLOAD.copy()
+        payload['userName'] = impersonation_target
+
         payload['query'] = 'show schemas'
         response = session.post(
             f'{base_url}/query.json',
@@ -464,7 +475,7 @@ def connect(host: str,
         logger.error('failed to authenticate to Drill.')
         raise AuthError(str(raw_data), response.status_code)
 
-    conn = Connection(host, port, proto, session)
+    conn = Connection(host, port, proto, impersonation_target, session)
     if db is not None:
         conn.submit_query(f'USE {db}')