added support for cmd specific env vars and auth

Author: jeffreyaven

.github/workflows/async_server_tests.yaml.disabled

@@ -46,10 +46,10 @@ jobs:
           pip install -r requirements.txt
 
       # Windows
-      - name: Install psycopg2-binary for Windows
+      - name: Install psycopg
         if: matrix.os == 'windows-latest'
         run: |
-          pip install psycopg
+          pip install psycopg[binary]
         shell: powershell
       # End Windows
 
@@ -58,16 +58,14 @@ jobs:
         if: matrix.os == 'ubuntu-latest'
         run: |
           sudo apt-get update
-          sudo apt-get install -y libpq-dev
-          pip install psycopg2
+          pip install psycopg
       # End Linux
 
       # macOS
       - name: Install PostgreSQL dev libraries on macOS
         if: matrix.os == 'macos-latest'
         run: |
-          brew install postgresql@14
-          pip install psycopg2
+          pip install psycopg
       # End macOS
 
       - name: Install pystackql
@@ -92,15 +90,15 @@ jobs:
         shell: bash
         if: matrix.os != 'windows-latest'
 
-      - name: Debug Shell Information
-        if: matrix.os == 'windows-latest'
-        shell: pwsh
-        run: |
-          Write-Host "Checking PowerShell version and environment variables:"
-          $PSVersionTable
-          Write-Host "Environment variables:"
-          Get-ChildItem Env:
-          Write-Host "Current Shell: $SHELL"
+      # - name: Debug Shell Information
+      #   if: matrix.os == 'windows-latest'
+      #   shell: pwsh
+      #   run: |
+      #     Write-Host "Checking PowerShell version and environment variables:"
+      #     $PSVersionTable
+      #     Write-Host "Environment variables:"
+      #     Get-ChildItem Env:
+      #     Write-Host "Current Shell: $SHELL"
 
       - name: Run tests on Windows
         env:

.github/workflows/test.yaml

@@ -166,7 +166,7 @@ Before testing, ensure you have all the required packages installed:
 ::
 
     pip install -r requirements.txt
-    pip install psycopg2-binary
+    pip install psycopg
 
 Once the dependencies are installed, you can run the tests using the provided script:
 

README.rst

@@ -1,5 +1,4 @@
 sphinx_rtd_theme
-psycopg2
 pandas
 requests
 IPython

docs/requirements.txt

@@ -79,30 +79,6 @@ def _download_file(url, path, showprogress=True):
 		print("ERROR: [_download_file] %s" % (str(e)))
 		exit(1)
 
-# def _setup(download_dir, platform, showprogress=False):
-#     try:
-#         print('installing stackql...')
-#         binary_name = _get_binary_name(platform)
-#         url = _get_url()
-#         print("downloading latest version of stackql from %s to %s" % (url, download_dir))
-#         archive_file_name = os.path.join(download_dir, os.path.basename(url))
-#         _download_file(url, archive_file_name, showprogress)
-#         # if Platform is starting with Darwin, then it is a MacOS
-#         if platform.startswith('Darwin'):
-#             unpacked_file_name = os.path.join(download_dir, 'stackql')
-#             command = 'pkgutil --expand-full {} {}'.format(archive_file_name, unpacked_file_name)
-#             # if there are files in unpacked_file_name, then remove them
-#             if os.path.exists(unpacked_file_name):
-#                 os.system('rm -rf {}'.format(unpacked_file_name))
-#             os.system(command)
-#         else:
-#             with zipfile.ZipFile(archive_file_name, 'r') as zip_ref:
-#                 zip_ref.extractall(download_dir) 
-#         os.chmod(os.path.join(download_dir, binary_name), 0o755)
-#     except Exception as e:
-#         print("ERROR: [_setup] %s" % (str(e)))
-#         exit(1)
-
 def _setup(download_dir, platform, showprogress=False):
     try:
         print('installing stackql...')

pystackql/_util.py

@@ -126,51 +126,51 @@ def _connect_to_server(self):
 
 		:return: Connection object if successful, or `None` if an error occurred.
 		:rtype: Connection or None
-		...
-		:raises `psycopg2.OperationalError`: Failed to connect to the server.
+		:raises `psycopg.OperationalError`: Failed to connect to the server.
 		"""
 		try:
-			conn = psycopg2.connect(
+			conn = psycopg.connect(
 				dbname='stackql',
 				user='stackql',
 				host=self.server_address,
-				port=self.server_port
+				port=self.server_port,
+				autocommit=True,
+				row_factory=dict_row  # Use dict_row to get rows as dictionaries
 			)
 			return conn
-		except psycopg2.OperationalError as oe:
+		except psycopg.OperationalError as oe:
 			print(f"OperationalError while connecting to the server: {oe}")
 		except Exception as e:
-			# Catching all other possible psycopg2 exceptions (and possibly other unexpected exceptions).
-			# You might want to log this or handle it differently in a real-world scenario.
 			print(f"Unexpected error while connecting to the server: {e}")
 		return None
 
 	def _run_server_query(self, query, is_statement=False):
-		"""Runs a query against the server using psycopg2.
-		
-		:param query: SQL query to be executed on the server.
-		:type query: str
-		:return: List of result rows if the query fetches results; empty list if there are no results.
-		:rtype: list of dict objects
-		:raises: psycopg2.ProgrammingError for issues related to the SQL query, 
-				unless the error is "no results to fetch", in which case an empty list is returned.
-		"""
+		"""Run a query against the server using the existing connection in server mode."""
+		if not self._conn:
+			raise ConnectionError("No active connection found. Ensure _connect_to_server is called.")
+
 		try:
-			cur = self._conn.cursor(cursor_factory=RealDictCursor)
-			cur.execute(query)
-			if is_statement:
-				# If the query is a statement, there are no results to fetch.
-				result_msg = cur.statusmessage
-				cur.close()
-				return [{'message': result_msg}]
-			rows = cur.fetchall()
-			cur.close()
-			return rows
-		except psycopg2.ProgrammingError as e:
-			if str(e) == "no results to fetch":
-				return []
-			else:
-				raise
+			with self._conn.cursor() as cur:
+				cur.execute(query)
+				if is_statement:
+					# Return status message for non-SELECT statements
+					result_msg = cur.statusmessage
+					return [{'message': result_msg}]
+				try:
+					# Fetch results for SELECT queries
+					rows = cur.fetchall()
+					return rows
+				except psycopg.ProgrammingError as e:
+					# Handle cases with no results
+					if "no results to fetch" in str(e):
+						return []
+					else:
+						raise
+		except psycopg.OperationalError as oe:
+			print(f"OperationalError during query execution: {oe}")
+		except Exception as e:
+			print(f"Unexpected error during query execution: {e}")
+
 
 	def _run_query(self, query, custom_auth=None, env_vars=None):
 		"""Internal method to execute a StackQL query using a subprocess.
@@ -330,13 +330,13 @@ def __init__(self,
 
 		if self.server_mode:
 			# server mode, connect to a server via the postgres wire protocol
-			# Attempt to import psycopg2 only if server_mode is True
-			global psycopg2, RealDictCursor
+			# Attempt to import psycopg only if server_mode is True
+			global psycopg, dict_row
 			try:
-				import psycopg2
-				from psycopg2.extras import RealDictCursor
+				import psycopg
+				from psycopg.rows import dict_row  # For returning results as dictionaries
 			except ImportError:
-				raise ImportError("psycopg2 is required in server mode but is not installed. Please install psycopg2 and try again.")
+				raise ImportError("psycopg is required in server mode but is not installed. Please install psycopg and try again.")
 
 			self.server_address = server_address
 			self.server_port = server_port
@@ -670,34 +670,30 @@ def execute(self, query, suppress_errors=True, custom_auth=None, env_vars=None):
 	def _run_server_query_with_new_connection(self, query):
 		"""Run a query against a StackQL postgres wire protocol server with a new connection.
 		"""
-		conn = None
 		try:
 			# Establish a new connection using credentials and configurations
-			conn = psycopg2.connect(
+			with psycopg.connect(
 				dbname='stackql',
 				user='stackql',
 				host=self.server_address,
-				port=self.server_port
-			)
-			# Create a new cursor and execute the query
-			with conn.cursor(cursor_factory=RealDictCursor) as cur:
-				cur.execute(query)
-				try:
-					rows = cur.fetchall()
-				except psycopg2.ProgrammingError as e:
-					if str(e) == "no results to fetch":
-						rows = []
-					else:
-						raise
-				return rows
-		except psycopg2.OperationalError as oe:
+				port=self.server_port,
+				row_factory=dict_row
+			) as conn:
+				# Execute the query with a new cursor
+				with conn.cursor() as cur:
+					cur.execute(query)
+					try:
+						rows = cur.fetchall()
+					except psycopg.ProgrammingError as e:
+						if str(e) == "no results to fetch":
+							rows = []
+						else:
+							raise
+					return rows
+		except psycopg.OperationalError as oe:
 			print(f"OperationalError while connecting to the server: {oe}")
 		except Exception as e:
 			print(f"Unexpected error while connecting to the server: {e}")
-		finally:
-			# Ensure the connection is always closed, even if an error occurs
-			if conn is not None:
-				conn.close()
 
 	def _sync_query(self, query, new_connection=False):
 		"""Synchronous function to perform the query.

pystackql/stackql.py

@@ -2,7 +2,7 @@
 
 # Install packages if they aren't already installed
 # pip3 install -r requirements.txt --user
-# pip3 install psycopg2-binary --user
+# pip3 install psycopg --user
 
 # Load environment variables
 source ./tests/creds/env_vars/test.env.sh

run_tests

@@ -1,6 +1,6 @@
 # install packages
 # pip.exe install -r requirements.txt --user
-# pip install psycopg2-binary --user
+# pip install psycopg[binary]
 
 # Load environment variables
 . .\tests\creds\env_vars\test.env.ps1