Add minimal support for querying multiple catalogs

pyhive/presto.py

@@ -17,6 +17,7 @@
 import getpass
 import logging
 import requests
+import datetime
 from requests.auth import HTTPBasicAuth
 
 try:  # Python 3
@@ -31,7 +32,25 @@
 paramstyle = 'pyformat'  # Python extended format codes, e.g. ...WHERE name=%(name)s
 
 _logger = logging.getLogger(__name__)
-_escaper = common.ParamEscaper()
+
+
+class PrestoParamEscaper(common.ParamEscaper):
+    def escape_item(self, item):
+        if isinstance(item, datetime.datetime):
+            return self.escape_datetime(item)
+        elif isinstance(item, datetime.date):
+            return self.escape_date(item)
+        else:
+            return super(PrestoParamEscaper, self).escape_item(item)
+
+    def escape_date(self, item):
+        return "date '{}'".format(item)
+
+    def escape_datetime(self, item):
+        return "timestamp  '{}'".format(item)
+
+
+_escaper = PrestoParamEscaper()
 
 
 def connect(*args, **kwargs):

pyhive/sqlalchemy_presto.py

@@ -12,11 +12,13 @@
 from sqlalchemy import exc
 from sqlalchemy import types
 from sqlalchemy import util
+
 # TODO shouldn't use mysql type
 from sqlalchemy.databases import mysql
 from sqlalchemy.engine import default
 from sqlalchemy.sql import compiler
 from sqlalchemy.sql.compiler import SQLCompiler
+from sqlalchemy.sql.expression import Alias
 
 from pyhive import presto
 from pyhive.common import UniversalSet
@@ -46,6 +48,37 @@
     def visit_char_length_func(self, fn, **kw):
         return 'length{}'.format(self.function_argspec(fn, **kw))
 
+    def visit_column(self, column, add_to_result_map=None, include_table=True, **kwargs):
+        sql = super(PrestoCompiler, self).visit_column(
+            column, add_to_result_map, include_table, **kwargs
+        )
+        table = column.table
+        return self.__add_catalog(sql, table)
+
+    def visit_table(self, table, asfrom=False, iscrud=False, ashint=False,
+                    fromhints=None, use_schema=True, **kwargs):
+        sql = super(PrestoCompiler, self).visit_table(
+            table, asfrom, iscrud, ashint, fromhints, use_schema, **kwargs
+        )
+        return self.__add_catalog(sql, table)
+
+    def __add_catalog(self, sql, table):
+        if table is None:
+            return sql
+
+        if isinstance(table, Alias):
+            return sql
+
+        if (
+            "presto" not in table.dialect_options
+            or "catalog" not in table.dialect_options["presto"]._non_defaults
+        ):
+            return sql
+
+        catalog = table.dialect_options["presto"]._non_defaults["catalog"]
+        sql = "\"{catalog}\".{sql}".format(catalog=catalog, sql=sql)
+        return sql
+
 
 class PrestoTypeCompiler(compiler.GenericTypeCompiler):
     def visit_CLOB(self, type_, **kw):