Optional DB drivers (#15)

* Optional DB drivers

Author: lloesche

README.md

@@ -15,9 +15,11 @@ source venv/bin/activate
 Once the virtual environment is activated, you can install cloud2sql:
 
 ```bash
-pip install cloud2sql
+pip install cloud2sql[all]
 ```
 
+If you only require support for a specific database, instead of `cloud2sql[all]` you can choose between `cloud2sql[snowflake]`, `cloud2sql[parquet]`, `cloud2sql[postgresql]`, `cloud2sql[mysql]`.
+
 ## Usage
 
 The sources and destinations for `cloud2sql` are configured via a configuration file. Create your own configuration by adjusting the [config template file](./config-template.yaml).

cloud2sql/__main__.py

@@ -7,7 +7,7 @@
 
 from cloud2sql.analytics import PosthogEventSender, NoEventSender, AnalyticsEventSender
 from cloud2sql.collect_plugins import collect_from_plugins, configure
-from cloud2sql.util import db_string_from_config
+from cloud2sql.util import db_string_from_config, check_parquet_driver
 
 # Will fail in case snowflake is not installed - which is fine.
 try:
@@ -58,8 +58,11 @@ def main() -> None:
         setup_logger("resoto.cloud2sql", level=args.log_level, force=True)
         sender = NoEventSender() if args.analytics_opt_out else PosthogEventSender()
         config = configure(args.config)
-        is_parquet = next(iter(config["destinations"].keys()), None) == "parquet"
-        engine = None if is_parquet else create_engine(db_string_from_config(config))
+        engine = None
+        if next(iter(config["destinations"].keys()), None) == "parquet":
+            check_parquet_driver()
+        else:
+            engine = create_engine(db_string_from_config(config))
         collect(engine, args, sender)
     except Exception as e:
         if args.debug:  # raise exception and show complete tracelog

cloud2sql/collect_plugins.py

@@ -31,7 +31,11 @@
 from cloud2sql.analytics import AnalyticsEventSender
 from cloud2sql.show_progress import CollectInfo
 from cloud2sql.sql import SqlUpdater, sql_updater
-from cloud2sql.parquet import ParquetModel, ParquetWriter
+
+try:
+    from cloud2sql.parquet import ParquetModel, ParquetWriter
+except ImportError:
+    pass
 
 
 log = getLogger("resoto.cloud2sql")

cloud2sql/util.py

@@ -1,6 +1,7 @@
 from typing import Union, List, Optional, Any
 
 from resotolib.types import JsonElement, Json
+from sqlalchemy.engine import create_engine
 
 
 def value_in_path(element: JsonElement, path_or_name: Union[List[str], str]) -> Optional[Any]:
@@ -26,6 +27,7 @@ def db_string_from_config(config: Json) -> str:
 
     db_type = list(destinations.keys())[0]
     db_config = destinations[db_type]
+    db_type = update_db_type(db_type)
     user = db_config.get("user")
     password = db_config.get("password")
     host = db_config.get("host")
@@ -52,4 +54,44 @@ def db_string_from_config(config: Json) -> str:
     if len(args) > 0:
         db_uri += "?" + "&".join([f"{k}={v}" for k, v in args.items()])
 
+    check_db_driver(db_uri)
+
     return db_uri
+
+
+def update_db_type(db_type: str) -> str:
+    if db_type == "mysql":
+        db_type = "mysql+pymysql"
+    elif db_type == "mariadb":
+        db_type = "mariadb+pymysql"
+    return db_type
+
+
+def check_db_driver(db_uri: str) -> None:
+    try:
+        create_engine(db_uri)
+    except ModuleNotFoundError:
+        err = "The database type you configured is not installed. "
+        if db_uri.startswith("mysql") or db_uri.startswith("mariadb"):
+            err += "Please run 'pip install cloud2sql[mysql]' and try again."
+        elif db_uri.startswith("postgresql"):
+            err += "Please run 'pip install cloud2sql[postgresql]' and try again."
+        elif db_uri.startswith("snowflake"):
+            err += "Please run 'pip install cloud2sql[snowflake]' and try again."
+        elif db_uri.startswith("mssql"):
+            err += "Please install the pymssql package and try again."
+        elif db_uri.startswith("oracle"):
+            err += "Please install the cx_oracle package and try again."
+        else:
+            err += "Please install the required dependencies and try again."
+        raise ModuleNotFoundError(err)
+
+
+def check_parquet_driver() -> None:
+    try:
+        import pyarrow  # noqa: F401
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "The parquet format you configured is not installed. "
+            "Please run 'pip install cloud2sql[parquet]' and try again."
+        )

requirements-mysql.txt

@@ -0,0 +1 @@
+pymysql>=1.0.2

requirements-parquet.txt

@@ -0,0 +1 @@
+pyarrow==10.0.1

requirements-postgresql.txt

@@ -0,0 +1 @@
+psycopg2-binary>=2.9.5

requirements-snowflake.txt

@@ -0,0 +1 @@
+snowflake-sqlalchemy>=1.4.5

requirements.txt

@@ -8,18 +8,6 @@ resotoclient>=1.2.1
 posthog>=2.2.0
 requests>=2.28.1
 
-# parquet support
-pyarrow==10.0.1
-
-# bundle popular db drivers
-psycopg2-binary>=2.9.5 # postgres
-pymysql>=1.0.2 # mysql + mariadb
-
-# Install snowflake connector
-# Note: this version does not support Python 3.11
-# See: https://github.com/snowflakedb/snowflake-connector-python/pull/1349 for reference
-snowflake-sqlalchemy>=1.4.5
-
 resotolib>=3.0.0
 # all collector plugins
 resoto-plugin-aws>=3.0.0

setup.py

@@ -9,6 +9,18 @@
 with open("requirements.txt") as f:
     required = f.read().splitlines()
 
+with open("requirements-mysql.txt") as f:
+    required_mysql = f.read().splitlines()
+
+with open("requirements-postgresql.txt") as f:
+    required_postgresql = f.read().splitlines()
+
+with open("requirements-parquet.txt") as f:
+    required_parquet = f.read().splitlines()
+
+with open("requirements-snowflake.txt") as f:
+    required_snowflake = f.read().splitlines()
+
 with open("requirements-test.txt") as f:
     test_required = f.read().splitlines()
 
@@ -28,6 +40,14 @@
     classifiers=["Programming Language :: Python :: 3"],
     entry_points={"console_scripts": ["cloud2sql=cloud2sql.__main__:main"]},
     install_requires=required,
+    extras_require={
+        "all": required_mysql + required_postgresql + required_snowflake + required_parquet,
+        "mysql": required_mysql,
+        "mariadb": required_mysql,
+        "postgresql": required_postgresql,
+        "snowflake": required_snowflake,
+        "parquet": required_parquet,
+    },
     license="Apache Software License 2.0",
     long_description=readme,
     long_description_content_type="text/markdown",