danallison · August 29, 2024 10:08 · May 16, 2018 · May 10, 2018 · May 4, 2018 · May 4, 2018
diff --git a/with_sql_session.py b/with_sql_session.py
@@ -1,6 +1,7 @@
 from sshtunnel import SSHTunnelForwarder
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
+from functools import wraps
 
 # secrets.py contains credentials, etc.
 import secrets
@@ -14,21 +15,21 @@ def get_engine_for_port(port):
         db=secrets.db
     ))
 
-def with_sql_session(function, engine=None):
+def with_sql_session(function, args, kwargs, engine=None):
     if engine is None:
         # Default to local port
         engine = get_engine_for_port(5432)
     Session = sessionmaker(bind=engine)
     session = Session()
     try:
-        return function(session)
+        return function(session, *args, **kwargs)
     finally:
         session.close()
 
-def with_local_sql_session(function):
-    return with_sql_session(function)
+def with_local_sql_session(function, *args, **kwargs):
+    return with_sql_session(function, args, kwargs)
 
-def with_remote_sql_session(function):
+def with_remote_sql_session(function, *args, **kwargs):
     # Hat tip: https://stackoverflow.com/a/38001815
     with SSHTunnelForwarder(
             (secrets.server_ip_address, 22),
@@ -39,4 +40,17 @@ def with_remote_sql_session(function):
         ) as tunnel:
         tunnel.start()
         engine = get_engine_for_port(tunnel.local_bind_port)
-        return with_sql_session(function, engine=engine)
+        return with_sql_session(function, args, kwargs, engine=engine)
+
+# Decorators
+def local_sql_session(function):
+    @wraps(function)
+    def wrapper(*args, **kwargs):
+        return with_local_sql_session(function, *args, **kwargs)
+    return wrapper
+
+def remote_sql_session(function):
+    @wraps(function)
+    def wrapper(*args, **kwargs):
+        return with_remote_sql_session(function, *args, **kwargs)
+    return wrapper
diff --git a/with_remote_sql_session.py b/with_remote_sql_session.py
@@ -1,33 +0,0 @@
-# Hat tip to https://stackoverflow.com/a/38001815
-from sshtunnel import SSHTunnelForwarder
-from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker
-
-# secrets.py contains credentials, etc.
-import secrets
-
-def with_remote_sql_session(function):
-    with SSHTunnelForwarder(
-            (secrets.server_ip_address, 22),
-            ssh_username=secrets.ssh_username,
-            ssh_pkey=secrets.ssh_private_key_path,
-            ssh_private_key_password=secrets.ssh_private_key_password,
-            remote_bind_address=('127.0.0.1', 5432)
-        ) as tunnel:
-
-        tunnel.start()
-
-        engine = create_engine('postgresql://{user}:{password}@{host}:{port}/{db}'.format(
-            user=secrets.pg_user,
-            password=secrets.pg_password,
-            host='127.0.0.1',
-            port=tunnel.local_bind_port,
-            db='my_database'
-        ))
-
-        Session = sessionmaker(bind=engine)
-        session = Session()
-        try:
-            return function(session)
-        finally:
-            session.close()

diff --git a/with_sql_session.py b/with_sql_session.py
@@ -0,0 +1,42 @@
+from sshtunnel import SSHTunnelForwarder
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+# secrets.py contains credentials, etc.
+import secrets
+
+def get_engine_for_port(port):
+    return create_engine('postgresql://{user}:{password}@{host}:{port}/{db}'.format(
+        user=secrets.pg_user,
+        password=secrets.pg_password,
+        host='127.0.0.1',
+        port=port,
+        db=secrets.db
+    ))
+
+def with_sql_session(function, engine=None):
+    if engine is None:
+        # Default to local port
+        engine = get_engine_for_port(5432)
+    Session = sessionmaker(bind=engine)
+    session = Session()
+    try:
+        return function(session)
+    finally:
+        session.close()
+
+def with_local_sql_session(function):
+    return with_sql_session(function)
+
+def with_remote_sql_session(function):
+    # Hat tip: https://stackoverflow.com/a/38001815
+    with SSHTunnelForwarder(
+            (secrets.server_ip_address, 22),
+            ssh_username=secrets.ssh_username,
+            ssh_pkey=secrets.ssh_private_key_path,
+            ssh_private_key_password=secrets.ssh_private_key_password,
+            remote_bind_address=('127.0.0.1', 5432)
+        ) as tunnel:
+        tunnel.start()
+        engine = get_engine_for_port(tunnel.local_bind_port)
+        return with_sql_session(function, engine=engine)
diff --git a/with_remote_sql_session.py b/with_remote_sql_session.py
@@ -1,5 +1,6 @@
 # Hat tip to https://stackoverflow.com/a/38001815
 from sshtunnel import SSHTunnelForwarder
+from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
 # secrets.py contains credentials, etc.

diff --git a/with_remote_sql_session.py b/with_remote_sql_session.py
@@ -12,24 +12,21 @@ def with_remote_sql_session(function):
             ssh_pkey=secrets.ssh_private_key_path,
             ssh_private_key_password=secrets.ssh_private_key_password,
             remote_bind_address=('127.0.0.1', 5432)
-        ) as server:
+        ) as tunnel:
 
-        server.start() #start ssh sever
+        tunnel.start()
 
-        #connect to PostgreSQL
-        local_port = str(server.local_bind_port)
         engine = create_engine('postgresql://{user}:{password}@{host}:{port}/{db}'.format(
             user=secrets.pg_user,
             password=secrets.pg_password,
-            host='localhost',
-            port=local_port,
+            host='127.0.0.1',
+            port=tunnel.local_bind_port,
             db='my_database'
         ))
 
         Session = sessionmaker(bind=engine)
         session = Session()
         try:
-            return_val = function(session)
+            return function(session)
         finally:
             session.close()
-            return return_val
diff --git a/with_remote_sql_session.py b/with_remote_sql_session.py
@@ -0,0 +1,35 @@
+# Hat tip to https://stackoverflow.com/a/38001815
+from sshtunnel import SSHTunnelForwarder
+from sqlalchemy.orm import sessionmaker
+
+# secrets.py contains credentials, etc.
+import secrets
+
+def with_remote_sql_session(function):
+    with SSHTunnelForwarder(
+            (secrets.server_ip_address, 22),
+            ssh_username=secrets.ssh_username,
+            ssh_pkey=secrets.ssh_private_key_path,
+            ssh_private_key_password=secrets.ssh_private_key_password,
+            remote_bind_address=('127.0.0.1', 5432)
+        ) as server:
+
+        server.start() #start ssh sever
+
+        #connect to PostgreSQL
+        local_port = str(server.local_bind_port)
+        engine = create_engine('postgresql://{user}:{password}@{host}:{port}/{db}'.format(
+            user=secrets.pg_user,
+            password=secrets.pg_password,
+            host='localhost',
+            port=local_port,
+            db='my_database'
+        ))
+
+        Session = sessionmaker(bind=engine)
+        session = Session()
+        try:
+            return_val = function(session)
+        finally:
+            session.close()
+            return return_val
No results found