Support Executing SQL (#6)

* Add fixtures for pg connection

* Test that we return a postgres client

* Support instantiating postgres connection

* Install psycopg2

* Ignore local configs

* Properly configure bearer token

Author: meierj-wayscript

.gitignore

@@ -1,3 +1,6 @@
+# local configs
+env.*.local
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

Pipfile

@@ -5,6 +5,7 @@ name = "pypi"
 
 [packages]
 slack-sdk = "*"
+psycopg2 = "*"
 
 [dev-packages]
 pytest = "*"

Pipfile.lock

@@ -0,0 +1,44 @@
+import json
+
+import psycopg2
+
+from wayscript import utils
+from wayscript.errors import MissingCredentialsError
+
+
+def get_connection_kwargs(_id: str) -> dict:
+    """
+    Return postgres connection kwargs
+    
+    If you want to instantiate your own client, use this method.
+    """
+    wayscript_client = utils.WayScriptClient()
+    response = wayscript_client.get_workspace_integration_detail(_id)
+    response.raise_for_status()
+    workspace_integration_data = response.json()
+    credentials_str = workspace_integration_data.get("credentials")
+    credentials = {}
+
+    try:
+        credentials = json.loads(credentials_str)
+    except json.decoder.JSONDecodeError:
+        credentials = {}
+
+    kwargs = {
+        "dbname": credentials.get("database_name"),
+        "user": credentials.get("database_user"),
+        "password": credentials.get("database_password"),
+        "host": credentials.get("database_host"),
+        "port": credentials.get("database_port", 5432),
+    }
+
+    if not credentials or not all(v for v in kwargs.values()):
+        raise MissingCredentialsError(f"Missing credentials for workspace_integration={_id}")
+
+    return kwargs
+
+
+def get_client_for_workspace_integration(_id: str):
+    """Instantiate connection from workspace integration kwargs"""
+    kwargs = get_connection_kwargs(_id)
+    return psycopg2.connect(**kwargs)

src/wayscript/integrations/sql/__init__.py

@@ -23,7 +23,9 @@ class WayScriptClient:
     def __init__(self, *args, **kwargs):
         """Init a wayscript client"""
         self.session = requests.Session()
-        self.session.headers["authorization"] = get_process_execution_user_token()
+        access_token =  get_process_execution_user_token()
+        self.session.headers["authorization"] = f"Bearer {access_token}"
+        self.session.headers["content-type"] = "application/json"
 
     def _get_url(self, subpath: str, route: str, template_args: dict=None):
         """Generate an url"""

src/wayscript/integrations/sql/postgres.py

@@ -101,3 +101,30 @@ def workspace_integrations_detail_response():
         })
     }
     return data
+
+@pytest.fixture
+def workspace_integration_sql_credentials():
+    """Data for sql credentials"""
+    data = {
+        "database_name": "my_db",
+        "database_user": "zach",
+        "database_port": 15432,
+        "database_password": "very-secret-password",
+        "database_host": "host.docker.internal"
+    }
+    return data
+
+
+@pytest.fixture
+def workspace_integrations_detail_response_sql(workspace_integration_sql_credentials):
+    """
+    Data from GET /workspaces-integrations/<id>
+    
+    For type=sql
+    """
+    data = {
+        "id": WORKSPACE_INTEGRATION_ID,
+        "type": "sql",
+        "credentials": json.dumps(workspace_integration_sql_credentials),
+    }
+    return data

src/wayscript/utils.py

@@ -0,0 +1,31 @@
+from unittest import mock
+
+import psycopg2
+import responses
+
+from wayscript.integrations.sql import postgres
+
+
+@responses.activate
+def test_get_postgres_client_for_workspace_integration(
+    monkeypatch,
+    patch_client_get_url,
+    workspace_integration_sql_credentials,
+    workspace_integrations_detail_response_sql,
+):
+    """Test getting postgres client kwargs"""
+    connection = "my-connection"
+    monkeypatch.setattr(psycopg2, "connect", mock.Mock(return_value=connection))
+
+    responses.add(
+        responses.GET,
+        patch_client_get_url,
+        json=workspace_integrations_detail_response_sql,
+        status=200,
+    )
+
+    _id = workspace_integrations_detail_response_sql["id"]
+
+    assert connection == postgres.get_client_for_workspace_integration(_id)
+
+    psycopg2.connect.assert_called_with(dbname='my_db', user='zach', password='very-secret-password', host='host.docker.internal', port=15432)