imports psycopg2 extensions

dlt-hub · Aug 14, 2024 · a39184e · a39184e
1 parent 84a4400
commit a39184e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sources/pg_replication/helpers.py b/sources/pg_replication/helpers.py
@@ -10,7 +10,7 @@
 from dataclasses import dataclass, field
 
 import psycopg2
-from psycopg2.extensions import cursor
+from psycopg2.extensions import cursor, connection
 from psycopg2.extras import (
     LogicalReplicationConnection,
     ReplicationCursor,
@@ -476,7 +476,7 @@ def _import_sql_table_resource() -> None:
 def _get_conn(
     credentials: ConnectionStringCredentials,
     connection_factory: Optional[Any] = None,
-) -> Union[psycopg2.extensions.connection, LogicalReplicationConnection]:
+) -> Union[connection, LogicalReplicationConnection]:
     """Returns a psycopg2 connection to interact with postgres."""
     return psycopg2.connect(  # type: ignore[no-any-return]
         database=credentials.database,