tests added. parser coded

Author: Arsalanjdev

src/rowgen/hf_api.py

@@ -30,7 +30,7 @@ def prompt_fake_data(self, db_schema: str, num_rows: int = 10) -> str:
         :return: str. the response of ai api.
         """
         prompt = f"""
-        Generate {num_rows} rows of fake sql data for the following table:column schema. Make the data look believable and realistic (not john doe or anything like that). 
+        Generate {num_rows} rows of fake sql data for the following table:column schema. Make the data look believable and realistic (not john doe or anything like that). make sure there's consistency between tables (such as foregin keys and so on).
         Your response should be only in plain sql insert statements with no words before or after. table:column schema: {db_schema}
         """
         return self.send_message_to_api(prompt)

src/rowgen/main.py

@@ -4,7 +4,8 @@
 import argparse
 
 from rowgen.extract_from_db import DBconnect
-from rowgen.parser import parse_sql_from_code_block
+from rowgen.sql_parser import parse_sql_from_code_block
+import os
 
 
 def main():
@@ -25,16 +26,21 @@ def main():
         action="store_true",
         help="Executes insert statements into the database directly.",
     )
-
+    parser.add_argument("--apikey", help="Enter your huggingface_hub api key.")
     args = parser.parse_args()
 
+    if args.apikey:
+        api_key = args.apikey
+    else:
+        api_key = get_api_key_from_config()
+
     if args.db_url:
         db_url = args.db_url
     else:
         db_url = "{args.db_type}://{args.user}:{args.password}@{args.host}:{args.port}/{args.database}"
 
     dbc = DBconnect(db_url)
-    hf = HFapi()
+    hf = HFapi(api_key=api_key)
     ai_sql_response = hf.prompt_fake_data(dbc.table_columns, 20)
     sql_statements = parse_sql_from_code_block(ai_sql_response)
     # execute
@@ -54,5 +60,46 @@ def main():
         print("Saved to sql file.")
 
 
+def get_api_key_from_config() -> str:
+    """
+    Retrieves API key from the config file (~/.config/rowgen/conf).
+    If no key is stored, prompts the user for input and saves it.
+
+    :return: str. The API key.
+    """
+    config_path = os.path.expanduser("~/.config/rowgen/conf")
+    apikey = None
+
+    try:
+        with open(config_path, "r") as file:
+            for line in file:
+                if line.startswith("apikey:"):
+                    apikey = line.split(":", 1)[1].strip()
+                    break  # no need to keep reading
+
+    except FileNotFoundError:
+        pass  # we'll prompt the user below
+    except PermissionError:
+        print("You do not have permission to read the config file.")
+        return ""
+    except Exception as e:
+        print(f"An unexpected error occurred while reading: {e}")
+        return ""
+
+    if not apikey:
+        # Prompt the user
+        apikey = input("Enter your API key: ").strip()
+
+        try:
+            os.makedirs(os.path.dirname(config_path), exist_ok=True)
+            with open(config_path, "w") as file:
+                file.write(f"apikey: {apikey}\n")
+                print(f"API key was store in the {config_path}")
+        except Exception as e:
+            print(f"Failed to save API key: {e}")
+
+    return apikey
+
+
 if __name__ == "__main__":
     main()

src/rowgen/parser.py src/rowgen/sql_parser.pysrc/rowgen/parser.py renamed to src/rowgen/sql_parser.py

@@ -39,16 +39,12 @@ def test_hf_hub_ping(hf_api_client):
         pytest.fail(f"Failed to ping Hugging Face Hub: {e}")
 
 
-def test_hf_inference_connection(inference_client):
+def test_hf_api_send_message_to_api(hf_api):
     prompt_message = "Hi! This is a test. say banana if you can hear me."
     try:
-        completion = inference_client.chat.completions.create(
-            model="deepseek-ai/DeepSeek-V3",
-            messages=[{"role": "user", "content": prompt_message}],
-        )
-        message: str = completion.choices[0].message["content"]
-        assert len(message) > 0
-        assert "banana" in message.casefold()
+        response_message = hf_api.send_message_to_api(prompt_message)
+        assert len(response_message) > 0
+        assert "banana" in response_message.casefold()
 
     except BadRequestError as e:
         print(e)
@@ -57,12 +53,13 @@ def test_hf_inference_connection(inference_client):
         pytest.fail(f"Failed inference call: {d}")
 
 
-def test_hf_api_send_message_to_api(hf_api):
-    prompt_message = "Hi! This is a test. say banana if you can hear me."
+def test_hf_sql_message(hf_api):
+    prompt_message = "Write five sql statements for the table Books made up of id, title, isbn, genre, publication_date, author. Your message should only consist of sql statements with no sentence before and after."
     try:
         response_message = hf_api.send_message_to_api(prompt_message)
         assert len(response_message) > 0
-        assert "banana" in response_message.casefold()
+        assert response_message.startswith("```sql")
+        assert response_message.endswith("```")
 
     except BadRequestError as e:
         print(e)

src/rowgen/testrowgendb.sqlite

@@ -1,5 +1,5 @@
 import json
-from rowgen.parser import parse_sql_from_code_block
+from rowgen.sql_parser import parse_sql_from_code_block
 from rowgen.hf_api import HFapi
 from pytest import fixture