Merge pull request #10 from dungdm93/support-pandas

Supporting pandas

Author: dungdm93

sqlalchemy_trino/compiler.py

@@ -85,7 +85,44 @@ class TrinoDDLCompiler(compiler.DDLCompiler):
 
 
 class TrinoTypeCompiler(compiler.GenericTypeCompiler):
-    pass
+    def visit_FLOAT(self, type_, **kw):
+        precision = type_.precision or 32
+        if 0 <= precision <= 32:
+            return self.visit_REAL(type_, **kw)
+        elif 32 < precision <= 64:
+            return self.visit_DOUBLE(type_, **kw)
+        else:
+            raise ValueError(f"type.precision={type_.precision} is invalid")
+
+    def visit_DOUBLE(self, type_, **kw):
+        return "DOUBLE"
+
+    def visit_NUMERIC(self, type_, **kw):
+        return self.visit_DECIMAL(type_, **kw)
+
+    def visit_NCHAR(self, type_, **kw):
+        return self.visit_CHAR(type_, **kw)
+
+    def visit_NVARCHAR(self, type_, **kw):
+        return self.visit_VARCHAR(type_, **kw)
+
+    def visit_TEXT(self, type_, **kw):
+        return self.visit_VARCHAR(type_, **kw)
+
+    def visit_BINARY(self, type_, **kw):
+        return self.visit_VARBINARY(type_, **kw)
+
+    def visit_CLOB(self, type_, **kw):
+        return self.visit_VARCHAR(type_, **kw)
+
+    def visit_NCLOB(self, type_, **kw):
+        return self.visit_VARCHAR(type_, **kw)
+
+    def visit_BLOB(self, type_, **kw):
+        return self.visit_VARBINARY(type_, **kw)
+
+    def visit_DATETIME(self, type_, **kw):
+        return self.visit_TIMESTAMP(type_, **kw)
 
 
 class TrinoIdentifierPreparer(compiler.IdentifierPreparer):

sqlalchemy_trino/dialect.py

@@ -7,6 +7,8 @@
 from sqlalchemy.engine.default import DefaultDialect, DefaultExecutionContext
 from sqlalchemy.engine.url import URL
 from trino.auth import BasicAuthentication
+from trino.client import TrinoQuery
+from trino.dbapi import Cursor
 
 from . import compiler
 from . import datatype
@@ -58,6 +60,7 @@ class TrinoDialect(DefaultDialect):
     # DML
     supports_empty_insert = False
     supports_multivalues_insert = True
+    postfetch_lastrowid = False
 
     # Version parser
     __version_pattern = re.compile(r'(\d+).*')
@@ -272,6 +275,19 @@ def _get_default_schema_name(self, connection: Connection) -> Optional[str]:
         dbapi_connection: trino_dbapi.Connection = connection.connection
         return dbapi_connection.schema
 
+    def do_execute(self, cursor: Cursor, statement: str, parameters: Tuple[Any, ...],
+                   context: DefaultExecutionContext = None):
+        cursor.execute(statement, parameters)
+        if context and context.should_autocommit:
+            # SQL statement only submitted to Trino server when cursor.fetch*() is called.
+            # For DDL (CREATE/ALTER/DROP) and DML (INSERT/UPDATE/DELETE) statement, call cursor.description
+            # to force submit statement immediately.
+            d = cursor.description
+            # old trino client does not support eager-loading cursor.description
+            if d is None:
+                query: TrinoQuery = cursor._query  # noqa
+                query._result._rows += query.fetch()  # noqa
+
     def do_rollback(self, dbapi_connection):
         if dbapi_connection.transaction is not None:
             dbapi_connection.rollback()