PostgreSQLConnection -> ServerConnection

ossdbtoolsservice/chat/plugin/postgres_plugin.py

@@ -14,7 +14,6 @@
 from ossdbtoolsservice.connection.connection_service import ConnectionService
 from ossdbtoolsservice.connection.contracts.common import ConnectionType
 from ossdbtoolsservice.driver.types.driver import ServerConnection
-from ossdbtoolsservice.driver.types.psycopg_driver import PostgreSQLConnection
 from ossdbtoolsservice.hosting import RequestContext
 
 from .postgres_utils import (
@@ -85,7 +84,7 @@ def get_schema_kernelfunc(
         connection = self._get_connection()
         if connection is None:
             return "Error. Could not connect to the database. No connection found."
-        assert isinstance(connection, PostgreSQLConnection)
+
         return fetch_schema_v1(connection._conn, schema_name=schema_name)
 
     @kernel_function(
@@ -112,7 +111,7 @@ def get_db_context(
         connection = self._get_connection()
         if connection is None:
             return "Error. Could not connect to the database. No connection found."
-        assert isinstance(connection, PostgreSQLConnection)
+
         return fetch_schema(connection._conn)
 
     @kernel_function(
@@ -155,7 +154,7 @@ async def execute_sql_query_readonly_kernelfunc(
         connection = self._get_connection()
         if connection is None:
             return "Error. Could not connect to the database. No connection found."
-        assert isinstance(connection, PostgreSQLConnection)
+
         try:
             result = execute_readonly_query(connection._conn, query, self._max_result_chars)
         except Exception:
@@ -237,7 +236,7 @@ async def execute_sql_statement_kernelfunc(
         connection = self._get_connection()
         if connection is None:
             return "Error. Could not connect to the database. No connection found."
-        assert isinstance(connection, PostgreSQLConnection)
+
         try:
             result = execute_statement(connection._conn, statement)
         except Exception:

ossdbtoolsservice/connection/connection_service.py

@@ -9,7 +9,8 @@
 import contextlib
 import threading
 import uuid
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple  # noqa
+from collections.abc import Iterable
+from typing import Any, Callable, Optional
 
 from ossdbtoolsservice.connection.contracts import (
     BUILD_CONNECTION_INFO_REQUEST,

ossdbtoolsservice/driver/connection_manager.py

@@ -4,7 +4,7 @@
 # --------------------------------------------------------------------------------------------
 
 
-from ossdbtoolsservice.driver.types import PostgreSQLConnection, ServerConnection
+from ossdbtoolsservice.driver.types import ServerConnection
 from ossdbtoolsservice.workspace.contracts import Configuration
 
 
@@ -28,7 +28,7 @@ def _create_connection(
         :param options: a dict containing connection parameters
         """
 
-        return PostgreSQLConnection(options, config)
+        return ServerConnection(options, config)
 
-    def get_connection(self) -> PostgreSQLConnection:
+    def get_connection(self) -> ServerConnection:
         return self._conn_object

ossdbtoolsservice/driver/types/__init__.py

@@ -4,6 +4,5 @@
 # --------------------------------------------------------------------------------------------
 
 from ossdbtoolsservice.driver.types.driver import ServerConnection
-from ossdbtoolsservice.driver.types.psycopg_driver import PostgreSQLConnection
 
-__all__ = ["ServerConnection", "PostgreSQLConnection"]
+__all__ = ["ServerConnection"]