Allow specification of Snowflake schema and database (#63)

* Allow specification of schema and database for Snowflake.

* Snowflake - Move to the last hour of query history by default.

* Version bump.

* Add handler to allow use of internal Pydantic field names.

* Update documentation.

README.md

@@ -57,6 +57,9 @@ isn't listed here, support can be added by creating a custom connector!
 * SalesForce Marketing Cloud audit event logs
 * SalesForce Marketing Cloud security event logs
 * Slack audit logs
+* Snowflake login history
+* Snowflake query history
+* Snowflake session history
 * Stripe events
 * Tines audit logs
 * Terraform Cloud audit trails

docs/index.rst

@@ -49,6 +49,7 @@ Currently the following log sources are supported by Grove out of the box. If a
 isn't listed here, support can be added by creating a custom connector!
 
 * Atlassian audit events (e.g. Confluence, Jira)
+* FleetDM host logs
 * GitHub audit logs
 * GSuite alerts
 * GSuite activity logs
@@ -62,6 +63,9 @@ isn't listed here, support can be added by creating a custom connector!
 * SalesForce Marketing Cloud audit event logs
 * SalesForce Marketing Cloud security event logs
 * Slack audit logs
+* Snowflake login history
+* Snowflake query history
+* Snowflake session history
 * Stripe events
 * Tines audit logs
 * Terraform Cloud audit trails

grove/__about__.py

@@ -1,6 +1,6 @@
 """Grove metadata."""
 
-__version__ = "1.5.0"
+__version__ = "1.6.0"
 __title__ = "grove"
 __license__ = "Mozilla Public License 2.0"
 __copyright__ = "Copyright 2023 HashiCorp, Inc."

grove/connectors/snowflake/common.py

@@ -51,14 +51,14 @@ def batch_size(self) -> int:
         This is used to control the maximum number of records which will be retrieved
         before they are flushed to the output handler.
 
-        The default is 500.
+        The default is 1000.
 
         :return: The "batch_size" portion of the connector's configuration.
         """
         try:
             candidate = self.configuration.batch_size
         except AttributeError:
-            return 500
+            return 1000
 
         try:
             candidate = int(candidate)
@@ -114,3 +114,32 @@ def passphrase(self) -> Optional[str]:
             return self.configuration.passphrase
         except AttributeError:
             return None
+
+    @property
+    def schema(self) -> Optional[str]:
+        """Fetches the optional schema name from the configuration.
+
+        The default is "SNOWFLAKE".
+
+        :return: The "schema" portion of the connector's configuration.
+        """
+        try:
+            # The trailing underscore is due to a limitation in Pydantic < 2.0 where
+            # 'schema' is an internal field. We automatically remap these internal
+            # fields with a trailing underscore while we migrate to Pydantic >= 2.0
+            return self.configuration.schema_
+        except AttributeError:
+            return "SNOWFLAKE"
+
+    @property
+    def database(self) -> Optional[str]:
+        """Fetches the optional database name from the configuration.
+
+        The default is "ADMIN".
+
+        :return: The "database" portion of the connector's configuration.
+        """
+        try:
+            return self.configuration.database
+        except AttributeError:
+            return "ADMIN"

grove/connectors/snowflake/login_history.py

@@ -14,7 +14,7 @@
 # Define the paramaterised Snowflake query to use to fetch login history records.
 SNOWFLAKE_QUERY_LOGIN_HISTORY = """
   SELECT *
-    FROM SNOWFLAKE.ACCOUNT_USAGE.LOGIN_HISTORY
+    FROM LOGIN_HISTORY
    WHERE EVENT_TIMESTAMP >  %(pointer)s
 ORDER BY EVENT_TIMESTAMP ASC;
  """
@@ -43,7 +43,9 @@ def collect(self):
             client = snowflake.connector.connect(
                 role=self.role,
                 user=self.identity,
+                schema=self.schema,
                 account=self.account,
+                database=self.database,
                 warehouse=self.warehouse,
                 private_key=private_key,
                 timezone="UTC",

grove/connectors/snowflake/query_history.py

@@ -18,7 +18,7 @@
          ROLE_NAME,
          START_TIME,
          END_TIME
-    FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY
+    FROM QUERY_HISTORY
    WHERE START_TIME > %(pointer)s
 ORDER BY START_TIME ASC;
  """
@@ -32,11 +32,12 @@ class Connector(SnowflakeConnector):
     def collect(self):
         """Collects query history records from Snowflake."""
 
-        # If no pointer is stored then start from 24-hours ago - due to volume.
+        # If no pointer is stored then start from 1 hour ago. We use a small value here
+        # due to volume.
         try:
             _ = self.pointer
         except NotFoundException:
-            self.pointer = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat()
+            self.pointer = (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat()
 
         # Decode the private key from the loaded PEM format PKCS#8 data.
         private_key = self._load_private_key()
@@ -47,7 +48,9 @@ def collect(self):
             client = snowflake.connector.connect(
                 role=self.role,
                 user=self.identity,
+                schema=self.schema,
                 account=self.account,
+                database=self.database,
                 warehouse=self.warehouse,
                 private_key=private_key,
                 timezone="UTC",

grove/connectors/snowflake/session_history.py

@@ -14,7 +14,7 @@
 # Define the paramaterised Snowflake query to use to fetch session history records.
 SNOWFLAKE_QUERY_SESSION_HISTORY = """
   SELECT *
-    FROM SNOWFLAKE.ACCOUNT_USAGE.SESSIONS
+    FROM SESSIONS
    WHERE CREATED_ON >  %(pointer)s
 ORDER BY CREATED_ON ASC;
  """
@@ -43,7 +43,9 @@ def collect(self):
             client = snowflake.connector.connect(
                 role=self.role,
                 user=self.identity,
+                schema=self.schema,
                 account=self.account,
+                database=self.database,
                 warehouse=self.warehouse,
                 private_key=private_key,
                 timezone="UTC",

grove/models.py

@@ -143,6 +143,19 @@ def _decode_fields(cls, values):  # noqa: B902
         Other encoding schemes may be supported in future, but for now only base64 is
         supported.
         """
+        # This is a horrible hack to allow fields with names that mask Pydantic
+        # internals. This can be removed once Grove is updated to use Pydantic >= 2.
+        INTERNAL_FIELDS = ["schema"]
+
+        for field in INTERNAL_FIELDS:
+            value = values.get(field, None)
+            if value is None:
+                continue
+
+            # Remap the field name to contain a trailing underscore.
+            values[f"{field}_"] = value
+            del values[field]
+
         for field, encoding in values.get("encoding", {}).items():
             # If the secret is externally stored decoding will be performed after the
             # secret has been retrieved. Right now, this field should not exist as it