improve columnMetaData (#57)

koxudaxi · web-flow · commit 19e9274c20b9 · 2020-04-14T02:05:06.000+09:00
* improve columnMetaData

* reformat
diff --git a/local_data_api/resources/jdbc/__init__.py b/local_data_api/resources/jdbc/__init__.py
@@ -54,6 +54,30 @@ def __init__(self, connection: Connection, transaction_id: Optional[str] = None)
         self.autocommit: bool = True
         super().__init__(connection, transaction_id)
 
+    def create_column_metadata_set(
+        self, cursor: jaydebeapi.Cursor
+    ) -> List[ColumnMetadata]:
+        meta = getattr(cursor, '_meta')
+        return [
+            ColumnMetadata(
+                arrayBaseColumnType=0,
+                isAutoIncrement=meta.isAutoIncrement(i),
+                isCaseSensitive=meta.isCaseSensitive(i),
+                isCurrency=meta.isCurrency(i),
+                isSigned=meta.isSigned(i),
+                label=meta.getColumnLabel(i),
+                name=meta.getColumnName(i),
+                nullable=meta.isNullable(i),
+                precision=meta.getPrecision(i),
+                scale=meta.getScale(i),
+                schema=meta.getSchemaName(i),
+                tableName=meta.getTableName(i),
+                type=meta.getColumnType(i),
+                typeName=meta.getColumnTypeName(i),
+            )
+            for i in range(1, meta.getColumnCount() + 1)
+        ]
+
     def autocommit_off(self, cursor: jaydebeapi.Cursor) -> None:
         self.connection.jconn.setAutoCommit(False)
         self.autocommit = False
@@ -96,8 +120,9 @@ def execute(
                         ],
                     )
                     if include_result_metadata:
-                        meta = getattr(cursor, '_meta')
-                        response.columnMetadata = create_column_metadata_set(meta)
+                        response.columnMetadata = self.create_column_metadata_set(
+                            cursor
+                        )
                     return response
                 else:
                     rowcount: int = cursor.rowcount
@@ -144,25 +169,3 @@ def create_connection_maker(
             {"user": user_name, "password": password},
             engine_kwargs['JAR_PATH'],
         )
-
-
-def create_column_metadata_set(meta: Any) -> List[ColumnMetadata]:
-    return [
-        ColumnMetadata(
-            arrayBaseColumnType=0,
-            isAutoIncrement=meta.isAutoIncrement(i),
-            isCaseSensitive=meta.isCaseSensitive(i),
-            isCurrency=meta.isCurrency(i),
-            isSigned=meta.isSigned(i),
-            label=meta.getColumnLabel(i),
-            name=meta.getColumnName(i),
-            nullable=meta.isNullable(i),
-            precision=meta.getPrecision(i),
-            scale=meta.getScale(i),
-            schema=meta.getSchemaName(i),
-            tableName=meta.getTableName(i),
-            type=meta.getColumnType(i),
-            typeName=meta.getColumnTypeName(i),
-        )
-        for i in range(1, meta.getColumnCount() + 1)
-    ]
diff --git a/local_data_api/resources/mysql.py b/local_data_api/resources/mysql.py
@@ -1,18 +1,49 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import pymysql
+from pymysql.constants import FIELD_TYPE
+from pymysql.protocol import FieldDescriptorPacket
 from sqlalchemy.dialects import mysql
 
+from local_data_api.models import ColumnMetadata
 from local_data_api.resources.resource import Resource, register_resource_type
 
 if TYPE_CHECKING:  # pragma: no cover
-    from local_data_api.resources.resource import ConnectionMaker
+    from local_data_api.resources.resource import ConnectionMaker, Cursor
+
+FIELD_TYPE_MAP: Dict[int, str] = {
+    getattr(FIELD_TYPE, k): k for k in dir(FIELD_TYPE) if not k.startswith('_')
+}
+
+
+def create_column_metadata(
+    field_descriptor_packet: FieldDescriptorPacket,
+) -> ColumnMetadata:
+    return ColumnMetadata(
+        arrayBaseColumnType=0,
+        isAutoIncrement=False,
+        isCaseSensitive=False,
+        isCurrency=False,
+        isSigned=False,
+        label=field_descriptor_packet.name,
+        name=field_descriptor_packet.org_name,
+        nullable=1 if field_descriptor_packet.flags % 2 == 0 else 0,
+        precision=field_descriptor_packet.get_column_length(),
+        scale=field_descriptor_packet.scale,
+        schema=None,
+        tableName=field_descriptor_packet.table_name,
+        type=None,  # JDBC Type unsupported
+        typeName=None,  # JDBC TypeName unsupported
+    )
 
 
 @register_resource_type
 class MySQL(Resource):
+    def create_column_metadata_set(self, cursor: Cursor) -> List[ColumnMetadata]:
+        return [create_column_metadata(f) for f in getattr(cursor, '_result').fields]
+
     DIALECT = mysql.dialect(paramstyle='named')
 
     @classmethod
diff --git a/local_data_api/resources/postgres.py b/local_data_api/resources/postgres.py
@@ -1,18 +1,24 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import psycopg2
 from sqlalchemy.dialects import postgresql
 
+from local_data_api.models import ColumnMetadata
 from local_data_api.resources.resource import Resource, register_resource_type
 
 if TYPE_CHECKING:  # pragma: no cover
-    from local_data_api.resources.resource import ConnectionMaker
+    from local_data_api.resources.resource import ConnectionMaker, Cursor
 
 
 @register_resource_type
 class PostgresSQL(Resource):
+    def create_column_metadata_set(
+        self, cursor: Cursor
+    ) -> List[ColumnMetadata]:  # pragma: no cover
+        raise NotImplementedError
+
     DIALECT = postgresql.dialect(paramstyle='named')
 
     @classmethod
diff --git a/local_data_api/resources/resource.py b/local_data_api/resources/resource.py
@@ -5,9 +5,11 @@
 import string
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from enum import Enum
 from hashlib import sha1
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Pattern, Tuple, Type, Union
 
+from pydantic.main import BaseModel
 from sqlalchemy import text
 from sqlalchemy.engine import Dialect
 from sqlalchemy.exc import ArgumentError, CompileError
@@ -222,31 +224,46 @@ def get_resource(
     return meta.resource_type(connection, transaction_id)
 
 
-def create_column_metadata(
-    name: str,
-    type_code: int,
-    display_size: Optional[int],
-    internal_size: int,
-    precision: int,
-    scale: int,
-    null_ok: bool,
-) -> ColumnMetadata:
-    return ColumnMetadata(
-        arrayBaseColumnType=0,
-        isAutoIncrement=False,
-        isCaseSensitive=False,
-        isCurrency=False,
-        isSigned=False,
-        label=name,
-        name=name,
-        nullable=1 if null_ok else 0,
-        precision=precision,
-        scale=scale,
-        schema=None,
-        tableName=None,
-        type=None,
-        typeName=None,
-    )
+class JDBCType(Enum):
+    BIT = -7
+    TINYINT = -6
+    SMALLINT = 5
+    INTEGER = 4
+    BIGINT = -5
+    FLOAT = 6
+    REAL = 7
+    DOUBLE = 8
+    NUMERIC = 2
+    DECIMAL = 3
+    CHAR = 1
+    VARCHAR = 12
+    LONGVARCHAR = -1
+    DATE = 91
+    TIME = 92
+    TIMESTAMP = 93
+    BINARY = -2
+    VARBINARY = -3
+    LONGVARBINARY = -4
+    NULL = 0
+    OTHER = 1111
+    JAVA_OBJECT = 2000
+    DISTINCT = 2001
+    STRUCT = 2002
+    ARRAY = 2003
+    BLOB = 2004
+    CLOB = 2005
+    REF = 2006
+    DATALINK = 70
+    BOOLEAN = 16
+    ROWID = -8
+    NCHAR = -15
+    NVARCHAR = -9
+    LONGNVARCHAR = -16
+    NCLOB = 2011
+    SQLXML = 2009
+    REF_CURSOR = 2012
+    TIME_WITH_TIMEZONE = 2013
+    TIMESTAMP_WITH_TIMEZONE = 2014
 
 
 class Resource(ABC):
@@ -309,6 +326,10 @@ def create_transaction_id() -> str:
             for _ in range(TRANSACTION_ID_LENGTH)
         )
 
+    @abstractmethod
+    def create_column_metadata_set(self, cursor: Cursor) -> List[ColumnMetadata]:
+        raise NotImplementedError
+
     def close(self) -> None:
         self.connection.close()
         if self.transaction_id in CONNECTION_POOL:
@@ -351,9 +372,9 @@ def execute(
                         ],
                     )
                     if include_result_metadata:
-                        response.columnMetadata = [
-                            create_column_metadata(*d) for d in cursor.description
-                        ]
+                        response.columnMetadata = self.create_column_metadata_set(
+                            cursor
+                        )
                     return response
                 else:
                     rowcount: int = cursor.rowcount
diff --git a/local_data_api/resources/sqlite.py b/local_data_api/resources/sqlite.py
@@ -1,18 +1,22 @@
 from __future__ import annotations
 
 import sqlite3
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from sqlalchemy.dialects import sqlite
 
+from local_data_api.models import ColumnMetadata
 from local_data_api.resources.resource import Resource, register_resource_type
 
 if TYPE_CHECKING:  # pragma: no cover
-    from local_data_api.resources.resource import ConnectionMaker
+    from local_data_api.resources.resource import ConnectionMaker, Cursor
 
 
 @register_resource_type
 class SQLite(Resource):
+    def create_column_metadata_set(self, cursor: Cursor) -> List[ColumnMetadata]:
+        raise NotImplementedError
+
     DIALECT = sqlite.dialect(paramstyle='named')
 
     @classmethod
diff --git a/tests/test_resource/test_jdbc/test_jdbc.py b/tests/test_resource/test_jdbc/test_jdbc.py
@@ -4,12 +4,7 @@
 import pytest
 
 from local_data_api.models import ColumnMetadata
-from local_data_api.resources.jdbc import (
-    JDBC,
-    attach_thread_to_jvm,
-    connection_maker,
-    create_column_metadata_set,
-)
+from local_data_api.resources.jdbc import JDBC, attach_thread_to_jvm, connection_maker
 
 
 class DummyJDBC(JDBC):
@@ -78,8 +73,9 @@ def test_create_column_metadata_set(mocker):
     mock_meta.getColumnType.side_effect = [5, 6]
     mock_meta.getColumnTypeName.side_effect = ['i', 'j']
     mock_meta.getColumnCount.return_value = 2
-
-    assert create_column_metadata_set(mock_meta) == [
+    mock_cursor = mocker.Mock()
+    mock_cursor._meta = mock_meta
+    assert DummyJDBC(mocker.Mock()).create_column_metadata_set(mock_cursor) == [
         ColumnMetadata(
             arrayBaseColumnType=0,
             isAutoIncrement=True,
diff --git a/tests/test_resource/test_jdbc/test_mysql.py b/tests/test_resource/test_jdbc/test_mysql.py
@@ -124,9 +124,7 @@ def test_execute_select_with_include_metadata(mocked_connection, mocked_cursor,
     mocked_cursor.description = (1, 2, 3, 4, 5, 6, 7), (8, 9, 10, 11, 12, 13, 14)
     mocked_cursor.fetchall.side_effect = [((1, 'abc'),)]
     dummy = MySQLJDBC(mocked_connection, transaction_id='123')
-    create_column_metadata_set_mock = mocker.patch(
-        'local_data_api.resources.jdbc.create_column_metadata_set'
-    )
+    dummy.create_column_metadata_set = create_column_metadata_set_mock = mocker.Mock()
     create_column_metadata_set_mock.side_effect = [
         [
             ColumnMetadata(
@@ -197,7 +195,7 @@ def test_execute_select_with_include_metadata(mocked_connection, mocked_cursor,
         ],
     )
 
-    create_column_metadata_set_mock.assert_called_once_with(meta_mock)
+    create_column_metadata_set_mock.assert_called_once_with(mocked_cursor)
     mocked_cursor.execute.assert_has_calls(
         [mocker.call('SELECT LAST_INSERT_ID(NULL)'), mocker.call('select * from users')]
     )
diff --git a/tests/test_resource/test_jdbc/test_postgres.py b/tests/test_resource/test_jdbc/test_postgres.py
@@ -99,14 +99,10 @@ def test_execute_select(mocked_connection, mocked_cursor, mocker):
 
 
 def test_execute_select_with_include_metadata(mocked_connection, mocked_cursor, mocker):
-    meta_mock = mocker.Mock()
-    mocked_cursor._meta = meta_mock
     mocked_cursor.description = (1, 2, 3, 4, 5, 6, 7), (8, 9, 10, 11, 12, 13, 14)
     mocked_cursor.fetchall.side_effect = [((1, 'abc'),)]
     dummy = PostgreSQLJDBC(mocked_connection, transaction_id='123')
-    create_column_metadata_set_mock = mocker.patch(
-        'local_data_api.resources.jdbc.create_column_metadata_set'
-    )
+    dummy.create_column_metadata_set = create_column_metadata_set_mock = mocker.Mock()
     create_column_metadata_set_mock.side_effect = [
         [
             ColumnMetadata(
@@ -177,7 +173,7 @@ def test_execute_select_with_include_metadata(mocked_connection, mocked_cursor,
         ],
     )
 
-    create_column_metadata_set_mock.assert_called_once_with(meta_mock)
+    create_column_metadata_set_mock.assert_called_once_with(mocked_cursor)
     mocked_cursor.execute.assert_has_calls([mocker.call('select * from users')])
     mocked_cursor.close.assert_called_once_with()
 
diff --git a/tests/test_resource/test_mysql.py b/tests/test_resource/test_mysql.py
diff --git a/tests/test_resource/test_resource.py b/tests/test_resource/test_resource.py

Original file line number	Diff line number	Diff line change
`@@ -124,9 +124,7 @@ def test_execute_select_with_include_metadata(mocked_connection, mocked_cursor,`
`124`	`124`	`mocked_cursor.description = (1, 2, 3, 4, 5, 6, 7), (8, 9, 10, 11, 12, 13, 14)`
`125`	`125`	`mocked_cursor.fetchall.side_effect = [((1, 'abc'),)]`
`126`	`126`	`dummy = MySQLJDBC(mocked_connection, transaction_id='123')`
`127`		`- create_column_metadata_set_mock = mocker.patch(`
`128`		`- 'local_data_api.resources.jdbc.create_column_metadata_set'`
`129`		`- )`
	`127`	`+ dummy.create_column_metadata_set = create_column_metadata_set_mock = mocker.Mock()`
`130`	`128`	`create_column_metadata_set_mock.side_effect = [`
`131`	`129`	`[`
`132`	`130`	`ColumnMetadata(`
`@@ -197,7 +195,7 @@ def test_execute_select_with_include_metadata(mocked_connection, mocked_cursor,`
`197`	`195`	`],`
`198`	`196`	`)`
`199`	`197`
`200`		`- create_column_metadata_set_mock.assert_called_once_with(meta_mock)`
	`198`	`+ create_column_metadata_set_mock.assert_called_once_with(mocked_cursor)`
`201`	`199`	`mocked_cursor.execute.assert_has_calls(`
`202`	`200`	`[mocker.call('SELECT LAST_INSERT_ID(NULL)'), mocker.call('select * from users')]`
`203`	`201`	`)`