port #38 to v3

agrc · Mar 13, 2023 · 5faab6b · 5faab6b
1 parent 4e11f5c
commit 5faab6b
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 5 deletions.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -4,11 +4,13 @@
     "authed",
     "auths",
     "ceildiv",
+    "cloudsql",
     "Codecov",
     "concatted",
     "Concatting",
     "Dataframe",
     "DataFrame",
+    "drivername",
     "Dtype",
     "featurelayer",
     "featureset",
@@ -26,6 +28,7 @@
     "Overwriter",
     "palletjack",
     "PARENTOBJECTID",
+    "PGSQL",
     "Polyline",
     "psycopg",
     "pylint",

diff --git a/setup.py b/setup.py
@@ -38,7 +38,8 @@
         'pygsheets==2.0.*',
         'geopandas==0.12.*',
         'SQLAlchemy==1.4.*',
-        'psycopg2==2.9.*',
+        'pg8000==1.29.*',
+        'psycopg2-binary==2.9.*',
         'numpy==1.23.*',  #: Pinned to fix "module 'numpy' has no attribute 'str'" error
     ],
     extras_require={

diff --git a/src/palletjack/extract.py b/src/palletjack/extract.py
@@ -15,7 +15,6 @@
 import requests
 import sqlalchemy
 from googleapiclient.http import MediaIoBaseDownload
-from sqlalchemy import create_engine
 
 from . import utils
 
@@ -440,9 +439,29 @@ class PostgresLoader:
 
     def __init__(self, host, database, username, password, port=5432):
         self._class_logger = logging.getLogger(__name__).getChild(self.__class__.__name__)
-
-        connection_string = f'postgresql://{username}:{password}@{host}:{port}/{database}'
-        self.engine = create_engine(connection_string)
+        if os.environ.get('FUNCTION_TARGET') is not None:  #: this is an env var specific to cloud functions
+            self._class_logger.info('running in GCF, using unix socket')
+            self.engine = sqlalchemy.create_engine(
+                sqlalchemy.engine.url.URL.create(
+                    drivername='postgresql+pg8000',
+                    username=username,
+                    password=password,
+                    database=database,
+                    query={'unix_sock': f'/cloudsql/{host}/.s.PGSQL.{port}'},  #: requires the pg8000 package
+                )
+            )
+        else:
+            self._class_logger.info('running locally, using traditional host connection')
+            self.engine = sqlalchemy.create_engine(
+                sqlalchemy.engine.url.URL.create(
+                    drivername='postgresql',
+                    username=username,
+                    password=password,
+                    database=database,
+                    host=host,
+                    port=port,
+                )
+            )
 
     def read_table_into_dataframe(self, table_name, index_column, crs, spatial_column):
         """Read a table into a dataframe