Initial Spec For Credential Management and SQLAlchemy Database Connectors

pyproject.toml

@@ -18,7 +18,10 @@ dependencies = [
   "databricks-labs-blueprint[yaml]>=0.2.3",
   "databricks-labs-lsql>=0.7.5,<0.14.0",  # TODO: Limit the LSQL version until dependencies are correct.
   "cryptography>=41.0.3",
+  "pyodbc",
+  "SQLAlchemy",
   "pygls>=2.0.0a2",
+
 ]
 
 [project.urls]

src/databricks/labs/remorph/connections/credential_manager.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+import yaml
+from databricks.labs.blueprint.wheels import ProductInfo
+
+
+class Credentials:
+    def __init__(self, product_info: ProductInfo) -> None:
+        self._product_info = product_info
+        self._credentials: dict[str, str] = self._load_credentials(self._get_local_version_file_path())
+
+    def _get_local_version_file_path(self) -> Path:
+        user_home = f"{Path(__file__).home()}"
+        return Path(f"{user_home}/.databricks/labs/{self._product_info.product_name()}/credentials.yml")
+
+    def _load_credentials(self, file_path: Path) -> dict[str, str]:
+        with open(file_path, encoding="utf-8") as f:
+            return yaml.safe_load(f)
+
+    def get(self, source: str) -> dict[str, str]:
+        error_msg = f"source system: {source} credentials not found not in file credentials.yml"
+        if source in self._credentials:
+            value = self._credentials[source]
+            if isinstance(value, dict):
+                return value
+            raise KeyError(error_msg)
+        raise KeyError(error_msg)

src/databricks/labs/remorph/connections/database_manager.py

@@ -0,0 +1,64 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from sqlalchemy import create_engine
+from sqlalchemy.engine import Engine, Result
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy import text
+
+
+class _ISourceSystemConnector(ABC):
+    @abstractmethod
+    def _connect(self) -> Engine:
+        pass
+
+    @abstractmethod
+    def execute_query(self, query: str) -> Result[Any]:
+        pass
+
+
+class _BaseConnector(_ISourceSystemConnector):
+    def __init__(self, config: dict[str, Any]):
+        self.config = config
+        self.engine: Engine = self._connect()
+
+    def _connect(self) -> Engine:
+        raise NotImplementedError("Subclasses should implement this method")
+
+    def execute_query(self, query: str) -> Result[Any]:
+        if not self.engine:
+            raise ConnectionError("Not connected to the database.")
+        session = sessionmaker(bind=self.engine)
+        connection = session()
+        return connection.execute(text(query))
+
+
+class SnowflakeConnector(_BaseConnector):
+    def _connect(self) -> Engine:
+        connection_string = (
+            f"snowflake://{self.config['user']}:{self.config['password']}@{self.config['account']}/"
+            f"{self.config['database']}/{self.config['schema']}?warehouse={self.config['warehouse']}&role={self.config['role']}"
+        )
+        self.engine = create_engine(connection_string)
+        return self.engine
+
+
+class MSSQLConnector(_BaseConnector):
+    def _connect(self) -> Engine:
+        connection_string = (
+            f"mssql+pyodbc://{self.config['user']}:{self.config['password']}@{self.config['server']}/"
+            f"{self.config['database']}?driver={self.config['driver']}"
+        )
+        self.engine = create_engine(connection_string)
+        return self.engine
+
+
+class SourceSystemConnectorFactory:
+    @staticmethod
+    def create_connector(db_type: str, config: dict[str, str]) -> _ISourceSystemConnector:
+        if db_type == "snowflake":
+            return SnowflakeConnector(config)
+        if db_type == "mssql":
+            return MSSQLConnector(config)
+
+        raise ValueError(f"Unsupported database type: {db_type}")

src/databricks/labs/remorph/resources/config/credentials.yml

@@ -0,0 +1,28 @@
+snowflake:
+  account: example_account
+  connect_retries: 1
+  connect_timeout: null
+  host: null
+  insecure_mode: false
+  oauth_client_id: null
+  oauth_client_secret: null
+  password: null
+  port: null
+  private_key: null
+  private_key_passphrase: null
+  private_key_path: null
+  role: null
+  token: null
+  user: null
+  warehouse: null
+
+msssql:
+  database: example_database
+  driver: ODBC Driver 18 for SQL Server
+  server: example_host
+  port: null
+  user: null
+  password: null
+
+
+