diff --git a/README.md b/README.md index 11caab8..0699fa5 100644 --- a/README.md +++ b/README.md @@ -517,6 +517,21 @@ df = fetch_table( ) ``` +#### Amazon Redshift + +```python +from src.connectors.redshift import fetch_table + +df = fetch_table( + host="my-cluster.xxxxx.eu-west-1.redshift.amazonaws.com", + port=5439, + database="dev", + user="awsuser", + password="my_password", + table="my_table" +) +``` + --- ### Connector Status @@ -528,8 +543,8 @@ df = fetch_table( | PostgreSQL | User/Pass | ✅ | ✅ | Stable | | MySQL | User/Pass | ✅ | ✅ | Stable | | BigQuery | Service Account JSON | ✅ | ✅ | Stable | +| Amazon Redshift | User/Pass | ✅ | ✅ | Stable | | MongoDB | — | 🔜 | 🔜 | Planned | -| Redshift | — | 🔜 | 🔜 | Planned | | DuckDB | — | 🔜 | 🔜 | Planned | | Microsoft Fabric | — | 🔜 | 🔜 | Planned | @@ -694,7 +709,6 @@ We want to support every major database. Next targets: | Database | Difficulty | Issue | |----------|-----------|-------| | MongoDB | Medium | #1 | -| Redshift | Easy | #2 | | DuckDB | Easy | #3 | | Microsoft Fabric | Medium | #4 | | Elasticsearch | Hard | #5 | @@ -883,7 +897,7 @@ tests/test_pipeline.py::TestCSVLoading::test_demo_csv_has_rows PASSED - [x] JSON export - [ ] pip package — `pip install multi-agent-data-pipeline` - [ ] MongoDB connector -- [ ] Redshift connector +- [x] Redshift connector - [ ] DuckDB connector - [ ] Microsoft Fabric connector - [ ] Async parallel agent execution diff --git a/app.py b/app.py index c6e46aa..5e528fd 100644 --- a/app.py +++ b/app.py @@ -630,7 +630,7 @@ def run_pipeline_ui(df): db_type = st.selectbox( "Database", - ["Azure Databricks", "Snowflake", "PostgreSQL", "MySQL", "BigQuery"], + ["Azure Databricks", "Snowflake", "PostgreSQL", "MySQL", "BigQuery", "Amazon Redshift"], label_visibility="collapsed" ) @@ -756,6 +756,30 @@ def run_pipeline_ui(df): else: st.warning("Please fill all fields") + elif db_type == "Amazon Redshift": + col1, col2 = st.columns(2) + with col1: + host = st.text_input("Host", placeholder="my-cluster.xxxxx.eu-west-1.redshift.amazonaws.com") + database = st.text_input("Database", placeholder="dev") + table = st.text_input("Table", placeholder="my_table") + with col2: + port = st.text_input("Port", value="5439") + user = st.text_input("Username", placeholder="awsuser") + password = st.text_input("Password", type="password") + + if st.button("🔌 Connect & Fetch Table"): + if host and database and user and password and table: + try: + from src.connectors.redshift import fetch_table + with st.spinner("Connecting to Redshift..."): + df = fetch_table(host, int(port), database, user, password, table) + st.success(f"Connected — {len(df)} rows fetched") + st.dataframe(df, use_container_width=True, height=240) + except Exception as e: + st.error(f"Connection failed: {e}") + else: + st.warning("Please fill all fields") + st.markdown('', unsafe_allow_html=True) if df is not None and not df.empty: diff --git a/src/connectors/redshift.py b/src/connectors/redshift.py new file mode 100644 index 0000000..8871bc1 --- /dev/null +++ b/src/connectors/redshift.py @@ -0,0 +1,44 @@ +import psycopg2 +import pandas as pd + +def connect(host: str, port: int, database: str, user: str, password: str): + conn = psycopg2.connect( + host=host, + port=port, + database=database, + user=user, + password=password + ) + return conn + +def list_tables(host: str, port: int, database: str, user: str, password: str) -> list: + conn = connect(host, port, database, user, password) + cursor = conn.cursor() + cursor.execute(""" + SELECT table_name + FROM information_schema.tables + WHERE table_schema = 'public' + ORDER BY table_name + """) + tables = [row[0] for row in cursor.fetchall()] + cursor.close() + conn.close() + return tables + +def fetch_table(host: str, port: int, database: str, user: str, password: str, table: str, limit: int = 1000) -> pd.DataFrame: + conn = connect(host, port, database, user, password) + cursor = conn.cursor() + + try: + # Validate table exists to prevent SQL injection + valid_tables = list_tables(host, port, database, user, password) + if table not in valid_tables: + raise ValueError(f"Table '{table}' not found in database. Available tables: {', '.join(valid_tables)}") + + cursor.execute(f"SELECT * FROM {table} LIMIT {limit}") + columns = [desc[0] for desc in cursor.description] + rows = cursor.fetchall() + return pd.DataFrame(rows, columns=columns) + finally: + cursor.close() + conn.close() diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index fb1e38e..4d72257 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -128,6 +128,75 @@ def test_demo_csv_has_rows(self): assert len(df) > 0, "Demo CSV is empty" +class FakeRedshiftCursor: + + def __init__(self): + self.description = None + self._rows = [] + self.executed_query = None + + def execute(self, query): + self.executed_query = query + if "information_schema.tables" in query: + self._rows = [("customers",), ("transactions",)] + self.description = [("table_name",)] + else: + self._rows = [ + ("TXN001", "Product A", 10.00), + ("TXN002", "Product B", 15.00), + ] + self.description = [("transaction_id",), ("product_name",), ("unit_price",)] + + def fetchall(self): + return self._rows + + def close(self): + pass + + +class FakeRedshiftConnection: + + def __init__(self): + self._cursor = FakeRedshiftCursor() + + def cursor(self): + return self._cursor + + def close(self): + pass + + +class TestRedshiftConnector: + + @pytest.fixture + def fake_connection(self, monkeypatch): + from src.connectors import redshift + conn = FakeRedshiftConnection() + monkeypatch.setattr(redshift.psycopg2, "connect", lambda **kwargs: conn) + return conn + + def test_list_tables(self, fake_connection): + from src.connectors.redshift import list_tables + tables = list_tables("host", 5439, "dev", "awsuser", "password") + assert tables == ["customers", "transactions"] + + def test_fetch_table(self, fake_connection): + from src.connectors.redshift import fetch_table + df = fetch_table("host", 5439, "dev", "awsuser", "password", "transactions") + assert len(df) == 2 + assert list(df.columns) == ["transaction_id", "product_name", "unit_price"] + + def test_fetch_table_applies_limit(self, fake_connection): + from src.connectors.redshift import fetch_table + fetch_table("host", 5439, "dev", "awsuser", "password", "transactions", limit=500) + assert "LIMIT 500" in fake_connection.cursor().executed_query + + def test_fetch_table_validates_table_name(self, fake_connection): + from src.connectors.redshift import fetch_table + with pytest.raises(ValueError, match="not found"): + fetch_table("host", 5439, "dev", "awsuser", "password", "nonexistent_table") + + class TestPDFExists: def test_demo_pdf_exists(self):