From ca691ae843a1911bb0d6ed6e1234df4f88d6edba Mon Sep 17 00:00:00 2001
From: Daniel Falbel <dfalbel@gmail.com>
Date: Fri, 21 Feb 2025 13:04:51 -0300
Subject: [PATCH] Add PostgreSQL for Python

---
 .../positron-connections/src/drivers.ts       | 71 +++++++++++++++++++
 .../posit/positron/connections.py             |  2 +-
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/extensions/positron-connections/src/drivers.ts b/extensions/positron-connections/src/drivers.ts
index e0dce4b83d8..4db81ff7a22 100644
--- a/extensions/positron-connections/src/drivers.ts
+++ b/extensions/positron-connections/src/drivers.ts
@@ -15,6 +15,7 @@ export function registerConnectionDrivers(context: vscode.ExtensionContext) {
 		new RPostgreSQLDriver(context),
 		new PythonSQLiteDriver(context),
 		new RSparkDriver(context),
+		new PythonPostgreSQLDriver(context),
 	];
 
 	for (const driver of drivers) {
@@ -351,3 +352,73 @@ conn = sqlite3.connect(${JSON.stringify(dbname) ?? JSON.stringify('')})
 `;
 	}
 }
+
+class PythonPostgreSQLDriver extends PythonDriver implements positron.ConnectionsDriver {
+
+	constructor(context: vscode.ExtensionContext) {
+		super();
+		const iconPath = path.join(context.extensionPath, 'media', 'logo', 'postgre.svg');
+		const iconData = readFileSync(iconPath, 'base64');
+		this.metadata.base64EncodedIconSvg = iconData;
+	}
+
+	driverId: string = 'py-postgres';
+	metadata: positron.ConnectionsDriverMetadata = {
+		languageId: 'python',
+		name: 'PostgresSQL',
+		inputs: [
+			{
+				'id': 'dbname',
+				'label': 'Database Name',
+				'type': 'string',
+				'value': 'localhost'
+			},
+			{
+				'id': 'host',
+				'label': 'Host',
+				'type': 'string',
+				'value': 'localhost'
+			},
+			{
+				'id': 'port',
+				'label': 'Port',
+				'type': 'number',
+				'value': '5432'
+			},
+			{
+				'id': 'user',
+				'label': 'User',
+				'type': 'string',
+				'value': 'postgres'
+			},
+			{
+				'id': 'password',
+				'label': 'Password',
+				'type': 'string',
+				'value': 'password'
+			},
+		]
+	};
+
+	generateCode(inputs: positron.ConnectionsInput[]) {
+		const dbname = inputs.find(input => input.id === 'dbname')?.value;
+		const host = inputs.find(input => input.id === 'host')?.value;
+		const port = inputs.find(input => input.id === 'port')?.value;
+		const user = inputs.find(input => input.id === 'user')?.value;
+		const password = inputs.find(input => input.id === 'password')?.value;
+
+		const connection_string = `postgresql+psycopg2://${user}:${password}@${host}:${port}/${dbname}`;
+
+		return `import sqlalchemy
+conn = sqlalchemy.create_engine(sqlalchemy.URL.create(
+	"postgresql+psycopg2",
+	username=${JSON.stringify(user)},
+	password=${JSON.stringify(password)},
+	host=${JSON.stringify(host)},
+	database=${JSON.stringify(dbname)},
+	port=${JSON.stringify(port)}
+))
+%connection_show conn
+`;
+	}
+}
diff --git a/extensions/positron-python/python_files/posit/positron/connections.py b/extensions/positron-python/python_files/posit/positron/connections.py
index 3ff0e7d0f66..d9b77fd0e04 100644
--- a/extensions/positron-python/python_files/posit/positron/connections.py
+++ b/extensions/positron-python/python_files/posit/positron/connections.py
@@ -622,7 +622,7 @@ class SQLAlchemyConnection(Connection):
     def __init__(self, conn: sqlalchemy.Engine):
         self.conn = conn
         self.display_name = f"SQLAlchemy ({conn.name})"
-        self.host = conn.url.render_as_string()
+        self.host = conn.url.render_as_string(hide_password=False)
         self.type = "SQLAlchemy"
         self.code = (
             "import sqlalchemy\n"